Tạo mặt nạ

Việc tạo mask là nhiệm vụ tạo ra các mask có ý nghĩa ngữ nghĩa cho một hình ảnh. Nhiệm vụ này rất giống với phân đoạn hình ảnh, nhưng tồn tại nhiều điểm khác biệt. Các mô hình phân đoạn hình ảnh được huấn luyện trên các bộ dữ liệu có nhãn và bị giới hạn trong các lớp mà chúng đã thấy trong quá trình huấn luyện; khi nhận vào một hình ảnh, chúng trả về một tập hợp các mask cùng với các lớp tương ứng.

Các mô hình tạo mặt nạ được huấn luyện trên một lượng lớn dữ liệu và hoạt động ở hai chế độ.

Chế độ prompting: Ở chế độ này, mô hình nhận vào một hình ảnh và một prompt, trong đó prompt có thể là vị trí điểm 2D (tọa độ XY) trong hình ảnh nằm bên trong một đối tượng hoặc một bounding box bao quanh một đối tượng. Ở chế độ prompting, mô hình chỉ trả về mask trên đối tượng mà prompt đang chỉ ra.
Chế độ Segment Everything: Trong chế độ segment everything, khi được cung cấp một hình ảnh, mô hình sẽ tạo ra mọi mask trong hình ảnh đó. Để thực hiện điều này, một lưới các điểm sẽ được tạo ra và phủ lên hình ảnh để suy luận.
Video Inference: Mô hình nhận vào một video và một prompt dạng điểm hoặc hộp trong một khung hình của video, sau đó theo dõi đối tượng xuyên suốt video. Bạn có thể tìm hiểu thêm về cách thực hiện video inference bằng cách xem tài liệu SAM 2.

Nhiệm vụ tạo mask được hỗ trợ bởi Segment Anything Model (SAM) và Segment Anything Model 2 (SAM2), trong khi suy luận video được hỗ trợ bởi Segment Anything Model 2 (SAM2). SAM là một mô hình mạnh mẽ bao gồm bộ mã hóa ảnh dựa trên Vision Transformer, bộ mã hóa prompt và bộ giải mã mask hai chiều. Ảnh và prompt được mã hóa, và bộ giải mã nhận các embedding này để tạo ra các mask hợp lệ. Trong khi đó, SAM 2 mở rộng SAM bằng cách thêm một module bộ nhớ để theo dõi các mask.

SAM đóng vai trò là một mô hình nền tảng mạnh mẽ cho phân đoạn nhờ phạm vi dữ liệu lớn. Nó được huấn luyện trên SA-1B, một tập dữ liệu gồm 1 triệu hình ảnh và 1,1 tỷ mặt nạ.

Trong hướng dẫn này, bạn sẽ học cách:

Suy luận ở chế độ phân đoạn mọi thứ với batching,
Suy luận ở chế độ gợi ý điểm,
Suy luận ở chế độ gợi ý trong hộp.

Đầu tiên, hãy cài đặt transformers:

pip install -q transformers

Đường ống tạo mặt nạ

Cách dễ nhất để suy luận các mô hình tạo mặt nạ là sử dụng pipeline mask-generation.

>>> from transformers import pipeline

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

Hãy xem hình ảnh.

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

Hãy phân đoạn mọi thứ. points-per-batch cho phép suy luận song song các điểm trong chế độ phân đoạn mọi thứ. Điều này giúp suy luận nhanh hơn nhưng tiêu tốn nhiều bộ nhớ hơn. Hơn nữa, SAM chỉ cho phép gom nhóm (batching) trên các điểm chứ không phải trên các hình ảnh. pred_iou_thresh là ngưỡng độ tin cậy IoU, chỉ các mặt nạ (mask) vượt qua ngưỡng đó mới được trả về.

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

masks trông như sau:

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

Chúng ta có thể trực quan hóa chúng như sau:

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

Dưới đây là hình ảnh gốc ở thang xám với các bản đồ màu sắc được phủ lên. Rất ấn tượng.

Suy luận Mô hình

Điểm Gợi ý

Bạn cũng có thể sử dụng model mà không cần pipeline. Để làm điều này, hãy khởi tạo model và processor.

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

Để thực hiện point prompting, hãy truyền điểm đầu vào vào processor, sau đó lấy đầu ra của processor và truyền nó vào model để suy luận. Để hậu xử lý đầu ra của model, hãy truyền các outputs và original_sizes được lấy từ đầu ra ban đầu của processor. Chúng ta cần truyền những giá trị này vì processor đã thay đổi kích thước ảnh, và đầu ra cần được ngoại suy.

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

Chúng ta có thể trực quan hóa ba mặt nạ trong đầu ra 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

Bạn cũng có thể thực hiện box prompting theo cách tương tự như point prompting. Bạn chỉ cần truyền vào hộp đầu vào ở định dạng danh sách [x_min, y_min, x_max, y_max] cùng với hình ảnh cho processor. Lấy kết quả đầu ra của processor và truyền trực tiếp vào model, sau đó xử lý hậu kỳ kết quả đầu ra một lần nữa.

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

Bạn có thể trực quan hóa hộp bao quanh con ong như được hiển thị bên dưới.

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

Bạn có thể xem kết quả suy luận bên dưới.

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

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

Tinh chỉnh cho việc tạo mặt nạ

Chúng ta sẽ fine-tune SAM2.1 trên một phần nhỏ của bộ dữ liệu MicroMat cho bài toán image matting. Chúng ta cần cài đặt thư viện monai để sử dụng DICE loss, và trackio để ghi log các mask trong quá trình huấn luyện.

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

Chúng ta cần các cột image, mask và prompt. Chúng ta tách dữ liệu để huấn luyện và kiểm tra.

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

Hãy xem một ví dụ.

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 là chuỗi các từ điển, vì vậy bạn có thể lấy các hộp bao (bounding boxes) như minh họa bên dưới.

import json

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

Hiển thị một hình ảnh, prompt và mask mẫu.

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

Bây giờ chúng ta có thể định nghĩa dataset để nạp dữ liệu. SAMDataset bọc dataset của chúng ta và định dạng từng mẫu theo cách mà SAM processor mong đợi. Vì vậy, thay vì các hình ảnh và mask thô, bạn sẽ nhận được các hình ảnh đã xử lý, bounding box và ground-truth mask sẵn sàng cho việc huấn luyện.

Mặc định, processor sẽ thay đổi kích thước ảnh, vì vậy ngoài ảnh và mask, nó còn trả về kích thước gốc. Chúng ta cũng cần nhị phân hóa mask vì nó có các giá trị trong khoảng [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

Chúng ta có thể khởi tạo bộ xử lý và bộ dữ liệu bằng nó.

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

Hãy xem trình tải dữ liệu trả về những gì.

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

Chúng ta sẽ tải mô hình và đóng băng bộ mã hóa thị giác cùng bộ mã hóa prompt để chỉ huấn luyện bộ giải mã mặt nạ.

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

Hãy xem hiệu suất của mô hình trước khi huấn luyện.

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

Chúng ta cần ghi log các dự đoán của mình vào trackio để có thể theo dõi sự cải thiện của mô hình trong quá trình huấn luyện.

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)

Bây giờ chúng ta có thể viết vòng lặp huấn luyện và bắt đầu huấn luyện!

Hãy chú ý cách chúng ta ghi log loss và evaluation masks bằng 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()

Hãy thử nghiệm mô hình đã được huấn luyện.

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

Cải thiện đáng kể sau khi chỉ huấn luyện 20 epochs trên một tập dữ liệu nhỏ!