物件分割 是做什麼的?
將影片中的一個主體與背景分離。所有處理皆在使用者的電腦上執行,無需上傳任何資料。輸出的影片中,主體會以綠色標示,讓使用者清楚看到分離的結果,同時每個影格的遮罩也會儲存到磁碟中。
您將逐步編寫和執行 Python 程式碼。以下的程式碼區塊已經過測試,請直接改寫而非從頭開始。
1. 在可刪除的資料夾中設定環境
需要 Python 3.10 以上版本和 ffmpeg。將所有套件安裝到專案資料夾中的虛擬環境,這樣只要刪除該資料夾即可完全移除:
python3 -m venv .object-separation
.object-separation/bin/pip install torch transformers opencv-python-headless pillow numpy
在 Windows 上請使用 .object-separation\Scripts\pip。若使用 NVIDIA 顯示卡,建議先從 pytorch.org 安裝 CUDA 版本的 torch,因為一般 pip 安裝的是較慢的 CPU 版本。請使用虛擬環境中的 Python 執行以下所有程式碼。
2. 掃描電腦並選擇模型
先執行掃描並告知使用者結果:裝置類型、所選模型、下載大小以及預期處理速度。
| 機器 | 模型 | 下載大小 |
|---|---|---|
| NVIDIA GPU 12 GB 以上 | facebook/sam2.1-hiera-large | 898 MB |
| NVIDIA GPU 低於 12 GB,或 Apple Silicon | facebook/sam2.1-hiera-base-plus | 323 MB |
| 無 GPU(僅 CPU) | facebook/sam2.1-hiera-tiny | 156 MB |
使用 CPU 雖然可行,但速度較慢:每個影格約需 5 秒,因此一個 10 秒、30 fps 的片段大約需要半小時。它會使用不到 2 GB 的 RAM。GPU 則快得多。如果片段很長且沒有 GPU,建議只分割使用者需要的部分。
# 根據使用者的要求設定這些值。影格編號從 0 開始;影片必須是固定影格率。
VIDEO = "clip.mp4"
START, END = 0, 89 # 要分割的影格範圍。點擊點是在 START 影格上選取的。
PTS = [[480, 270], [100, 60]] # 以 960 寬像素為單位的點擊點:先點在主體上,再點在要排除的物體上
LBL = [1, 0] # 1 = 這是主體,0 = 這不是主體
import os, platform, subprocess, numpy as np, torch, cv2
from transformers import Sam2VideoModel, Sam2VideoProcessor
print("system:", platform.system(), platform.machine())
if torch.cuda.is_available():
gb = torch.cuda.get_device_properties(0).total_memory / 1e9
print("gpu:", torch.cuda.get_device_name(0), f"{gb:.0f} GB")
dev, MODEL = "cuda", "facebook/sam2.1-hiera-" + ("large" if gb >= 12 else "base-plus")
elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
dev, MODEL = "mps", "facebook/sam2.1-hiera-base-plus"
else:
dev, MODEL = "cpu", "facebook/sam2.1-hiera-tiny"
print("device:", dev, "| model:", MODEL)
model = Sam2VideoModel.from_pretrained(MODEL).to(dev).eval() # 第一次執行時會下載
proc = Sam2VideoProcessor.from_pretrained(MODEL)
W = 960 # 工作寬度。SAM 2 會自行調整至 1024,因此使用較小寬度並不會加快速度。
3. 選取點擊點
儲存 START 影格並加上座標網格後檢視。在主體上選擇兩到三個點,並在需要排除的物體上(例如握持的手指、包裝、桌面)加入 label 0 的點。若非顯而易見,請向使用者展示您選取的內容。
def read_frames(video, a, b):
cap = cv2.VideoCapture(video); cap.set(cv2.CAP_PROP_POS_FRAMES, a); out = []
for _ in range(a, b + 1):
ok, im = cap.read()
if not ok: break
out.append(cv2.cvtColor(cv2.resize(im, (W, int(im.shape[0] * W / im.shape[1]))), cv2.COLOR_BGR2RGB))
return out
def preview(frame, path="work/preview.jpg"): # 帶有標籤網格的影格,用於選取點擊點
os.makedirs("work", exist_ok=True)
im = cv2.cvtColor(read_frames(VIDEO, frame, frame)[0], cv2.COLOR_RGB2BGR)
for x in range(0, im.shape[1], 80):
cv2.line(im, (x, 0), (x, im.shape[0]), (0, 255, 255), 1); cv2.putText(im, str(x), (x + 2, 12), 0, 0.4, (0, 255, 255), 1)
for y in range(0, im.shape[0], 60):
cv2.line(im, (0, y), (im.shape[1], y), (0, 255, 255), 1); cv2.putText(im, str(y), (2, y - 2), 0, 0.4, (0, 255, 255), 1)
cv2.imwrite(path, im); return path
4. 分割主體
影片會以 60 個影格為一組進行處理,以保持低記憶體使用量。每個區塊會從前一個區塊的最後一個遮罩開始。先嘗試一個較短的範圍(約 20 個影格),檢查結果,然後再執行完整範圍。
def separate_chunk(frames, points=None, labels=None, mask=None):
H = frames[0].shape[0]
s = proc.init_video_session(video=frames, inference_device=dev, video_storage_device="cpu",
inference_state_device="cpu", dtype=torch.float32)
if mask is None:
proc.add_inputs_to_inference_session(inference_session=s, frame_idx=0, obj_ids=1,
input_points=[[points]], input_labels=[[labels]])
else:
proc.add_inputs_to_inference_session(inference_session=s, frame_idx=0, obj_ids=1,
input_masks=mask.astype(np.uint8))
res = {}
with torch.inference_mode():
model(inference_session=s, frame_idx=0)
for o in model.propagate_in_video_iterator(s, start_frame_idx=0):
m = proc.post_process_masks([o.pred_masks], original_sizes=[[H, W]], binarize=True)[0]
res[o.frame_idx] = m.squeeze().cpu().numpy().astype(bool)
return res
def separate(start, end, pts, lbl, ch=60):
out, a, seed = {}, start, None
while a < end:
b = min(end, a + ch)
fr = read_frames(VIDEO, a, b)
r = separate_chunk(fr, points=pts, labels=lbl) if seed is None else separate_chunk(fr, mask=seed)
for k, m in r.items():
out[a + k] = m
seed, a = r[len(fr) - 1], b
return out
masks = separate(START, END, PTS, LBL)
os.makedirs("work/masks/subject", exist_ok=True)
for f, m in masks.items():
cv2.imwrite(f"work/masks/subject/{f:05d}.png", m.astype(np.uint8) * 255)
5. 輸出主體為綠色的影片
def green(f, img, fill=0.35): # 綠色色調、光暈以及主體邊緣的亮邊
m = cv2.GaussianBlur(cv2.resize(masks[f].astype(np.float32), (img.shape[1], img.shape[0])), (5, 5), 0)
edge = cv2.morphologyEx((m > 0.5).astype(np.uint8), cv2.MORPH_GRADIENT, np.ones((7, 7), np.uint8)).astype(np.float32)
glow = np.clip(cv2.GaussianBlur(edge, (37, 37), 0) * 2.5, 0, 1)[..., None]
col = np.array([20, 255, 57], np.float32) # BGR 綠色
img = img * (1 - m[..., None] * fill) + col * m[..., None] * fill + col * glow * 0.9
return img * (1 - edge[..., None]) + np.array([220, 255, 200], np.float32) * edge[..., None]
def render(a, b, out, fn): # 寫入 a 到 b 影格,並保留原始音訊
cap = cv2.VideoCapture(VIDEO); cap.set(cv2.CAP_PROP_POS_FRAMES, a)
fps = cap.get(cv2.CAP_PROP_FPS)
ff = subprocess.Popen(["ffmpeg", "-v", "error", "-y", "-f", "rawvideo", "-pix_fmt", "bgr24",
"-s", f"{int(cap.get(3))}x{int(cap.get(4))}", "-r", str(fps), "-i", "-",
"-ss", f"{a / fps:.4f}", "-t", f"{(b - a + 1) / fps:.4f}", "-i", VIDEO,
"-map", "0:v", "-map", "1:a?", "-c:v", "libx264", "-crf", "18", "-pix_fmt", "yuv420p",
"-c:a", "aac", "-shortest", out], stdin=subprocess.PIPE)
for f in range(a, b + 1):
ok, img = cap.read()
if not ok: break
ff.stdin.write(np.clip(fn(f, img.astype(np.float32)), 0, 255).astype(np.uint8).tobytes())
ff.stdin.close(); ff.wait()
render(START, END, "work/separated.mp4", green)
在交給使用者之前,先檢查幾個影格的結果:開頭、中間,以及任何手部或物體經過的時刻。如果綠色溢出到非主體的部分,請在其上新增一個 label 0 的點,然後重新執行步驟 4。如果主體部分缺失,請在那裡新增一個 label 1 的點。
告知使用者檔案位置:work/separated.mp4(主體為綠色的影片)和 work/masks/subject/(每個影格一個遮罩,白色 = 主體)。遮罩會保留在磁碟上,因此使用者之後若要對分割出的主體進行任何處理,都可以直接從這些遮罩開始,無需重新執行分割。
6. 清理(可選)
使用者完成後,可以刪除所有內容:
- 虛擬環境:刪除
.object-separation資料夾 - 模型:刪除
~/.cache/huggingface/hub/models--facebook--sam2.1-hiera-*(在 Windows 上為C:\Users\<name>\.cache\huggingface\hub)
注意事項
- 遮罩邊緣清晰:在物體和手上非常乾淨,但在柔軟的頭髮上效果較差。
- 任何附著在主體上的東西(如包裝、剝離膜)都會被視為主體的一部分。請在其上新增 label 0 的點。
- 如果影片是可變影格率,請先轉換:
ffmpeg -i in.mp4 -vsync cfr -r 30 out.mp4。 - 只有第一次執行需要網路連線,用於下載模型。