Génération de masques

La génération de masques est la tâche qui consiste à générer des masques sémantiquement significatifs pour une image. Cette tâche est très similaire à la segmentation d'image, mais de nombreuses différences existent. Les modèles de segmentation d'image sont entraînés sur des jeux de données étiquetés et sont limités aux classes qu'ils ont vues pendant l'entraînement ; ils renvoient un ensemble de masques et les classes correspondantes, étant donné une image.

Les modèles de génération de masques sont entraînés sur de grandes quantités de données et fonctionnent en deux modes.

Mode de sollicitation : dans ce mode, le modèle reçoit une image et une sollicitation, où la sollicitation peut être un emplacement de point 2D (coordonnées XY) dans l'image à l'intérieur d'un objet, ou une boîte englobante entourant un objet. En mode de sollicitation, le modèle ne renvoie que le masque de l'objet que la sollicitation désigne.
Mode Segment Everything : dans le mode segment everything, à partir d'une image, le modèle génère tous les masques présents dans l'image. Pour ce faire, une grille de points est générée et superposée à l'image pour l'inférence.
Inférence vidéo : le modèle accepte une vidéo ainsi qu’un point ou une boîte de prompt dans une image de la vidéo, qui est suivie tout au long de la séquence. Vous pouvez obtenir plus d’informations sur la manière d’effectuer une inférence vidéo en consultant la documentation de SAM 2.

La tâche de génération de masques est prise en charge par Segment Anything Model (SAM) et Segment Anything Model 2 (SAM2), tandis que l’inférence vidéo est prise en charge par Segment Anything Model 2 (SAM2). SAM est un modèle puissant qui comprend un encodeur d’images basé sur Vision Transformer, un encodeur de prompts et un décodeur de masques bidirectionnel. Les images et les prompts sont encodés, puis le décodeur utilise ces embeddings pour générer des masques valides. SAM 2 étend SAM en ajoutant un module mémoire pour suivre les masques.

SAM constitue un modèle de fondation puissant pour la segmentation, grâce à sa large couverture de données. Il est entraîné sur SA-1B, un jeu de données composé de 1 million d'images et de 1,1 milliard de masques.

Dans ce guide, vous apprendrez à :

Inférez en mode « tout segmenter » avec lot de traitement,
Inférez en mode de sollicitation par points,
Inférez en mode de sollicitation par boîte.

Tout d'abord, installons transformers :

pip install -q transformers

Pipeline de génération de masques

La façon la plus simple d’inférer des modèles de génération de masque est d’utiliser le pipeline mask-generation.

>>> from transformers import pipeline

>>> checkpoint = "facebook/sam2-hiera-base-plus"
>>> mask_generator = pipeline(model=checkpoint, task="mask-generation")

Regardons l'image.

from PIL import Image
import requests

img_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee.jpg"
image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")

Segmentons tout. points-per-batch permet l’inférence parallèle des points en mode « segmenter tout ». Cela accélère l’inférence, mais consomme davantage de mémoire. De plus, SAM n’active le batching que sur les points, et non sur les images. pred_iou_thresh est le seuil de confiance IoU au-delà duquel seules les masques sont retournés.

masks = mask_generator(image, points_per_batch=128, pred_iou_thresh=0.88)

Les masks ressemblent à ce qui suit :

