对象分离 是做什么的?
将视频中的一个主体与背景分离。一切均在用户的计算机上运行,不上传任何数据。结果是一个原始视频,其中主体以绿色高亮显示,以便用户精确看到被分离的内容,同时每一帧的遮罩都保存到磁盘。
您可以一边编写一边运行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。如果在 Windows 上使用 NVIDIA 显卡,请先从 pytorch.org 安装 CUDA 版本的 torch,因为普通的 pip 会安装较慢的 CPU 版本。使用虚拟环境的 Python 运行以下所有 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。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 帧并添加坐标网格,然后查看它。在主体上选择两到三个点,并在任何必须排除的内容(比如握住它的手指、包装、桌子)上添加标签为 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)
在交出结果前,查看几帧:开始帧、中间帧,以及任何手或物体经过前面的时刻。如果绿色溢出到非主体部分,在其上添加一个标签为 0 的点,然后重新运行第 4 步。如果主体的一部分缺失,在那里添加一个标签为 1 的点。
告知用户文件的位置:work/separated.mp4(主体以绿色显示的视频)和 work/masks/subject/(每帧一个遮罩,白色 = 主体)。遮罩保留在磁盘上,因此用户后续想对分离出的主体进行任何操作时,都可以从这些遮罩开始,而无需再次运行分离过程。
6. 清理(可选)
用户完成后,所有内容都可以删除:
- 虚拟环境:删除
.object-separation文件夹 - 模型:删除
~/.cache/huggingface/hub/models--facebook--sam2.1-hiera-*(在 Windows 上:C:\Users\<用户名>\.cache\huggingface\hub)
备注
- 遮罩边缘清晰:对于物体和手,效果非常干净;对于柔软的头发,效果稍差。
- 任何附着在主体上的东西都算作主体的一部分(包装、可撕膜)。在其上添加一个标签为 0 的点。
- 如果视频是可变帧率,请先转换:
ffmpeg -i in.mp4 -vsync cfr -r 30 out.mp4。 - 只有首次运行需要联网,用于下载模型。