# マスクの生成

マスク生成は、画像に対して意味のあるマスクを生成するタスクです。このタスクは[画像セグメンテーション](semantic_segmentation)と非常に似ていますが、多くの違いがあります。画像セグメンテーションモデルはラベル付きデータセットで学習され、学習中に確認したクラスに限定されます。画像が与えられると、マスクのセットと対応するクラスを返します。

マスク生成モデルは大量のデータで学習され、2つのモードで動作します。

- Prompting mode: このモードでは、モデルは画像とプロンプトを入力として受け取ります。プロンプトは、画像内のオブジェクト内の2次元の点の位置（XY座標）または、オブジェクトを囲むバウンディングボックスです。Prompting modeでは、モデルはオブジェクト上のマスクのみを返します。
プロンプトが指摘している点です。
- Segment Everythingモード: Segment Everythingでは、画像が与えられると、モデルは画像内のすべてのマスクを生成します。そのために、推論用に点のグリッドを生成し、画像上に重ねます。
- Video Inference: モデルは動画と、動画フレーム内の点またはボックスのプロンプトを受け取り、その対象を動画全体にわたってトラッキングします。Video Inferenceの実行方法については、[SAM 2 docs](../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/)という、100万枚の画像と11億個のマスクを含むデータセットで学習されています。

このガイドでは、次のことを学びます。

- バッチ処理を有効にして、セグメントエブリシングモードで推論を実行します。
- ポイントプロンプティングモードで推論を実行します。
- ボックスプロンプティングモードで推論を実行します。

まず、`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()
```

以下は、グレースケールで表示された元の画像にカラフルなマップを重ねたものです。非常に印象的です。

     

## モデル推論

### ポイントプロンプティング

パイプラインなしでモデルを使用することもできます。そのためには、モデルとプロセッサーを初期化します。

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

ポイントプロンプティングを行うには、入力ポイントをプロセッサに渡し、プロセッサの出力をモデルに渡して推論を実行します。モデルの出力を後処理するには、出力と、プロセッサの初期出力から取得した`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` の出力で、3つのマスクを視覚化できます。

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

ポイントプロンプティングと同様の方法で、ボックスプロンプティングも実行できます。`processor`に画像とともに、`[x_min, y_min, x_max, y_max]`形式のリストとして入力ボックスを渡すだけです。プロセッサの出力を取得してモデルに直接渡し、その後、出力を再度後処理します。

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

     

## マスク生成のためのファインチューニング

画像マッティングのために、MicroMat データセットの一部で SAM2.1 をファインチューニングします。DICE loss を使用するには [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
#})
```

image、mask、prompt の列が必要です。train と test に分割します。

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

プロンプトは辞書の文字列なので、以下のようにバウンディングボックスを取得できます。

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

サイズが異なる正解マスクを、同じ形状にリシェイプされたマスクのバッチに変換するデータコラテーターを定義する必要があります。リシェイプには最近傍補間を使用します。また、バッチ内の他の要素についてもバッチ化されたテンソルを作成します。マスクがすべて同じサイズである場合は、このステップをスキップして構いません。

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

これでトレーニングループを書き、トレーニングを実行できます。

loss と評価マスクを 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)