{'masks': [tensor([[False, False, False,  ...,  True,  True,  True],
          [False, False, False,  ...,  True,  True,  True],
          [False, False, False,  ...,  True,  True,  True],
          ...,
          [False, False, False,  ..., False, False, False], .. 
 'scores': tensor([0.9874, 0.9793, 0.9780, 0.9776, ... 0.9016])}

Nous pouvons les visualiser ainsi :

import matplotlib.pyplot as plt

plt.imshow(image, cmap='gray')

for i, mask in enumerate(masks["masks"]):
    plt.imshow(mask, cmap='viridis', alpha=0.1, vmin=0, vmax=1)

plt.axis('off')
plt.show()

Voici l'image d'origine en niveaux de gris avec des cartes colorées superposées. Très impressionnant.

Inférence du modèle

Point Prompting

Vous pouvez également utiliser le modèle sans le pipeline. Pour ce faire, initialisez le modèle et le processeur.

from transformers import SamModel, SamProcessor
from accelerate import Accelerator
import torch
device = Accelerator().device
model = SamModel.from_pretrained("facebook/sam-vit-base").to(device)
processor = SamProcessor.from_pretrained("facebook/sam-vit-base")

Pour effectuer un point prompting, transmettez le point d’entrée au processeur, puis prenez la sortie du processeur et transmettez-la au modèle pour l’inférence. Pour post-traiter la sortie du modèle, transmettez les sorties et original_sizes sont prises de la sortie initiale du processeur. Il est nécessaire de les transmettre car le processeur redimensionne l’image, et la sortie doit être extrapolée.

input_points = [[[2592, 1728]]] # point location of the bee

inputs = processor(image, input_points=input_points, return_tensors="pt").to(device)
with torch.no_grad():
    outputs = model(**inputs)
masks = processor.image_processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"].cpu())

Nous pouvons visualiser les trois masques dans la sortie masks.

import matplotlib.pyplot as plt
import numpy as np

fig, axes = plt.subplots(1, 4, figsize=(15, 5))

axes[0].imshow(image)
axes[0].set_title('Original Image')
mask_list = [masks[0][0][0].numpy(), masks[0][0][1].numpy(), masks[0][0][2].numpy()]

for i, mask in enumerate(mask_list, start=1):
    overlayed_image = np.array(image).copy()

    overlayed_image[:,:,0] = np.where(mask == 1, 255, overlayed_image[:,:,0])
    overlayed_image[:,:,1] = np.where(mask == 1, 0, overlayed_image[:,:,1])
    overlayed_image[:,:,2] = np.where(mask == 1, 0, overlayed_image[:,:,2])
    
    axes[i].imshow(overlayed_image)
    axes[i].set_title(f'Mask {i}')
for ax in axes:
    ax.axis('off')

plt.show()

Box Prompting

Vous pouvez également effectuer du box prompting de manière similaire au point prompting. Il suffit de transmettre la boîte d'entrée au format d'une liste [x_min, y_min, x_max, y_max] ainsi que l'image au processor. Prenez la sortie du processor et transmettez-la directement au modèle, puis post-traitez à nouveau la sortie.

# bounding box around the bee
box = [2350, 1600, 2850, 2100]

inputs = processor(
        image,
        input_boxes=[[[box]]],
        return_tensors="pt"
    ).to(model.device)

with torch.no_grad():
    outputs = model(**inputs)

mask = processor.image_processor.post_process_masks(
    outputs.pred_masks.cpu(),
    inputs["original_sizes"].cpu(),
)[0][0][0].numpy()

Vous pouvez visualiser la boîte englobante autour de l’abeille, comme indiqué ci-dessous.

import matplotlib.patches as patches

fig, ax = plt.subplots()
ax.imshow(image)

rectangle = patches.Rectangle((2350, 1600), 500, 500, linewidth=2, edgecolor='r', facecolor='none')
ax.add_patch(rectangle)
ax.axis("off")
plt.show()

Vous pouvez voir la sortie d'inférence ci-dessous.

fig, ax = plt.subplots()
ax.imshow(image)
ax.imshow(mask, cmap='viridis', alpha=0.4)

ax.axis("off")
plt.show()

Ajustement fin pour la génération de masques

Nous allons effectuer un fine-tuning de SAM2.1 sur une petite partie du jeu de données MicroMat pour le matting d’images. Nous devons installer la bibliothèque monai pour utiliser la perte DICE, ainsi que trackio pour journaliser les masques pendant l’entraînement.

pip install -q datasets monai trackio
``` 
We can now load our dataset and take a look.

```python
from datasets import load_dataset

dataset = load_dataset("merve/MicroMat-mini", split="train")
dataset
# Dataset({
#    features: ['image', 'mask', 'prompt', 'image_id', 'object_id', 'sample_idx', 'granularity', 
# 'image_path', 'mask_path', 'prompt_path'],  num_rows: 94
#})

Nous avons besoin des colonnes image, masque et prompt. Nous effectuons le découpage pour l'entraînement et le test.

dataset = dataset.train_test_split(test_size=0.1)
train_ds = dataset["train"]
val_ds = dataset["test"]

Examinons un exemple.

train_ds[0]

 {'image': <PIL.PngImagePlugin.PngImageFile image mode=RGB size=2040x1356>,
 'mask': <PIL.PngImagePlugin.PngImageFile image mode=L size=2040x1356>,
 'prompt': '{"point": [[137, 1165, 1], [77, 1273, 0], [58, 1351, 0]], "bbox": [0, 701, 251, 1356]}',
 'image_id': '0034',
 'object_id': '34',
 'sample_idx': 1,
 'granularity': 'fine',
 'image_path': '/content/MicroMat-mini/img/0034.png',
 'mask_path': '/content/MicroMat-mini/mask/0034_34.png',
 'prompt_path': '/content/MicroMat-mini/prompt/0034_34.json'}

Les prompts sont des chaînes de dictionnaires, vous pouvez donc obtenir les boîtes englobantes comme indiqué ci-dessous.

import json

json.loads(train_ds["prompt"][0])["bbox"]
# [0, 701, 251, 1356]

Visualiser une image d’exemple, un prompt et un masque.

import matplotlib.pyplot as plt
import numpy as np

def show_mask(mask, ax):
    color = np.array([0.12, 0.56, 1.0, 0.6])
    mask = np.array(mask)
    h, w = mask.shape
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, 4)
    ax.imshow(mask_image)
    x0, y0, x1, y1 = eval(train_ds["prompt"][0])["bbox"]
    ax.add_patch(
        plt.Rectangle((x0, y0), x1 - x0, y1 - y0,
                      fill=False, edgecolor="lime", linewidth=2))

example = train_ds[0]
image = np.array(example["image"])
ground_truth_mask = np.array(example["mask"])

fig, ax = plt.subplots()
ax.imshow(image)
show_mask(ground_truth_mask, ax)
ax.set_title("Ground truth mask")
ax.set_axis_off()

plt.show()

Nous pouvons maintenant définir notre jeu de données pour le chargement des données. SAMDataset encapsule notre jeu de données et met en forme chaque échantillon de la manière attendue par le processeur SAM. Ainsi, au lieu d’images et de masques bruts, vous obtenez des images traitées, des boîtes englobantes et des masques de référence prêts pour l’entraînement.

Par défaut, le processeur redimensionne les images, il renvoie donc, en plus des images et des masques, les dimensions d'origine. Il est également nécessaire de binariser le masque, car ses valeurs sont comprises dans l'intervalle [0, 255].

from torch.utils.data import Dataset
import torch

class SAMDataset(Dataset):
  def __init__(self, dataset, processor):
    self.dataset = dataset
    self.processor = processor

  def __len__(self):
    return len(self.dataset)

  def __getitem__(self, idx):
    item = self.dataset[idx]
    image = item["image"]
    prompt = eval(item["prompt"])["bbox"]
    inputs = self.processor(image, input_boxes=[[prompt]], return_tensors="pt")
    inputs["ground_truth_mask"] = (np.array(item["mask"]) > 0).astype(np.float32)
    inputs["original_image_size"] = torch.tensor(image.size[::-1])


    return inputs

Nous pouvons initialiser le processeur et le jeu de données à l'aide de celui-ci.

from transformers import Sam2Processor

processor = Sam2Processor.from_pretrained("facebook/sam2.1-hiera-small")
train_dataset = SAMDataset(dataset=train_ds, processor=processor)
``` 

We need to define a data collator that will turn varying size of ground truth masks to batches of reshaped masks in same shape. We reshape them using nearest neighbor interpolation. We also make batched tensors for rest of the elements in the batch. If your masks are all of same size, feel free to skip this step.

```python
import torch.nn.functional as F

def collate_fn(batch, target_hw=(256, 256)):

    pixel_values = torch.cat([item["pixel_values"] for item in batch], dim=0)
    original_sizes = torch.stack([item["original_sizes"] for item in batch])
    input_boxes = torch.cat([item["input_boxes"] for item in batch], dim=0)
    ground_truth_masks = torch.cat([
        F.interpolate(
            torch.as_tensor(x["ground_truth_mask"]).unsqueeze(0).unsqueeze(0).float(),
            size=(256, 256),
            mode="nearest"
        )
        for x in batch
    ], dim=0).long()

    return {
        "pixel_values": pixel_values,
        "original_sizes": original_sizes,
        "input_boxes": input_boxes,
        "ground_truth_mask": ground_truth_masks,
        "original_image_size": torch.stack([item["original_image_size"] for item in batch]),
    }

from torch.utils.data import DataLoader
train_dataloader = DataLoader(
    train_dataset,
    batch_size=4,
    shuffle=True,
    collate_fn=collate_fn,
)

Examinons ce que le chargeur de données produit.

batch = next(iter(train_dataloader))
for k,v in batch.items():
  print(k,v.shape)

# pixel_values torch.Size([4, 3, 1024, 1024])
# original_sizes torch.Size([4, 1, 2])
# input_boxes torch.Size([4, 1, 4])
# ground_truth_mask torch.Size([4, 1, 256, 256])
#original_image_size torch.Size([4, 2])

Nous allons maintenant charger le modèle et geler le vision et l'encodeur de prompt afin de n'entraîner que le décodeur de masque.

from transformers import Sam2Model

model = Sam2Model.from_pretrained("facebook/sam2.1-hiera-small")

for name, param in model.named_parameters():
  if name.startswith("vision_encoder") or name.startswith("prompt_encoder"):
    param.requires_grad_(False)
``` 

We can now define the optimizer and the loss function.
```python 
from torch.optim import Adam
import monai

optimizer = Adam(model.mask_decoder.parameters(), lr=1e-5, weight_decay=0)
seg_loss = monai.losses.DiceCELoss(sigmoid=True, squared_pred=True, reduction='mean')

Voyons comment le modèle se comporte avant l’entraînement.

import matplotlib.pyplot as plt

item = val_ds[1]
img = item["image"]
bbox = json.loads(item["prompt"])["bbox"]
inputs = processor(images=img, input_boxes=[[bbox]], return_tensors="pt").to(model.device)

with torch.no_grad():
  outputs = model(**inputs)

masks = processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"])[0]
preds = masks.squeeze(0)
mask = (preds[0] > 0).cpu().numpy()

overlay = np.asarray(img, dtype=np.uint8).copy()
overlay[mask] = 0.55 * overlay[mask] + 0.45 * np.array([0, 255, 0], dtype=np.float32)

plt.imshow(overlay)
plt.axis("off")
plt.show()

SAM2 result after training

Nous devons consigner nos prédictions dans trackio afin de pouvoir suivre l’amélioration du modèle en cours d’entraînement.

from PIL import Image
import trackio
import json


@torch.no_grad()
def predict_fn(img, bbox):

  inputs = processor(images=img, input_boxes=[[bbox]], return_tensors="pt").to(model.device)

  with torch.no_grad():
      outputs = model(**inputs)

  masks = processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"])[0]
  return masks

def log_eval_masks_trackio(dataset, indices, step, predict_fn,  project=None, sample_cap=8):
    logs = {"eval/step": int(step)}
    for idx in indices[:sample_cap]:
        item = dataset[idx] 
        img = item["image"]
        bbox = json.loads(item["prompt"])["bbox"]
        preds = predict_fn(img, bbox)
        preds = preds.squeeze(0)
        mask = (preds[0] > 0).cpu().numpy()  

        overlay = np.asarray(img, dtype=np.uint8).copy()
        overlay[mask] = 0.55 * overlay[mask] + 0.45 * np.array([0, 255, 0], dtype=np.float32)
        logs[f"{idx}/overlay"] = trackio.Image(overlay, caption="overlay")
        
    trackio.log(logs)

Nous pouvons maintenant écrire notre boucle d'entraînement et lancer l'entraînement !

Remarquez que nous enregistrons notre perte et nos masques d'évaluation avec trackio.

from tqdm import tqdm
from statistics import mean
import trackio
import torch

num_epochs = 30

device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else "cpu"
model.to(device)

model.train()
trackio.init(project="mask-eval")
for epoch in range(num_epochs):
    epoch_losses = []
    for batch in tqdm(train_dataloader):
      outputs = model(pixel_values=batch["pixel_values"].to(device),
                      input_boxes=batch["input_boxes"].to(device),
                      multimask_output=False)

      predicted_masks = outputs.pred_masks.squeeze(1)
      ground_truth_masks = batch["ground_truth_mask"].float().to(device)
      loss = seg_loss(predicted_masks, ground_truth_masks)

      optimizer.zero_grad()
      loss.backward()

      optimizer.step()
      epoch_losses.append(loss.item())
      
    log_eval_masks_trackio(dataset=val_ds, indices=[0, 3, 6, 9], step=epoch, predict_fn=predict_fn, project="mask-eval")
    print(f'Epoch: {epoch}')
    print(f'Mean loss: {mean(epoch_losses)}')
    trackio.log({"loss": mean(epoch_losses)})

trackio.finish()

Passons le modèle entraîné à l’épreuve.

import matplotlib.pyplot as plt

item = val_ds[1]
img = item["image"]
bbox = json.loads(item["prompt"])["bbox"]

inputs = processor(images=img, input_boxes=[[bbox]], return_tensors="pt").to(model.device)

with torch.no_grad():
  outputs = model(**inputs)

preds = processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"])[0]

preds = preds.squeeze(0)
mask = (preds[0] > 0).cpu().numpy()

overlay = np.asarray(img, dtype=np.uint8).copy()
overlay[mask] = 0.55 * overlay[mask] + 0.45 * np.array([0, 255, 0], dtype=np.float32)

plt.imshow(overlay)
plt.axis("off")
plt.show()

Grosse amélioration après seulement 20 époques d'entraînement sur un petit jeu de données !