object-separation

Compare original and translation side by side

🇺🇸

Original

English

🇨🇳

Translation

Chinese

Object separation

目标分离

Separate one subject from the background in a video. Everything runs on the user's computer and nothing is uploaded. The result is the same video with the subject highlighted in green, so the user can see exactly what was separated, plus a mask for every frame saved to disk.

You write and run the Python as you go. The code blocks below are tested; adapt them rather than starting from scratch.

将一个主体与视频背景分离。所有内容都在用户的计算机上运行，不会上传任何数据。输出结果为同一段视频，其中主体以绿色突出显示，方便用户准确查看已分离的主体；同时还会将每一帧的蒙版保存到磁盘。

你需要在执行过程中编写并运行 Python 代码。下面的代码块均经过测试；应基于这些代码进行调整，而不是从头编写。

1. Set up in a folder that can be deleted

1. 在可删除的文件夹中完成设置

Needs python 3.10+ and ffmpeg. Install everything into a virtual environment in the project folder, so removing that one folder removes it all:

python3 -m venv .object-separation
.object-separation/bin/pip install torch transformers opencv-python-headless pillow numpy

On Windows use

.object-separation\Scripts\pip

. With an NVIDIA card on Windows, install the CUDA build of torch from pytorch.org first, because plain pip gives the slower CPU build. Run all the Python below with the venv's python.

需要 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. Scan the computer and pick the model

2. 扫描计算机并选择模型

Run the scan first and tell the user what it found: the device, the model it picked, the download size, and what speed to expect.

Machine	Model	Download
NVIDIA GPU with 12 GB+	`facebook/sam2.1-hiera-large`	898 MB
NVIDIA GPU under 12 GB, or Apple Silicon	`facebook/sam2.1-hiera-base-plus`	323 MB
No GPU (CPU only)	`facebook/sam2.1-hiera-tiny`	156 MB

On a CPU it works but it is slow: about 5 seconds per frame, so a 10 second clip at 30 fps takes around half an hour. It uses under 2 GB of RAM. A GPU is much faster. If the clip is long and there is no GPU, suggest separating only the part the user needs.

python

undefined

先运行扫描，并告知用户检测结果：设备、所选模型、下载大小以及预期速度。

机器	模型	下载大小
配备 12 GB 以上显存的 NVIDIA GPU	`facebook/sam2.1-hiera-large`	898 MB
显存低于 12 GB 的 NVIDIA GPU，或 Apple Silicon	`facebook/sam2.1-hiera-base-plus`	323 MB
无 GPU（仅 CPU）	`facebook/sam2.1-hiera-tiny`	156 MB

在 CPU 上可以运行，但速度较慢：每帧约 5 秒，因此一段 30 fps、10 秒的片段大约需要半小时。内存占用低于 2 GB。GPU 会快得多。如果片段较长且没有 GPU，建议只分离用户需要的部分。

python

undefined

Set these from the user's request. Frame numbers start at 0; the video must be constant frame rate.

VIDEO = "clip.mp4" START, END = 0, 89 # frames to separate. Click points are picked on the START frame. PTS = [[480, 270], [100, 60]] # click points in 960-wide pixels: on the subject first, then on things to leave out LBL = [1, 0] # 1 = this is the subject, 0 = this is NOT the subject


```python
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()     # downloads on first run
proc = Sam2VideoProcessor.from_pretrained(MODEL)
W = 960   # working width. SAM 2 resizes to 1024 itself, so a smaller width is NOT faster.


```python
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()     # downloads on first run
proc = Sam2VideoProcessor.from_pretrained(MODEL)
W = 960   # working width. SAM 2 resizes to 1024 itself, so a smaller width is NOT faster.

3. Pick the click points

3. 选择点击点

Save frame START with a coordinate grid and look at it. Choose two or three points on the subject and add points with label 0 on anything that must be left out (fingers holding it, packaging, the table). Show the user what you picked if it is not obvious.

python

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"):       # frame with a labelled grid, to pick click points
    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

保存 START 帧并添加坐标网格，然后查看该帧。在主体上选择两到三个点，并为需要排除的任何内容（拿着它的手指、包装、桌子等）添加标签为 0 的点。如果所选位置不明显，请向用户展示。

python

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"):       # frame with a labelled grid, to pick click points
    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. Separate the subject

4. 分离主体

The video is processed in chunks of 60 frames so memory stays low. Each chunk starts from the last mask of the one before. Try a short range first (about 20 frames), check it, then run the full range.

python

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)

视频按每 60 帧一块进行处理，以降低内存占用。每个块都以上一块的最后一个蒙版为起点。先尝试较短的帧范围（约 20 帧），检查结果后，再运行完整范围。

python

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. Output the video with the subject in green

5. 输出主体呈绿色的视频

python

def green(f, img, fill=0.35):                       # green tint, glow and a bright edge on the subject
    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 green
    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):                          # writes the frames a..b with the original audio
    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)

Look at a few frames of the result before handing it over: the start, the middle, and any moment where a hand or object passes in front. If green spills onto something that is not the subject, add a label-0 point on it and run step 4 again. If part of the subject is missing, add a label-1 point there.

Tell the user where the files are:

work/separated.mp4

(the video with the subject in green) and

work/masks/subject/

(one mask per frame, white = subject). The masks stay on disk, so anything the user wants to do with the separated subject next can start from them without running the separation again.

python

def green(f, img, fill=0.35):                       # green tint, glow and a bright edge on the subject
    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 green
    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):                          # writes the frames a..b with the original audio
    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. Clean up (optional)

6. 清理（可选）

Once the user is done, everything can be removed:

the virtual environment: delete the
```
.object-separation
```
folder

the model: delete

~/.cache/huggingface/hub/models--facebook--sam2.1-hiera-*

(on Windows:

C:\Users\<name>\.cache\huggingface\hub

)

用户使用完毕后，可以删除所有内容：

虚拟环境：删除
```
.object-separation
```
文件夹

模型：删除

~/.cache/huggingface/hub/models--facebook--sam2.1-hiera-*

（在 Windows 上：

C:\Users\<name>\.cache\huggingface\hub

）

Notes

注意

Masks are hard-edged: very clean on objects and hands, less so on soft hair.
Anything attached to the subject counts as part of it (packaging, peel-off film). Add a label-0 point on it.
If the video is variable frame rate, convert it first:
```
ffmpeg -i in.mp4 -vsync cfr -r 30 out.mp4
```
.
Only the first run needs the internet, to download the model.

蒙版为硬边：对物体和手部的效果非常干净，对细软发丝的效果则稍差。
与主体相连的所有内容都会被视为主体的一部分（包装、可撕膜等）。请在其上添加一个标签 0 的点。
如果视频为可变帧率，请先转换：
```
ffmpeg -i in.mp4 -vsync cfr -r 30 out.mp4
```
。
仅首次运行需要联网以下载模型。