Transformers documentation
Генерация масок
Генерация масок
Генерация масок — это задача создания семантически значимых масок для изображения. Эта задача очень похожа на сегментацию изображений, но между ними есть существенные различия. Модели сегментации изображений обучаются на размеченных наборах данных и ограничены классами, которые они видели в процессе обучения; по заданному изображению они возвращают набор масок и соответствующие им классы.
Модели генерации масок обучаются на больших объёмах данных и работают в двух режимах.
- Режим подсказок (prompting): в этом режиме модель принимает изображение и подсказку, где подсказкой может быть точка (координаты XY) внутри объекта или ограничивающая рамка (bounding box) вокруг объекта. В режиме подсказок модель возвращает только маску объекта, на который указывает подсказка.
- Режим Segment Everything: в режиме segment everything, по заданному изображению модель генерирует все маски на изображении. Для этого создаётся сетка точек, которая накладывается на изображение для выполнения инференса.
- Видеоинференс: модель принимает видео и точку или рамку в кадре, которые отслеживаются на протяжении всего видео. Подробнее о том, как выполнять видеоинференс, можно узнать в документации SAM 2.
Задача генерации масок поддерживается моделями Segment Anything Model (SAM) и Segment Anything Model 2 (SAM2), а инференс для видео — моделью Segment Anything Model 2 (SAM2). SAM — это мощная модель, состоящая из кодировщика изображений на основе Vision Transformer, кодировщика промптов и двунаправленного трансформерного декодера масок. Изображения и промпты кодируются, а декодер принимает эти эмбеддинги и генерирует валидные маски. В свою очередь, SAM 2 расширяет SAM за счёт добавления модуля памяти для отслеживания масок.
SAM является мощной базовой моделью для сегментации благодаря широкому охвату данных. Она обучена на SA-1B — наборе данных, содержащем 1 миллион изображений и 1,1 миллиарда масок.
В этом руководстве вы узнаете, как:
- Выполняйте инференс в режиме segment everything с батчингом,
- Выполнить инференс в режиме point prompting,
- Выполнить инференс в режиме box prompting.
Сначала установим 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")
Разделим всё на сегменты. Параметр 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()Ниже представлена исходная изображение в оттенках серого с наложенными на неё цветными картами. Очень впечатляюще.
Вывод модели
Point Prompting
Вы также можете использовать модель без pipeline. Для этого инициализируйте модель и процессор.
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, которые берутся из начального выхода процессора. Эти значения необходимо передать, поскольку процессор изменяет размер изображения, а выход требует экстраполяции.
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.
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. Затем возьмите выход процессора и передайте его напрямую в модель, после чего выполните постобработку результата.
# 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()
Ниже вы можете увидеть вывод при инференсе.
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, а для логирования масок в процессе обучения — 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
#})Нам нужны столбцы с изображениями, масками и промптами. Мы разбиваем данные на обучающую и тестовую выборки.
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'}Prompts — это строки словарей, поэтому вы можете получить ограничивающие прямоугольники, как показано ниже.
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()
Нам нужно отправлять наши предсказания в 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)Теперь мы можем написать цикл обучения и запустить тренировку!
Обратите внимание, что мы записываем нашу потерю и маски оценки с помощью 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 эпох обучения на небольшом наборе данных!
