Synergie entre YOLO et SAM en Vision par Ordinateur

L'association de YOLO (You Only Look Once) et du Segment Anything Model (SAM) crée une synergie puissante pour l'analyse d'images. Tandis que YOLO excelle dans la détection rapide d'objets, SAM se distingue par sa capacité à segmenter avec une grande précision. En combinant ces deux architectures, il est possible d'obtenir non seulement la localisation des objets, mais également leurs contours exacts, offrant ainsi un contexte riche pour les tâches en aval.

Architecture et Caractéristiques de SAM

Introduit par Meta, SAM repose sur une approche multi-étapes pour la segmentation d'images. Son architecture intègre des mécanismes d'attention qui permettent au modèle de se concentrer dynamiquement sur les régions pertinentes, ainsi que des connexions résiduelles pour préserver les détails spatiaux tout en extrayant des caractéristiques sémantiques de haut niveau.

Les points forts de ce modèle incluent :

• Polyvalence : Capacité à segmenter n'importe quel objet, indépendamment de sa catégorie.
• Robustesse : Performances stables face à des conditions d'éclairage variées et des orientations complexes.
• Compréhension contextuelle : Utilisation du contexte global pour différencier les objets dans des scènes encombrées.

Implémentation de SAM pour la Segmentation d'Images

Voici comment configurer et utiliser SAM pour générer des masques de segmentation sur une image statique.

import os
import urllib.request
import torch
import cv2
import numpy as np
import matplotlib.pyplot as plt
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator

# Configuration des chemins et téléchargement des poids
BASE_DIR = os.path.abspath(".")
WEIGHTS_DIR = os.path.join(BASE_DIR, "modeles")
os.makedirs(WEIGHTS_DIR, exist_ok=True)

WEIGHT_URL = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
WEIGHT_PATH = os.path.join(WEIGHTS_DIR, "sam_vit_h.pth")

if not os.path.exists(WEIGHT_PATH):
    urllib.request.urlretrieve(WEIGHT_URL, WEIGHT_PATH)

# Initialisation du dispositif et du modèle
device = "cuda" if torch.cuda.is_available() else "cpu"
sam_model = sam_model_registry["vit_h"](checkpoint=WEIGHT_PATH).to(device)
mask_gen = SamAutomaticMaskGenerator(sam_model)

# Chargement et traitement de l'image
img_path = "exemple.jpg"
img_bgr = cv2.imread(img_path)
img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

# Génération des masques
resultats_masques = mask_gen.generate(img_rgb)

# Visualisation des résultats
plt.figure(figsize=(10, 5))
plt.subplot(1, 2, 1)
plt.imshow(img_rgb)
plt.title("Image Originale")
plt.axis("off")

plt.subplot(1, 2, 2)
mask_overlay = np.zeros_like(img_rgb)
for m in resultats_masques:
    mask_overlay[m['segmentation']] = np.random.randint(0, 255, 3)
blended = cv2.addWeighted(img_rgb, 0.6, mask_overlay, 0.4, 0)
plt.imshow(blended)
plt.title("Masques SAM")
plt.axis("off")
plt.show()

Traitement Vidéo avec YOLO et SAM

Pour appliquer cette combinaison à un flux vidéo, YOLO est utilisé pour extraire les boîtes englobentes, qui servent ensuite d'invites (prompts) à SAM pour générer les masques précis.

import cv2
import numpy as np
import torch
from ultralytics import YOLO
from segment_anything import SamPredictor

# Initialisation des modèles
yolo_model = YOLO("yolov8n.pt")
sam_predictor = SamPredictor(sam_model)

# Paramètres de la vidéo
video_src = "entree.mp4"
video_dst = "sortie.avi"
cap = cv2.VideoCapture(video_src)
fps = cap.get(cv2.CAP_PROP_FPS)
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

out = cv2.VideoWriter(video_dst, cv2.VideoWriter_fourcc(*'MJPG'), fps, (w, h))

# Définition des couleurs
COULEUR_BOITE = (255, 165, 0)
COULEUR_TEXTE = (255, 255, 255)
COULEUR_MASQUE = (0, 200, 0)

while cap.isOpened():
    succes, frame = cap.read()
    if not succes:
        break
        
    # Détection des objets avec YOLO
    resultats_yolo = yolo_model(frame, conf=0.5)[0]
    boites = resultats_yolo.boxes.xyxy.cpu().numpy()
    classes = resultats_yolo.boxes.cls.cpu().numpy()
    
    if len(boites) == 0:
        out.write(frame)
        continue
        
    # Préparation des boîtes pour SAM
    sam_predictor.set_image(frame)
    tenseurs_boites = torch.tensor(boites, device=sam_predictor.device)
    boites_transformees = sam_predictor.transform.apply_boxes_torch(tenseurs_boites, frame.shape[:2])
    
    # Prédiction des masques avec SAM
    masques, _, _ = sam_predictor.predict_torch(
        point_coords=None,
        point_labels=None,
        boxes=boites_transformees,
        multimask_output=False,
    )
    
    # Superposition des masques et des boîtes englobantes
    canvas = frame.copy()
    for i, masque in enumerate(masques):
        m = masque[0].cpu().numpy()
        
        # Application du masque coloré
        couleur_masque = np.zeros_like(frame, dtype=np.uint8)
        couleur_masque[m] = COULEUR_MASQUE
        canvas = cv2.addWeighted(canvas, 1, couleur_masque, 0.4, 0)
        
        # Dessin de la boîte et du label
        x1, y1, x2, y2 = map(int, boites[i])
        cv2.rectangle(canvas, (x1, y1), (x2, y2), COULEUR_BOITE, 2)
        
        nom_classe = yolo_model.names[int(classes[i])]
        cv2.putText(canvas, nom_classe, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, COULEUR_TEXTE, 2)
        
    out.write(canvas)

cap.release()
out.release()
cv2.destroyAllWindows()

Cas d'Usage et Applications

• Véhicules autonomes : Amélioration de la perception environnementale pour une navigation sécurisée.
• Imagerie médicale : Segmentation précise des anomalies dans les radiographies ou les IRM.
• Vidéosurveillance : Suivi et isolation d'individus ou d'objets spécifiques dans des flux en temps réel.
• Automatisation industrielle : Détection de défauts microscopiques sur les chaînes de production.
• Agriculture de précision : Identification et segmentation des cultures et des nuisibles par drone.
• Réalité augmentée : Isolation des objets du monde réel pour une intégration fluide d'éléments virtuels.