separación-de-objetos

Separación de objetos

Separa un sujeto del fondo en un video. Todo se ejecuta en el ordenador del usuario y no se sube nada. El resultado es el mismo video con el sujeto resaltado en verde, para que el usuario pueda ver exactamente qué se separó, además de una máscara por cada fotograma guardada en disco.

Escribe y ejecuta el código Python sobre la marcha. Los bloques de código siguientes están probados; adáptalos en lugar de empezar desde cero.

1. Configuración en una carpeta que se pueda eliminar

Necesita python 3.10+ y ffmpeg. Instala todo en un entorno virtual dentro de la carpeta del proyecto, de modo que eliminar esa carpeta lo borre todo:

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

En Windows usa .object-separation\Scripts\pip. Con una tarjeta NVIDIA en Windows, instala primero la compilación CUDA de torch desde pytorch.org, porque pip simple da la compilación CPU más lenta. Ejecuta todo el código Python siguiente con el python del entorno virtual.

2. Escanea el ordenador y elige el modelo

Ejecuta primero el escaneo e informa al usuario de lo que encontró: el dispositivo, el modelo elegido, el tamaño de descarga y la velocidad esperada.

Máquina	Modelo	Descarga
GPU NVIDIA con 12 GB+	facebook/sam2.1-hiera-large	898 MB
GPU NVIDIA con menos de 12 GB, o Apple Silicon	facebook/sam2.1-hiera-base-plus	323 MB
Sin GPU (solo CPU)	facebook/sam2.1-hiera-tiny	156 MB

En CPU funciona pero es lento: aproximadamente 5 segundos por fotograma, por lo que un clip de 10 segundos a 30 fps tarda alrededor de media hora. Usa menos de 2 GB de RAM. Una GPU es mucho más rápida. Si el clip es largo y no hay GPU, sugiere separar solo la parte que el usuario necesita.

# Establece estos valores a partir de la solicitud del usuario. Los números de fotograma empiezan en 0; el video debe tener velocidad de fotogramas constante.
VIDEO = "clip.mp4"
START, END = 0, 89              # fotogramas a separar. Los puntos de clic se eligen en el fotograma START.
PTS = [[480, 270], [100, 60]]   # puntos de clic en píxeles de 960 de ancho: primero en el sujeto, luego en cosas que dejar fuera
LBL = [1, 0]                    # 1 = esto es el sujeto, 0 = esto NO es el sujeto

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()     # se descarga en la primera ejecución
proc = Sam2VideoProcessor.from_pretrained(MODEL)
W = 960   # ancho de trabajo. SAM 2 redimensiona a 1024 por sí mismo, por lo que un ancho menor NO es más rápido.

3. Elige los puntos de clic

Guarda el fotograma START con una cuadrícula de coordenadas y míralo. Elige dos o tres puntos en el sujeto y añade puntos con etiqueta 0 en cualquier cosa que deba excluirse (dedos que lo sostienen, embalaje, la mesa). Muestra al usuario lo que elegiste si no es obvio.

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"):       # fotograma con una cuadrícula etiquetada, para elegir puntos de clic
    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. Separa el sujeto

El video se procesa en fragmentos de 60 fotogramas para que la memoria se mantenga baja. Cada fragmento comienza desde la última máscara del anterior. Prueba primero un rango corto (unos 20 fotogramas), verifica el resultado y luego ejecuta el rango 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. Genera el video con el sujeto en verde

def green(f, img, fill=0.35):                       # tinte verde, brillo y un borde brillante en el sujeto
    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):                          # escribe los fotogramas a..b con el audio 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)

Mira algunos fotogramas del resultado antes de entregarlo: el inicio, el medio y cualquier momento en que una mano u objeto pase por delante. Si el verde se extiende sobre algo que no es el sujeto, añade un punto con etiqueta 0 sobre ello y ejecuta el paso 4 de nuevo. Si falta parte del sujeto, añade un punto con etiqueta 1 allí.

Informa al usuario de dónde están los archivos: work/separated.mp4 (el video con el sujeto en verde) y work/masks/subject/ (una máscara por fotograma, blanco = sujeto). Las máscaras permanecen en disco, por lo que cualquier cosa que el usuario quiera hacer después con el sujeto separado puede comenzar desde ellas sin volver a ejecutar la separación.

6. Limpieza (opcional)

Una vez que el usuario haya terminado, se puede eliminar todo:

• el entorno virtual: elimina la carpeta .object-separation
• el modelo: elimina ~/.cache/huggingface/hub/models--facebook--sam2.1-hiera-* (en Windows: C:\Users\<nombre>\.cache\huggingface\hub)

Notas

• Las máscaras tienen bordes definidos: muy limpias en objetos y manos, menos en cabello suave.
• Cualquier cosa unida al sujeto cuenta como parte de él (embalaje, película despegable). Añade un punto con etiqueta 0 sobre ello.
• Si el video tiene velocidad de fotogramas variable, conviértelo primero: ffmpeg -i in.mp4 -vsync cfr -r 30 out.mp4.
• Solo la primera ejecución necesita internet, para descargar el modelo.