Hugging Face's logo

Transformers documentation

マスクの生成

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v5.22.2).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

マスクの生成

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

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

  • Prompting モード: このモードでは、モデルは画像とプロンプトを入力として受け取ります。プロンプトは、対象オブジェクト内の2次元の点位置（XY 座標）または、オブジェクトを囲むバウンディングボックスです。Prompting モードでは、モデルはプロンプトが指し示しているオブジェクトのマスクのみを返します。
  • Segment Everythingモード: Segment Everythingでは、画像が与えられると、モデルは画像内のすべてのマスクを生成します。そのために、推論用に点のグリッドを生成し、画像上に重ねます。
  • Video Inference: モデルは動画と、動画フレーム内の点またはボックスのプロンプトを受け取り、その対象を動画全体にわたってトラッキングします。Video Inferenceの実行方法については、SAM 2 docsを参照してください。

Segment Anything Model (SAM) および Segment Anything Model 2 (SAM2) によってマスク生成タスクがサポートされており、Segment Anything Model 2 (SAM2) によって動画推論がサポートされています。SAM は、Vision Transformer ベースの画像エンコーダ、プロンプトエンコーダ、および双方向トランスフォーマーマスクデコーダから構成される強力なモデルです。画像とプロンプトがエンコードされ、デコーダはこれらの埋め込みを受け取って有効なマスクを生成します。一方、SAM 2 はマスクを追跡するためのメモリモジュールを追加することで SAM を拡張しています。

SAM Architecture

SAMは、大規模なデータカバレッジを有するセグメンテーション用の強力な基盤モデルとして機能します。SA-1Bという、100万枚の画像と11億個のマスクを含むデータセットで学習されています。

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

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

まず、transformersをインストールします。

pip install -q transformers

マスク生成パイプライン

マスク生成モデルの推論を行う最も簡単な方法は、mask-generation パイプラインを使用することです。

>>> from transformers import pipeline

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

画像を見てみましょう。

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

すべてをセグメントしましょう。points-per-batch は、すべてをセグメントするモードにおけるポイントの並列推論を有効にします。これにより推論が高速化されますが、メモリ消費量が増加します。また、SAM は画像ではなく、ポイントに対してのみバッチ処理を有効にします。pred_iou_thresh は IoU の信頼度閾値であり、この閾値を超えるマスクのみが返されます。

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

masks は次のようになります。

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

これらは次のように可視化できます。

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

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

Visualized

モデル推論

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

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

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を渡します。プロセッサーが画像をリサイズするため、出力を補間する必要があるため、これらを渡す必要があります。

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つのマスクを視覚化できます。

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

Box Prompting

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

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

以下に示すように、ハチの周囲のバウンディングボックスを視覚化できます。

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

以下に推論の出力を確認できます。

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

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

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

画像マッティングのために、MicroMat データセットの一部で SAM2.1 をファインチューニングします。DICE loss を使用するには monai ライブラリのインストールが必要であり、トレーニング中のマスクのログ記録には trackio が必要です。

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

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

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

サンプルを見てみましょう。

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

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

import json

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

例の画像、プロンプト、マスクを可視化します。

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] の範囲にあるため、マスクを二値化する必要があります。

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

プロセッサーとデータセットをこれを使って初期化できます。

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

データローダーが生成する内容を確認してみましょう。

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

次に、モデルを読み込み、マスクデコーダーのみを学習するために、ビジョンとプロンプトエンコーダーをフリーズします。

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

学習前にモデルの性能を確認してみましょう。

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

トレーニングの途中でモデルの改善を監視できるようにするため、予測結果をtrackioに記録する必要があります。

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 で記録している点に注目してください。

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

学習したモデルをテストしてみましょう。

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

Update on GitHub

←画像特徴抽出 キーポイント検出→