Séparation d'objet
Séparez un sujet de l'arrière-plan dans une vidéo. Tout s'exécute sur l'ordinateur de l'utilisateur et rien n'est téléchargé. Le résultat est la même vidéo avec le sujet mis en évidence en vert, afin que l'utilisateur puisse voir exactement ce qui a été séparé, ainsi qu'un masque pour chaque image enregistré sur le disque.
Vous écrivez et exécutez le code Python au fur et à mesure. Les blocs de code ci-dessous sont testés ; adaptez-les plutôt que de partir de zéro.
1. Configuration dans un dossier pouvant être supprimé
Nécessite Python 3.10+ et ffmpeg. Installez tout dans un environnement virtuel dans le dossier du projet, de sorte que la suppression de ce seul dossier supprime tout :
python3 -m venv .object-separation
.object-separation/bin/pip install torch transformers opencv-python-headless pillow numpy
Sous Windows, utilisez .object-separation\Scripts\pip. Avec une carte NVIDIA sous Windows, installez d'abord la version CUDA de torch depuis pytorch.org, car pip standard donne la version CPU plus lente. Exécutez tout le code Python ci-dessous avec le python de l'environnement virtuel.
2. Analyse de l'ordinateur et choix du modèle
Exécutez d'abord l'analyse et informez l'utilisateur de ce qui a été trouvé : le périphérique, le modèle sélectionné, la taille du téléchargement et la vitesse attendue.
| Machine | Modèle | Téléchargement |
|---|---|---|
| GPU NVIDIA avec 12 Go+ | facebook/sam2.1-hiera-large | 898 Mo |
| GPU NVIDIA < 12 Go, ou Apple Silicon | facebook/sam2.1-hiera-base-plus | 323 Mo |
| Pas de GPU (CPU uniquement) | facebook/sam2.1-hiera-tiny | 156 Mo |
Sur un CPU, cela fonctionne mais c'est lent : environ 5 secondes par image, donc un clip de 10 secondes à 30 ips prend environ une demi-heure. Cela utilise moins de 2 Go de RAM. Un GPU est beaucoup plus rapide. Si le clip est long et qu'il n'y a pas de GPU, suggérez de séparer uniquement la partie dont l'utilisateur a besoin.
# Définissez ces valeurs à partir de la demande de l'utilisateur. Les numéros d'image commencent à 0 ; la vidéo doit avoir un débit d'images constant.
VIDEO = "clip.mp4"
START, END = 0, 89 # images à séparer. Les points de clic sont sélectionnés sur l'image START.
PTS = [[480, 270], [100, 60]] # points de clic en pixels de largeur 960 : d'abord sur le sujet, puis sur les éléments à exclure
LBL = [1, 0] # 1 = c'est le sujet, 0 = ce n'est PAS le sujet
import os, platform, subprocess, numpy as np, torch, cv2
from transformers import Sam2VideoModel, Sam2VideoProcessor
print("système :", 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} Go")
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("périphérique :", dev, "| modèle :", MODEL)
model = Sam2VideoModel.from_pretrained(MODEL).to(dev).eval() # téléchargement lors de la première exécution
proc = Sam2VideoProcessor.from_pretrained(MODEL)
W = 960 # largeur de travail. SAM 2 redimensionne lui-même à 1024, donc une largeur plus petite n'est PAS plus rapide.
3. Choix des points de clic
Sauvegardez l'image START avec une grille de coordonnées et examinez-la. Choisissez deux ou trois points sur le sujet et ajoutez des points avec le label 0 sur tout ce qui doit être exclu (doigts le tenant, emballage, table). Montrez à l'utilisateur ce que vous avez sélectionné si ce n'est pas évident.
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"): # image avec une grille étiquetée, pour choisir les points 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. Séparation du sujet
La vidéo est traitée par blocs de 60 images pour que la mémoire reste faible. Chaque bloc commence à partir du dernier masque du précédent. Essayez d'abord une courte plage (environ 20 images), vérifiez-la, puis exécutez la plage complète.
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. Sortie de la vidéo avec le sujet en vert
def green(f, img, fill=0.35): # teinte verte, lueur et bordure brillante sur le sujet
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 vert
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): # écrit les images a..b avec l'audio d'origine
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)
Examinez quelques images du résultat avant de le remettre : le début, le milieu et tout moment où une main ou un objet passe devant. Si le vert déborde sur quelque chose qui n'est pas le sujet, ajoutez un point de label 0 dessus et réexécutez l'étape 4. Si une partie du sujet est manquante, ajoutez un point de label 1 à cet endroit.
Indiquez à l'utilisateur où se trouvent les fichiers : work/separated.mp4 (la vidéo avec le sujet en vert) et work/masks/subject/ (un masque par image, blanc = sujet). Les masques restent sur le disque, de sorte que tout ce que l'utilisateur souhaite faire ensuite avec le sujet séparé peut en partir sans relancer la séparation.
6. Nettoyage (optionnel)
Une fois que l'utilisateur a terminé, tout peut être supprimé :
- l'environnement virtuel : supprimez le dossier
.object-separation - le modèle : supprimez
~/.cache/huggingface/hub/models--facebook--sam2.1-hiera-*(sous Windows :C:\Users\<nom>\.cache\huggingface\hub)
Remarques
- Les masques ont des bords nets : très propres sur les objets et les mains, moins sur les cheveux fins.
- Tout ce qui est attaché au sujet compte comme en faisant partie (emballage, film pelable). Ajoutez un point de label 0 dessus.
- Si la vidéo a un débit d'images variable, convertissez-la d'abord :
ffmpeg -i in.mp4 -vsync cfr -r 30 out.mp4. - Seule la première exécution nécessite Internet, pour télécharger le modèle.