객체 분리

객체 분리은(는) 무엇을 하나요?

비디오에서 하나의 주체를 배경으로부터 분리합니다. 모든 작업은 사용자 컴퓨터에서 실행되며 아무것도 업로드되지 않습니다. 결과는 주체가 녹색으로 강조된 동일한 비디오이며, 각 프레임에 대한 마스크가 디스크에 저장됩니다.

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초, 따라서 30fps 10초 클립은 약 30분이 소요됩니다. RAM은 2 GB 미만을 사용합니다. GPU는 훨씬 빠릅니다. 클립이 길고 GPU가 없는 경우, 필요한 부분만 분리하도록 제안하세요.

# 사용자 요청에 따라 설정하세요. 프레임 번호는 0부터 시작합니다. 비디오는 고정 프레임 레이트여야 합니다.
VIDEO = "clip.mp4"
START, END = 0, 89              # 분리할 프레임. 클릭 포인트는 START 프레임에서 선택됩니다.
PTS = [[480, 270], [100, 60]]   # 960px 너비 기준 클릭 포인트: 주체를 먼저, 그 다음 제외할 부분
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 프레임을 좌표 그리드와 함께 저장하고 확인합니다. 주체에 2~3개의 포인트를 선택하고, 제외해야 할 부분(손가락, 포장, 테이블)에 레이블 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\<name>\.cache\huggingface\hub)

참고 사항

• 마스크는 경계가 뚜렷합니다: 물체와 손은 매우 깔끔하고, 부드러운 머리카락은 덜 깔끔합니다.
• 주체에 부착된 모든 것은 주체의 일부로 간주됩니다(포장, 필름). 해당 부분에 레이블 0 포인트를 추가하세요.
• 비디오의 프레임 레이트가 가변적인 경우, 먼저 변환하세요: ffmpeg -i in.mp4 -vsync cfr -r 30 out.mp4.
• 첫 번째 실행만 인터넷이 필요합니다(모델 다운로드).