# Генерация масок

Генерация масок — это задача создания семантически значимых масок для изображения. Эта задача очень похожа на [сегментацию изображений](semantic_segmentation), но между ними есть существенные различия. Модели сегментации изображений обучаются на размеченных наборах данных и ограничены классами, которые они видели во время обучения; по заданному изображению они возвращают набор масок и соответствующие им классы.

Модели генерации масок обучаются на больших объёмах данных и работают в двух режимах.

- Режим подсказок: в этом режиме модель принимает изображение и подсказку, где подсказкой может быть точка (координаты XY) внутри объекта или ограничивающая рамка вокруг объекта. В режиме подсказок модель возвращает только маску объекта
на что указывает промпт.
- Режим Segment Everything: в режиме segment everything, по заданному изображению модель генерирует все маски на изображении. Для этого создаётся сетка точек, которая накладывается на изображение для выполнения инференса.
- Видеоинференс: модель принимает видео и точку или рамку в кадре, которые отслеживаются на протяжении всего видео. Подробнее о том, как выполнять видеоинференс, можно узнать в [документации SAM 2](../model_doc/sam2).

Задача генерации масок поддерживается моделями [Segment Anything Model (SAM)](../model_doc/sam) и [Segment Anything Model 2 (SAM2)](../model_doc/sam2), а инференс для видео — моделью [Segment Anything Model 2 (SAM2)](../model_doc/sam2). SAM — это мощная модель, состоящая из кодировщика изображений на основе Vision Transformer, кодировщика промптов и двунаправленного трансформерного декодера масок. Изображения и промпты кодируются, а декодер принимает эти эмбеддинги и генерирует валидные маски. В свою очередь, SAM 2 расширяет SAM за счёт добавления модуля памяти для отслеживания масок.

     

SAM является мощной базовой моделью для сегментации благодаря широкому охвату данных. Она обучена на [SA-1B](https://ai.meta.com/datasets/segment-anything/) — наборе данных, содержащем 1 миллион изображений и 1,1 миллиарда масок.

В этом руководстве вы узнаете, как:

- Выполняйте инференс в режиме segment everything с батчингом,
- Выполнить инференс в режиме point prompting,
- Выполнить инференс в режиме box prompting.

Сначала установим `transformers`:

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

## Конвейер генерации масок

Самый простой способ выполнить инференс моделей генерации масок — использовать пайплайн `mask-generation`.

```python
>>> from transformers import pipeline

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

Рассмотрим изображение.

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

     

Сегментируем всё. `points-per-batch` позволяет выполнять параллельный вывод точек в режиме сегментации всего. Это ускоряет вывод, но увеличивает потребление памяти. Кроме того, SAM поддерживает батчинг только для точек, но не для изображений. `pred_iou_thresh` — это порог уверенности по IoU, при котором возвращаются только маски, превышающие заданный порог.

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

Маски `masks` выглядят следующим образом:

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

Мы можем визуализировать их следующим образом:

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

Ниже представлена исходная изображение в оттенках серого с наложенными на неё цветными картами. Очень впечатляюще.

     

## Вывод модели

### Point Prompting

Вы также можете использовать модель без pipeline. Для этого инициализируйте модель и процессор.

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

Чтобы выполнить point prompting, передайте входную точку в процессор, затем возьмите выход процессора и передайте его в модель для инференса. Для постобработки выхода модели передайте выходные данные и `original_sizes`, которые берутся из начального выхода процессора. Эти значения необходимо передать, поскольку процессор изменяет размер изображения, а выход требует экстраполяции.

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

Мы можем визуализировать три маски в выводе `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

Вы также можете выполнять box prompting по аналогии с point prompting. Достаточно передать входной box в формате списка `[x_min, y_min, x_max, y_max]` вместе с изображением в `processor`. Затем возьмите выход процессора и передайте его напрямую в модель, после чего снова выполните постобработку результата.

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

Вы можете визуализировать ограничивающий прямоугольник вокруг пчелы, как показано ниже.

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

     

Ниже вы можете увидеть вывод при инференсе.

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

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

     

## Доработка для генерации масок

Мы будем дообучать SAM2.1 на небольшой части набора данных MicroMat для задачи маттинга изображений. Для использования функции потерь DICE необходимо установить библиотеку [monai](https://github.com/Project-MONAI/MONAI), а для логирования масок в процессе обучения — [trackio](https://huggingface.co/docs/trackio/index).

```bash
pip install -q datasets monai trackio
``` 
Теперь мы можем загрузить наш набор данных и посмотреть на него.

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

Нам нужны столбцы с изображениями, масками и промптами. Мы разбиваем данные на обучающую и тестовую выборки.

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

Рассмотрим пример.

```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 — это строки словарей, поэтому вы можете получить ограничивающие прямоугольники, как показано ниже.

```python
import json

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

Визуализируйте пример изображения, промпта и маски.

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

Теперь мы можем определить наш датасет для загрузки данных. SAMDataset оборачивает наш датасет и форматирует каждый образец так, как этого ожидает процессор SAM. Таким образом, вместо сырых изображений и масок вы получаете обработанные изображения, ограничивающие рамки и эталонные маски, готовые для обучения.

По умолчанию processor изменяет размер изображений, поэтому, помимо изображений и масок, он также возвращает исходные размеры. Кроме того, маску необходимо бинаризовать, так как её значения лежат в диапазоне [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
```

Мы можем инициализировать процессор и датасет с его помощью.

```python
from transformers import Sam2Processor

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

Необходимо определить data collator, который преобразует ground truth маски разного размера в батчи масок одинаковой формы. Мы изменяем их размер с помощью интерполяции ближайшего соседа. Также мы создаём батченные тензоры для остальных элементов в батче. Если все ваши маски имеют одинаковый размер, вы можете пропустить этот шаг.

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

Рассмотрим, что возвращает загрузчик данных.

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

Теперь мы загрузим модель и заморозим модуль зрения и кодировщик промпта, чтобы обучать только декодер маски.

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

Теперь мы можем определить оптимизатор и функцию потерь.
```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')
```

Посмотрим, как модель работает до обучения.

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

Нам нужно отправлять наши предсказания в trackio, чтобы отслеживать улучшение модели в процессе обучения.

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

Теперь мы можем написать цикл обучения и запустить тренировку!

Обратите внимание, что мы записываем нашу потерю и маски оценки с помощью 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()
```

Давайте проверим обученную модель в деле.

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

Отличное улучшение после всего лишь 20 эпох обучения на небольшом наборе данных!

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