Objekt-Trennung

Was macht Objekt-Trennung?

Trennen Sie ein Motiv vom Hintergrund in einem Video. Alles läuft auf dem Computer des Benutzers und es wird nichts hochgeladen. Das Ergebnis ist dasselbe Video mit dem in Grün hervorgehobenen Motiv, sodass der Benutzer genau sehen kann, was getrennt wurde, plus eine Maske für jedes Bild, die auf der Festplatte gespeichert wird.

Sie schreiben und führen das Python während des Vorgangs aus. Die folgenden Codeblöcke sind getestet; passen Sie sie an, anstatt von Grund auf neu zu beginnen.

1. Einrichtung in einem Ordner, der gelöscht werden kann

Benötigt Python 3.10+ und ffmpeg. Installieren Sie alles in einer virtuellen Umgebung im Projektordner, sodass das Entfernen dieses einen Ordners alles entfernt:

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

Unter Windows verwenden Sie .object-separation\Scripts\pip. Wenn Sie eine NVIDIA-Karte unter Windows haben, installieren Sie zuerst den CUDA-Build von torch von pytorch.org, da einfaches pip die langsamere CPU-Version liefert. Führen Sie das gesamte folgende Python mit dem Python der virtuellen Umgebung aus.

2. Computer durchsuchen und Modell auswählen

Führen Sie zuerst den Scan durch und teilen Sie dem Benutzer mit, was gefunden wurde: das Gerät, das ausgewählte Modell, die Downloadgröße und die zu erwartende Geschwindigkeit.

Maschine	Modell	Download
NVIDIA GPU mit 12 GB+	facebook/sam2.1-hiera-large	898 MB
NVIDIA GPU unter 12 GB, oder Apple Silicon	facebook/sam2.1-hiera-base-plus	323 MB
Keine GPU (nur CPU)	facebook/sam2.1-hiera-tiny	156 MB

Auf einer CPU funktioniert es, ist aber langsam: etwa 5 Sekunden pro Bild, also dauert ein 10-Sekunden-Clip bei 30 fps etwa eine halbe Stunde. Es verbraucht unter 2 GB RAM. Eine GPU ist viel schneller. Wenn der Clip lang ist und keine GPU vorhanden ist, schlagen Sie vor, nur den Teil zu trennen, den der Benutzer benötigt.

# Legen Sie diese Werte aus der Benutzeranfrage fest. Bildnummern beginnen bei 0; das Video muss eine konstante Bildrate haben.
VIDEO = "clip.mp4"
START, END = 0, 89              # zu trennende Bilder. Klickpunkte werden auf dem START-Bild ausgewählt.
PTS = [[480, 270], [100, 60]]   # Klickpunkte in 960 Pixel Breite: zuerst auf das Motiv, dann auf Dinge, die ausgelassen werden sollen
LBL = [1, 0]                    # 1 = das ist das Motiv, 0 = das ist NICHT das Motiv

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()     # wird beim ersten Ausführen heruntergeladen
proc = Sam2VideoProcessor.from_pretrained(MODEL)
W = 960   # Arbeitsbreite. SAM 2 skaliert selbst auf 1024, daher ist eine kleinere Breite NICHT schneller.

3. Klickpunkte auswählen

Speichern Sie das Bild START mit einem Koordinatengitter und betrachten Sie es. Wählen Sie zwei oder drei Punkte auf dem Motiv und fügen Sie Punkte mit der Bezeichnung 0 auf allem hinzu, das ausgelassen werden muss (Finger, die es halten, Verpackung, der Tisch). Zeigen Sie dem Benutzer, was Sie ausgewählt haben, wenn es nicht offensichtlich ist.

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"):       # Bild mit beschriftetem Gitter zum Auswählen von Klickpunkten
    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. Das Motiv trennen

Das Video wird in Blöcken von 60 Bildern verarbeitet, damit der Speicher niedrig bleibt. Jeder Block beginnt mit der letzten Maske des vorherigen. Versuchen Sie zuerst einen kurzen Bereich (ca. 20 Bilder), überprüfen Sie ihn und führen Sie dann den vollständigen Bereich aus.

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. Das Video mit dem grün hervorgehobenen Motiv ausgeben

def green(f, img, fill=0.35):                       # grüne Tönung, Leuchten und eine helle Kante am Motiv
    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 grün
    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):                          # schreibt die Bilder a..b mit dem 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)

Sehen Sie sich einige Bilder des Ergebnisses an, bevor Sie es übergeben: den Anfang, die Mitte und jeden Moment, in dem eine Hand oder ein Objekt davor vorbeikommt. Wenn Grün auf etwas übergeht, das nicht das Motiv ist, fügen Sie einen Punkt mit der Bezeichnung 0 darauf hinzu und führen Sie Schritt 4 erneut aus. Wenn ein Teil des Motivs fehlt, fügen Sie dort einen Punkt mit der Bezeichnung 1 hinzu.

Teilen Sie dem Benutzer mit, wo sich die Dateien befinden: work/separated.mp4 (das Video mit dem grünen Motiv) und work/masks/subject/ (eine Maske pro Bild, weiß = Motiv). Die Masken bleiben auf der Festplatte, sodass alles, was der Benutzer als Nächstes mit dem getrennten Motiv tun möchte, von ihnen ausgehen kann, ohne die Trennung erneut ausführen zu müssen.

6. Aufräumen (optional)

Wenn der Benutzer fertig ist, kann alles entfernt werden:

• die virtuelle Umgebung: löschen Sie den Ordner .object-separation
• das Modell: löschen Sie ~/.cache/huggingface/hub/models--facebook--sam2.1-hiera-* (unter Windows: C:\Users\<Name>\.cache\huggingface\hub)

Hinweise

• Masken haben harte Kanten: sehr sauber bei Objekten und Händen, weniger bei weichem Haar.
• Alles, was am Motiv befestigt ist, gilt als Teil davon (Verpackung, Abziehfolie). Fügen Sie einen Punkt mit der Bezeichnung 0 darauf hinzu.
• Wenn das Video eine variable Bildrate hat, konvertieren Sie es zuerst: ffmpeg -i in.mp4 -vsync cfr -r 30 out.mp4.
• Nur der erste Durchlauf benötigt das Internet, um das Modell herunterzuladen.