separacao-de-objetos

Separação de objetos

Separa um sujeito do fundo em um vídeo. Tudo é executado no computador do usuário e nada é enviado para servidores. O resultado é o mesmo vídeo com o sujeito destacado em verde, permitindo que o usuário veja exatamente o que foi separado, além de uma máscara para cada quadro salva em disco.

Você escreve e executa o Python conforme avança. Os blocos de código abaixo são testados; adapte-os em vez de começar do zero.

1. Configurar em uma pasta que pode ser deletada

Requer python 3.10+ e ffmpeg. Instale tudo em um ambiente virtual na pasta do projeto, de modo que remover essa pasta remova tudo:

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

No Windows use .object-separation\Scripts\pip. Com uma placa NVIDIA no Windows, instale a build CUDA do torch a partir do pytorch.org primeiro, porque o pip simples fornece a build de CPU mais lenta. Execute todo o Python abaixo com o python do venv.

2. Escanear o computador e escolher o modelo

Execute a varredura primeiro e informe ao usuário o que foi encontrado: o dispositivo, o modelo escolhido, o tamanho do download e a velocidade esperada.

Máquina	Modelo	Download
GPU NVIDIA com 12 GB+	facebook/sam2.1-hiera-large	898 MB
GPU NVIDIA abaixo de 12 GB, ou Apple Silicon	facebook/sam2.1-hiera-base-plus	323 MB
Sem GPU (apenas CPU)	facebook/sam2.1-hiera-tiny	156 MB

Na CPU funciona, mas é lento: cerca de 5 segundos por quadro, então um clipe de 10 segundos a 30 fps leva cerca de meia hora. Usa menos de 2 GB de RAM. Uma GPU é muito mais rápida. Se o clipe for longo e não houver GPU, sugira separar apenas a parte que o usuário precisa.

# Defina estes a partir da solicitação do usuário. Os números dos quadros começam em 0; o vídeo deve ter taxa de quadros constante.
VIDEO = "clip.mp4"
START, END = 0, 89              # quadros a separar. Os pontos de clique são escolhidos no quadro START.
PTS = [[480, 270], [100, 60]]   # pontos de clique em pixels de largura 960: primeiro no sujeito, depois em coisas a ignorar
LBL = [1, 0]                    # 1 = este é o sujeito, 0 = este NÃO é o sujeito

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()     # baixa na primeira execução
proc = Sam2VideoProcessor.from_pretrained(MODEL)
W = 960   # largura de trabalho. O SAM 2 redimensiona para 1024 internamente, então uma largura menor NÃO é mais rápida.

3. Escolher os pontos de clique

Salve o quadro START com uma grade de coordenadas e observe-a. Escolha dois ou três pontos no sujeito e adicione pontos com rótulo 0 em qualquer coisa que deva ser ignorada (dedos segurando-o, embalagem, a mesa). Mostre ao usuário o que você escolheu se não for óbvio.

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"):       # quadro com uma grade rotulada, para escolher pontos de clique
    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. Separar o sujeito

O vídeo é processado em blocos de 60 quadros para manter a memória baixa. Cada bloco começa a partir da última máscara do anterior. Teste primeiro um intervalo curto (cerca de 20 quadros), verifique-o e depois execute o intervalo completo.

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. Gerar o vídeo com o sujeito em verde

def green(f, img, fill=0.35):                       # tom verde, brilho e uma borda brilhante no sujeito
    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 verde
    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):                          # escreve os quadros a..b com o áudio original
    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)

Observe alguns quadros do resultado antes de entregá-lo: o início, o meio e qualquer momento em que uma mão ou objeto passe na frente. Se o verde respingar em algo que não é o sujeito, adicione um ponto com rótulo 0 nele e execute a etapa 4 novamente. Se parte do sujeito estiver faltando, adicione um ponto com rótulo 1 lá.

Informe ao usuário onde estão os arquivos: work/separated.mp4 (o vídeo com o sujeito em verde) e work/masks/subject/ (uma máscara por quadro, branco = sujeito). As máscaras permanecem no disco, então qualquer coisa que o usuário queira fazer com o sujeito separado em seguida pode começar a partir delas sem executar a separação novamente.

6. Limpeza (opcional)

Quando o usuário terminar, tudo pode ser removido:

• o ambiente virtual: delete a pasta .object-separation
• o modelo: delete ~/.cache/huggingface/hub/models--facebook--sam2.1-hiera-* (no Windows: C:\Users\<nome>\.cache\huggingface\hub)

Notas

• As máscaras têm bordas nítidas: muito limpas em objetos e mãos, menos em cabelos finos.
• Qualquer coisa anexada ao sujeito conta como parte dele (embalagem, filme de remoção). Adicione um ponto com rótulo 0 nela.
• Se o vídeo tiver taxa de quadros variável, converta-o primeiro: ffmpeg -i in.mp4 -vsync cfr -r 30 out.mp4.
• Apenas a primeira execução precisa da internet, para baixar o modelo.