Découvrez SAM 2 pour une segmentation vidéo précise et performante en traitement d'image

Plongez dans l'univers de la vision par ordinateur et de la segmentation vidéo avec la puissance du modèle Segment Anything Model 2 (SAM 2). Ce tutoriel vous guidera à travers les étapes nécessaires pour utiliser efficacement SAM 2, un modèle unifié de segmentation d'image et vidéo qui marque une avancée significative en termes de précision, de vitesse et de sensibilité au contexte.

Comprendre SAM 2

SAM 2 adresse les défis uniques de la segmentation vidéo en gérant des facteurs variables tels que le mouvement des objets, les déformations, les occultations, les changements d'éclairage et bien plus. Doté d'une précision accrue et de moins d'interactions manuelles par rapport aux méthodes antérieures, SAM 2 se démarque également par sa rapidité pour la segmentation d'image, avec une vitesse six fois supérieure au modèle SAM original.

Installation et configuration

Pour commencer à exploiter SAM 2 pour le traitement vidéo, initiez d'abord le clonage du dépôt et l'installation des dépendances nécessaires :

git clone https://github.com/facebookresearch/segment-anything-2.git
cd segment-anything-2
pip install -e .
python setup.py build_ext --inplace

En raison d'un bug précédemment identifié, exécutez la commande supplémentaire python setup.py build_ext --inplace après l'installation initiale. Notez que l’installation de SAM 2 en local peut prendre un certain temps.

Prétraitement des données avec SAM 2

Le modèle SAM 2 inclut une mémoire incorporée permettant de stocker les informations relatives aux objets et aux interactions passées. Cela favorise la génération de masques de prédiction tout au long de la vidéo.

Avant de débuter la segmentation, les cadres doivent être sauvegardés sur disque au format JPEG :

import supervision as sv
 
frames_generator = sv.get_video_frames_generator()
sink = sv.ImageSink(target_dir_path='path/to/frames', image_name_pattern="{:05d}.jpeg")
 
with sink:
    for frame in frames_generator:
        sink.save_image(frame)

Initialisation de l'état d'inférence pour SAM 2

Pour une segmentation vidéo interactive, une initialisation de l'état d'inférence est requise :

inference_state = sam2_model.init_state()

Si un suivi antérieur a été effectué avec cet état d'inférence, assurez-vous de le réinitialiser avec sam2_model.reset_state(inference_state).

Segmentation et suivi d'un objet avec SAM 2

Démarrez par tenter de segmenter un objet dans le premier cadre de la vidéo :

import numpy as np
import torch
from sam2.build_sam import build_sam2_video_predictor
 
CHECKPOINT = "checkpoints/sam2_hiera_large.pt"
CONFIG = "sam2_hiera_l.yaml"
sam2_model = build_sam2_video_predictor(CONFIG, CHECKPOINT)
 
points = np.array([[703, 303]], dtype=np.float32)
labels = np.array([1])  # Label 1 pour une inclusion
frame_idx = 0
tracker_id = 1
 
_, object_ids, mask_logits = sam2_model.add_new_points(
    inference_state=inference_state,
    frame_idx=frame_idx,
    obj_id=tracker_id,
    points=points,
    labels=labels,
)

SAM 2 permet aussi l'utilisation de points négatifs pour exclure certaines régions de la segmentation :

points = np.array([[703, 303], [731, 256], [713, 356], [740, 297]], dtype=np.float32)
labels = np.array([1, 0, 0, 0])  # les points négatifs sont marqués par 0
 
_, object_ids, mask_logits = sam2_model.add_new_points(
    inference_state=inference_state,
    frame_idx=frame_idx,
    obj_id=tracker_id,
    points=points,
    labels=labels,
)

Propagation des prompts à travers la vidéo

Avec le générateur propagate_in_video, appliquez vos prompts à tous les cadres de la vidéo :

import cv2
import supervision as sv
 
color_palette = ['#FF1493', '#00BFFF', '#FF6347', '#FFD700']
mask_annotator = sv.MaskAnnotator(color=sv.ColorPalette.from_hex(color_palette), color_lookup=sv.ColorLookup.TRACK)
video_info = sv.VideoInfo.from_video_path('path/to/video')
frames_paths = sorted(sv.list_files_with_extensions(directory='path/to/frames', extensions=["jpeg"]))
 
with sv.VideoSink('output/path/of/video', video_info=video_info) as sink:
    for frame_idx, object_ids, mask_logits in sam2_model.propagate_in_video(inference_state):
        frame = cv2.imread(frames_paths[frame_idx])
        masks = (mask_logits > 0.0).cpu().numpy()
        N, X, H, W = masks.shape
        masks = masks.reshape(N * X, H, W)
        detections = sv.Detections(
            xyxy=sv.mask_to_xyxy(masks=masks),
            mask=masks,
            tracker_id=np.array(object_ids)
        )
        frame = mask_annotator.annotate(frame, detections)
        sink.write_frame(frame)

Conclusions

Bien que SAM 2 puisse rencontrer des difficultés dans certaines situations, comme les scènes encombrées ou les objets de détail fin en mouvement rapide, il reste un outil performant dans le domaine de la segmentation image et vidéo, avec une gamme étendue d'applications à travers divers domaines.

Ce tutoriel a pour but de vous initier à SAM 2 et de vous accompagner dans vos premiers pas vers la maîtrise de la segmentation vidéo. Comme l’intelligence artificielle et la vision par ordinateur continuent d'évoluer, les capacités de SAM 2 représentent un atout pour les professionnels cherchant à simplifier et à améliorer leurs workflows de traitement d'image et vidéo.

Pour plus d'informations sur SAM 2, référez-vous à l'article original écrit par Piotr Skalski : How to Use SAM 2 for Video Segmentation.