Generación de máscaras

La generación de máscaras es la tarea de generar máscaras con significado semántico para una imagen. Esta tarea es muy similar a la segmentación de imágenes, pero existen muchas diferencias. Los modelos de segmentación de imágenes se entrenan con conjuntos de datos etiquetados y están limitados a las clases que han visto durante el entrenamiento; devuelven un conjunto de máscaras y las clases correspondientes, dada una imagen.

Los modelos de generación de máscaras se entrenan con grandes cantidades de datos y operan en dos modos.

Modo de indicación: en este modo, el modelo recibe una imagen y una indicación, donde la indicación puede ser una ubicación de punto 2D (coordenadas XY) dentro de un objeto en la imagen o un recuadro delimitador que rodea un objeto. En el modo de indicación, el modelo solo devuelve la máscara sobre el objeto al que la indicación hace referencia.
Modo Segment Everything: en segment everything, a partir de una imagen, el modelo genera todas las máscaras presentes en la imagen. Para ello, se genera una cuadrícula de puntos que se superpone a la imagen para la inferencia.
Inferencia de video: el modelo acepta un video y un punto o caja de entrada en un fotograma del video, que se rastrea a lo largo de todo el video. Puede obtener más información sobre cómo realizar la inferencia de video siguiendo la documentación de SAM 2.

La tarea de generación de máscaras es compatible con Segment Anything Model (SAM) y Segment Anything Model 2 (SAM2), mientras que la inferencia de vídeo es compatible con Segment Anything Model 2 (SAM2). SAM es un modelo potente que consta de un codificador de imágenes basado en Vision Transformer, un codificador de prompts y un decodificador de máscaras bidireccional. Las imágenes y los prompts se codifican, y el decodificador toma estos embeddings y genera máscaras válidas. Mientras tanto, SAM 2 extiende SAM añadiendo un módulo de memoria para rastrear las máscaras.

SAM funciona como un modelo de base potente para la segmentación, ya que cuenta con una amplia cobertura de datos. Se entrena en SA-1B, un conjunto de datos con 1 millón de imágenes y 1.100 millones de máscaras.

En esta guía aprenderás a:

Inferir en el modo segmentar todo con lotes,
Inferir en el modo de indicación por puntos,
Inferir en el modo de indicación en caja.

Primero, instalemos transformers:

pip install -q transformers

Pipeline de generación de máscaras

La forma más sencilla de realizar inferencia con modelos de generación de máscaras es utilizar la pipeline mask-generation.

>>> from transformers import pipeline

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

Veamos la imagen.

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")

Segmentemos todo. points-per-batch permite la inferencia paralela de puntos en el modo segment everything. Esto permite una inferencia más rápida, pero consume más memoria. Además, SAM solo permite el agrupamiento (batching) sobre puntos y no sobre imágenes. pred_iou_thresh es el umbral de confianza de IoU, donde solo se devuelven las máscaras que superan dicho umbral.

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

Las masks son las siguientes:

{'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])}

Podemos visualizarlos de la siguiente manera:

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()

A continuación se muestra la imagen original en escala de grises con mapas de colores superpuestos. Muy impresionante.

Inferencia del modelo

Point Prompting

También puedes usar el modelo sin la pipeline. Para ello, inicializa el modelo y el procesador.

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")

Para realizar point prompting, pase el punto de entrada al procesador y luego tome la salida del procesador y pásela al modelo para la inferencia. Para postprocesar la salida del modelo, pase las salidas y original_sizes se toman de la salida inicial del procesador. Es necesario pasar estos valores porque el procesador redimensiona la imagen y la salida debe ser extrapolada.

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())

Podemos visualizar las tres máscaras en la salida 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

También puedes realizar box prompting de manera similar a point prompting. Simplemente pasa la caja de entrada en el formato de lista [x_min, y_min, x_max, y_max] junto con la imagen al processor. Toma la salida del processor y pásala directamente al modelo, y luego realiza el postprocesado de la salida nuevamente.

# 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()

Puede visualizar el recuadro delimitador alrededor de la abeja, como se muestra a continuación.

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()

Puede ver la salida de inferencia a continuación.

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

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

Ajuste fino para la generación de máscaras

Ajustaremos finamente SAM2.1 sobre una pequeña parte del conjunto de datos MicroMat para el recorte de imágenes. Necesitamos instalar la biblioteca monai para usar la pérdida DICE, y trackio para registrar las máscaras durante el entrenamiento.

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
#})

Necesitamos columnas de imagen, máscara y prompt. Dividimos para entrenamiento y prueba.

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

Veamos un ejemplo.

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'}

Los prompts son cadenas de diccionarios, por lo que puedes obtener las cajas delimitadoras como se muestra a continuación.

import json

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

Visualiza una imagen de ejemplo, un prompt y una máscara.

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()

Ahora podemos definir nuestro dataset para cargar los datos. SAMDataset envuelve nuestro dataset y formatea cada muestra de la manera que el procesador de SAM espera. Así, en lugar de imágenes y máscaras sin procesar, obtienes imágenes procesadas, bounding boxes y máscaras de ground truth listas para el entrenamiento.

Por defecto, el procesador redimensiona las imágenes, por lo que, además de las imágenes y las máscaras, también devuelve los tamaños originales. También necesitamos binarizar la máscara, ya que tiene valores en el rango [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

Podemos inicializar el procesador y el conjunto de datos con él.

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,
)

Veamos lo que produce el cargador de datos.

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])

Ahora cargaremos el modelo y congelaremos el codificador de visión y el codificador del prompt para entrenar únicamente el decodificador de la máscara.

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')

Veamos cómo rinde el modelo antes del entrenamiento.

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

Necesitamos registrar nuestras predicciones en trackio para poder monitorear la mejora del modelo a mitad del entrenamiento.

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)

Ahora podemos escribir nuestro bucle de entrenamiento ¡y entrenar!

Observa cómo registramos nuestra pérdida y las máscaras de evaluación con 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()

Pongamos a prueba el modelo entrenado.

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()

¡Gran mejora después de entrenar solo 20 épocas en un conjunto de datos pequeño!