Segmentation d’images avec SAM 2 - Guide ultime 📷

Graphique

Introduction

Dans ce guide complet, nous allons explorer comment mettre en œuvre la segmentation d’images à l’aide du modèle Segment Anything Model (SAM) version 2. Cette technologie de pointe nous permet de segmenter n’importe quel objet dans une image sans avoir besoin de données d’apprentissage extensives ou de configurations compliquées. À la fin de ce tutoriel, vous disposerez d’un pipeline robuste pour effectuer des tâches de segmentation d’images précises et efficaces.

Conditions préalables

Python 3.10+
torch version >=2.0.1
transformers version >=4.26.0
matplotlib
opencv-python

📺 Regarder : Les réseaux de neurones expliqués

Vidéo par 3Blue1Brown

Installez les packages requis avec :

pip install torch==2.0.1 transformers==4.26.0 matplotlib opencv-python

Étape 1 : Configuration du projet

Créez d’abord un nouveau répertoire pour votre projet et initialisez-le en tant que package Python si nécessaire.

Dans ce répertoire, configurez l’environnement en créant les fichiers suivants :

requirements.txt : Liste des packages requis.
setup.py (si vous prévoyez de distribuer votre code).

Ajoutez le contenu suivant à votre requirements.txt :

torch>=2.0.1
transformers>=4.26.0
matplotlib
opencv-python

Puis, dans votre terminal ou invite de commandes, accédez au répertoire de votre projet et exécutez :

pip install -r requirements.txt

Étape 2 : Mise en œuvre centrale

La partie centrale de notre mise en œuvre consiste à intégrer SAM dans un pipeline de segmentation d’images de base. Voici un exemple complet :

import torch
from transformers import SamPredictor, sam_model_registry
from PIL import Image
import matplotlib.pyplot as plt
import cv2

def fonction_principale():
    # Définir le type de modèle et le chemin vers le point de contrôle
    model_type = "vit_h"
    checkpoint_path = "chemin/vers/sam_vit_h_4b8939.pth"

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Initialiser le prédicteur SAM
    predictor = SamPredictor(sam_model_registry[model_type].to(device))

    # Charger une image à partir d'un fichier et la convertir en format RGB
    image_path = 'chemin/vers/image.jpg'
    img_pil = Image.open(image_path).convert("RGB")
    img_cv2 = cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)

    # Définir l'image dans le prédicteur SAM pour la segmentation
    predictor.set_image(img_cv2)

    # Facultativement, fournir des points de prompt et des masques d'entrée ici.
    # Pour simplifier, supposons que nous avons quelques prompts prêts.

    # Générer un masque à l'aide de prompts (le cas échéant)
    # Voici un exemple où vous pourriez spécifier les coordonnées d'un point sur l'image
    point_coords = [500, 375]  # Exemple de coordonnée pour la segmentation

    masks, _, _ = predictor.predict(point_coords=point_coords)

    # Visualiser le masque sur l'image d'origine
    plt.figure(figsize=(10, 10))
    plt.imshow(img_cv2)
    show_mask(masks[0], plt.gca(), random_color=True)  # Fonction pour visualiser le masque

def show_mask(mask, ax=None, random_color=False):
    """Visualise un masque binaire en le plaçant sur l'image d'origine."""
    if not ax:
        fig, ax = plt.subplots(1, figsize=(20, 20))

    if random_color:
        color = np.concatenate([np.random.random(3), np.array([.3, .3, .3])], axis=0)
    else:
        color = np.array([251/255, 49/255, 47/255, 0.6])

    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)

fonction_principale()

Étape 3 : Configuration

Personnalisez les variables model_type et checkpoint_path en fonction de la version de votre modèle SAM. Ajoutez également le chemin vers votre image d’entrée et les coordonnées du point pour la segmentation.

# Exemple de personnalisation
model_type = "vit_b" # Passer à 'vit_h' ou à d'autres variantes SAM si vous utilisez un point de contrôle différent.
checkpoint_path = "chemin/vers/sam_vit_b_01ec64.pth"
image_path = 'votre/chemin/vers/image.jpg'
point_coords = [750, 325]  # Ajuster les coordonnées du point en fonction de votre image et de l'objet cible

Étape 4 : Exécution du code

Pour exécuter votre script de segmentation :

python main.py
# Sortie attendue :
# Une fenêtre matplotlib avec l'image d'entrée et le masque superposé.

Assurez-vous de remplacer main.py par le nom réel de votre fichier Python.

Étape 5 : Astuces avancées

Pour une performance optimale, envisagez de pré-traiter les images pour qu’elles soient de tailles standard avant de les envoyer à SAM. Expérimentez également avec différents points de prompt et des masques pour obtenir des résultats de segmentation plus précis. Enfin, integre

Segmentation d'image avec SAM 2 - Le guide ultime 📸