# Geração de máscaras

A geração de máscaras é a tarefa de gerar máscaras semanticamente significativas para uma imagem. Essa tarefa é muito semelhante à [segmentação de imagem](semantic_segmentation), mas existem muitas diferenças. Os modelos de segmentação de imagem são treinados em conjuntos de dados rotulados e são limitados às classes que viram durante o treinamento; eles retornam um conjunto de máscaras e as classes correspondentes, dada uma imagem.

Os modelos de geração de máscaras são treinados com grandes quantidades de dados e operam em dois modos.

- Modo de prompt: neste modo, o modelo recebe uma imagem e um prompt, em que um prompt pode ser uma localização de ponto 2D (coordenadas XY) na imagem dentro de um objeto ou uma caixa delimitadora ao redor de um objeto. No modo de prompt, o modelo retorna apenas a máscara sobre o objeto
que o prompt está apontando.
- Modo Segment Everything: no segment everything, dado uma imagem, o modelo gera todas as máscaras presentes na imagem. Para isso, uma grade de pontos é gerada e sobreposta à imagem para a inferência.
- Inferência de vídeo: o modelo aceita um vídeo e um prompt de ponto ou caixa em um frame do vídeo, que é rastreado ao longo de todo o vídeo. Você pode obter mais informações sobre como realizar a inferência de vídeo seguindo a [documentação do SAM 2](../model_doc/sam2).

A tarefa de geração de máscaras é suportada por [Segment Anything Model (SAM)](../model_doc/sam) e [Segment Anything Model 2 (SAM2)](../model_doc/sam2), enquanto a inferência de vídeo é suportada por [Segment Anything Model 2 (SAM2)](../model_doc/sam2). O SAM é um modelo poderoso que consiste em um codificador de imagem baseado em Vision Transformer, um codificador de prompts e um decodificador de máscaras com transformador bidirecional. As imagens e os prompts são codificados, e o decodificador recebe esses embeddings e gera máscaras válidas. Enquanto isso, o SAM 2 estende o SAM adicionando um módulo de memória para rastrear as máscaras.

     

O SAM funciona como um modelo de fundação poderoso para segmentação, pois possui uma ampla cobertura de dados. Ele é treinado no [SA-1B](https://ai.meta.com/datasets/segment-anything/), um conjunto de dados com 1 milhão de imagens e 1,1 bilhão de máscaras.

Neste guia, você aprenderá como:

- Faça inferência no modo segmentar tudo com loteamento,
- Infira no modo de prompt de ponto,
- Infira no modo de prompt em caixa.

Primeiro, vamos instalar o `transformers`:

```bash
pip install -q transformers
```

## Pipeline de Geração de Máscaras

A maneira mais fácil de inferir modelos de geração de máscaras é usar o pipeline `mask-generation`.

```python
>>> from transformers import pipeline

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

Vamos ver a imagem.

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

     

Vamos segmentar tudo. `points-per-batch` permite a inferência paralela de pontos no modo segment everything. Isso permite uma inferência mais rápida, mas consome mais memória. Além disso, o SAM só permite o agrupamento (batching) de pontos e não das imagens. `pred_iou_thresh` é o limiar de confiança de IoU, no qual apenas as máscaras acima desse determinado limiar são retornadas.

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

As `masks` têm o seguinte formato:

```bash
{'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 visualizá-los da seguinte forma:

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

Abaixo está a imagem original em escala de cinza com mapas coloridos sobrepostos. Muito impressionante.

     

## Inferência de Modelo

### Prompting de Pontos

Você também pode usar o modelo sem o pipeline. Para isso, inicialize o modelo e o processador.

```python
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 fazer point prompting, passe o ponto de entrada para o processador, em seguida, pegue a saída do processador e passe-a para o modelo para inferência. Para pós-processar a saída do modelo, passe as saídas e `original_sizes` são obtidos da saída inicial do processador. Precisamos passar esses valores porque o processador redimensiona a imagem e a saída precisa ser extrapolada.

```python
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 as três máscaras na saída `masks`.

```python
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

Você também pode fazer box prompting de forma semelhante ao point prompting. Basta passar a caixa de entrada no formato de uma lista `[x_min, y_min, x_max, y_max]` junto com a imagem para o `processor`. Pegue a saída do processor e passe diretamente para o modelo, depois faça o pós-processamento da saída novamente.

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

Você pode visualizar a caixa delimitadora ao redor da abelha, conforme mostrado abaixo.

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

     

Você pode ver a saída de inferência abaixo.

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

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

     

## Fine-tuning para geração de máscaras

Vamos fazer o fine-tuning do SAM2.1 em uma pequena parte do conjunto de dados MicroMat para matting de imagem. Precisamos instalar a biblioteca [monai](https://github.com/Project-MONAI/MONAI) para usar a perda DICE e o [trackio](https://huggingface.co/docs/trackio/index) para registrar as máscaras durante o treinamento.

```bash
pip install -q datasets monai trackio
``` 
Agora podemos carregar nosso conjunto de dados e dar uma olhada.

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

Precisamos das colunas de imagem, máscara e prompt. Fazemos a divisão para treino e teste.

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

Vamos dar uma olhada em uma amostra.

```python
train_ds[0]
```

```json
 {'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'}
```

Prompts são strings de dicionários, portanto, você pode obter as bounding boxes conforme mostrado abaixo.

```python
import json

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

Visualize uma imagem de exemplo, um prompt e uma máscara.

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

Agora podemos definir nosso dataset para carregar os dados. O SAMDataset encapsula nosso dataset e formata cada amostra da maneira que o processador do SAM espera. Assim, em vez de imagens e máscaras brutas, você obtém imagens processadas, bounding boxes e máscaras de ground-truth prontas para o treinamento.

Por padrão, o processor redimensiona as imagens, portanto, além das imagens e máscaras, ele também retorna os tamanhos originais. Também precisamos binarizar a máscara, pois ela possui valores [0, 255].

```python
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 o processador e o conjunto de dados com ele.

```python
from transformers import Sam2Processor

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

Precisamos definir um data collator que converterá máscaras de ground truth de tamanhos variados em lotes de máscaras redimensionadas com a mesma forma. Redimensionamos-as usando interpolação de vizinho mais próximo. Também criamos tensores em lote para os demais elementos do lote. Se suas máscaras forem todas do mesmo tamanho, sinta-se à vontade para pular esta etapa.

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

Vamos dar uma olhada no que o data loader produz.

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

Agora, carregaremos o modelo e congelaremos o codificador de visão e o codificador de prompt para treinar apenas o decodificador de máscara.

```python
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)
``` 

Agora podemos definir o otimizador e a função de perda.
```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')
```

Vamos ver como o modelo se comporta antes do treinamento.

```python
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](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/sam2_before_training.png)

Precisamos registrar nossas previsões no trackio para que possamos monitorar a melhoria do modelo no meio do treinamento.

```python
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)
```

Agora podemos escrever nosso loop de treinamento e treinar!

Observe como registramos nossa perda e as máscaras de avaliação com trackio.

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

Vamos colocar o modelo treinado à prova.

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

Melhora significativa após apenas 20 épocas de treinamento em um pequeno conjunto de dados!

![SAM2 result after training](https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/sam2_after_training.png)