Image segmentation

Mask trả lời một câu hỏi cho mỗi điểm

Có ba kiểu segmentation:

• Semantic: mỗi pixel một lớp, chẳng hạn nền, mực, giấy. Hai vật thể cùng lớp mang cùng ID.
• Instance: tách ID cho từng vật thể, dù chúng cùng lớp. Nếu cần đếm hai chai chạm nhau, semantic mask “chai” chưa đủ để tách chúng.
• Panoptic (toàn cảnh): gộp hai yêu cầu, mọi pixel có lớp, và những vật đếm được như người, xe còn có ID riêng.

Mask cũng là số đo: hồ 400 pixel, mỗi pixel ứng với ô $5 \times 5$ m, rộng khoảng 10 000 m².

Instance mask phải giữ ranh giới giữa các vật thể cùng lớp, nên cần model có khái niệm “từng vật”. Một cách phổ biến là Mask R-CNN: lấy detector hai giai đoạn Faster R-CNN của bài 5.5 và thêm một nhánh dự đoán mask nhị phân nhỏ, cỡ $28\times28$, bên trong mỗi box. Mỗi vật thể nhận lớp, box, score và mask riêng, nên hai chai chạm nhau vẫn là hai instance.

U-Net trong bài này là model semantic: mỗi pixel nhận một lớp, không có ID vật thể.

Với bài nhị phân vùng mực, đầu ra thường là một logit cho mỗi pixel: (N,1,H,W). BCEWithLogitsLoss nhận logits và target float cùng shape. Khi suy luận, sigmoid thành xác suất rồi threshold thành mask.

Với nhiều lớp loại trừ nhau, logits là (N,K,H,W), target là (N,H,W) long, dùng cross-entropy rồi argmax theo chiều lớp: 21 lớp trên ảnh $512 \times 512$ cho logits $(N, 21, 512, 512)$. Mask lưu dạng ảnh thường mang 0/255; phải đổi về 0/1 kiểu long trước cross-entropy, vì 255 vượt quá số lớp.

Nhãn segmentation đắt hơn nhãn cả ảnh: phải xác định từng pixel hoặc vùng. Biên có thể mơ hồ do bóng, mờ hoặc chính sách annotation. Model chỉ có thể học theo quy tắc nhãn đã cung cấp; cần xem vài cặp ảnh–mask ở độ phân giải thật.

Vì sao accuracy chưa đủ?

Giả sử 99% pixel là nền. Dự đoán tất cả nền có pixel accuracy 99% nhưng bỏ mọi vùng mực. Ta cần metric đo chồng lấn vùng mục tiêu.

Với $P$ là tập pixel dự đoán foreground, $G$ là tập thật:

Viết bằng TP, FP, FN: IoU $=\mathrm{TP}/(\mathrm{TP}+\mathrm{FP}+\mathrm{FN})$, Dice $=2\mathrm{TP}/(2\mathrm{TP}+\mathrm{FP}+\mathrm{FN})$. TN không xuất hiện, nên nhiều nền đúng không che hết lỗi foreground.

Ví dụ: ground truth [1,1,1,0], prediction [1,1,0,1] có TP = 2, FP = 1, FN = 1. IoU $2/4=0.5$, Dice $4/6=2/3$.

Với mask không rỗng, Dice $=2\operatorname{IoU}/(1+\operatorname{IoU})$, nên hai chỉ số đi cùng chiều nhưng giá trị khác. Ví dụ mask dự đoán 30 pixel, mask thật 20, chung 15 cho Dice $30/50 = 0.6$; còn IoU 0.6 ứng với Dice $1.2/1.6 = 0.75$. Ảnh vệ tinh 95% nền mà model đoán toàn nền thì pixel accuracy 95% nhưng IoU của lớp “nhà” bằng 0.

import numpy as np

def mask_scores(predicted, target):
    predicted = np.asarray(predicted, dtype=bool)
    target = np.asarray(target, dtype=bool)
    assert predicted.shape == target.shape
    intersection = np.logical_and(predicted, target).sum()
    union = np.logical_or(predicted, target).sum()
    total = predicted.sum() + target.sum()
    iou = 1.0 if union == 0 else intersection / union
    dice = 1.0 if total == 0 else 2 * intersection / total
    return float(iou), float(dice)

print(tuple(round(v, 4) for v in mask_scores([1, 1, 0, 1], [1, 1, 1, 0])))
print(mask_scores(np.zeros(4), np.zeros(4)))
print(mask_scores([0, 0, 0, 0], [1, 0, 0, 0]))

Kết quả(0.5, 0.6667)
(1.0, 1.0)
(0.0, 0.0)

Hàm trên cho điểm 1 khi cả hai mask đều rỗng: không có pixel foreground nào để đoán sai. Scorer của đề có thể bỏ ảnh hoặc lớp rỗng, hoặc dùng quy tắc khác; khi nhiều ảnh không có vật thể, chính quy tắc này quyết định một phần lớn điểm.

Cách gộp cũng đổi kết quả: mean Dice theo ảnh cho mỗi ảnh một phiếu, còn Dice tính sau khi gộp pixel của mọi ảnh cho vùng lớn trọng lượng lớn hơn.

U-Net giữ cả ngữ cảnh lẫn chi tiết

Encoder giảm độ phân giải để đọc ngữ cảnh rộng. Decoder tăng độ phân giải để xuất mask. Nếu chỉ phóng to biểu diễn nhỏ, thông tin biên đã mất không tự trở lại. U-Net nối feature map encoder với decoder ở độ phân giải tương ứng, thường bằng phép nối kênh (concatenate).

Khác residual $x+F(x)$ giữ nguyên số kênh, skip của U-Net cat hai tensor theo chiều kênh: số kênh cộng lại, còn $N,H,W$ của hai tensor phải khớp. Conv ngay sau phép nối phải được khai báo với tổng số kênh đó; mạng trong bài nối 4 kênh skip với 8 kênh từ đáy thành 12 kênh.

Decoder phóng to bằng một trong hai cách:

• Nội suy (F.interpolate, nearest hoặc bilinear): không có tham số học.
• Transposed convolution (nn.ConvTranspose2d): kernel học được. Với kernel $k$, stride $s$, padding $p$, chiều ra là $H_{\text{out}}=(H-1)s-2p+k$ (cộng thêm output_padding nếu đặt, mặc định 0), nên cấu hình hay gặp $k=s=2$, $p=0$ cho đúng $2H$. Ngược với conv thường, stride ở đây làm ảnh lớn lên: nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2) đưa (1, 64, 16, 16) thành (1, 32, 32, 32).

Với 4 lần pooling, cạnh ảnh nên chia hết cho 16; ảnh $101 \times 101$ bị làm tròn xuống ở mỗi lần pooling nên lúc nối hai tensor lệch kích thước, trừ khi pad ảnh lên 112 hoặc phóng theo đúng size của skip như mục sau.

Bài báo U-Net (Ronneberger và cộng sự, 2015) giảm độ phân giải bốn lần, số kênh tăng từ 64 lên 1024, và dùng conv không padding: ảnh vào $572\times572$ cho mask $388\times388$, nên feature map của encoder phải được cắt phần giữa (crop) trước khi nối. Mạng trong bài chỉ có một mức, conv có padding và phóng to bằng nội suy, đủ nhỏ để train trên CPU.

U-Net chỉ gồm conv, pooling và phép phóng to, không có tầng fully connected gắn với kích thước ảnh, nên cùng bộ trọng số chạy được trên ảnh lớn hơn lúc train.

Giới hạn thật là bộ nhớ: ảnh vệ tinh hay ảnh hiển vi vài nghìn pixel mỗi chiều có thể không vừa GPU. Khi đó cắt ảnh thành các tile chồng lấn, dự đoán từng tile, rồi ghép logits và chỉ giữ (hoặc lấy trung bình) phần giữa mỗi tile, vì pixel sát mép tile thiếu ngữ cảnh ở một phía.

Dữ liệu nhỏ để nhìn được sai đúng

Ảnh tự tạo có hình chữ nhật sáng, nhiễu nhỏ và nền tối; mask đúng là hình chữ nhật. Train và validation dùng seed khác. Độ sáng gần như cho thẳng nhãn, nên đây là bài dễ: nó dùng để kiểm cả đường đi từ dữ liệu, shape, loss đến metric, chưa phải để so sánh model.

import torch
from torch import nn
import torch.nn.functional as F

torch.set_num_threads(1)
torch.manual_seed(56)

def make_rectangles(count, seed):
    rng = np.random.default_rng(seed)
    masks = np.zeros((count, 1, 12, 12), dtype=np.float32)
    for index in range(count):
        row, col = rng.integers(1, 6, size=2)
        height, width = rng.integers(3, 6, size=2)
        masks[index, 0, row:row + height, col:col + width] = 1
    noise = rng.normal(0, 0.08, masks.shape).astype(np.float32)
    images = np.clip(0.1 + 0.8 * masks + noise, 0, 1)
    return torch.from_numpy(images), torch.from_numpy(masks)

train_x, train_y = make_rectangles(24, 10)
valid_x, valid_y = make_rectangles(8, 11)
print(tuple(train_x.shape), tuple(train_y.shape))
print(int(train_y[0].sum()), int(valid_y[0].sum()))

Kết quả(24, 1, 12, 12) (24, 1, 12, 12)
9 20

Hai tổng cuối là số pixel foreground trong hai mask mẫu: hình chữ nhật $3\times3$ và $4\times5$. Nhìn chúng phát hiện ngay nhãn vô tình toàn 0, hoặc nhãn mang giá trị lạ, trước khi train. Mask ở đây là float 0/1 vì BCEWithLogitsLoss cần target float; bài nhiều lớp dùng ID kiểu long như đã nói.

Viết và train U-Net nhỏ

Max pooling làm tròn xuống, nên với chiều lẻ, phóng gấp đôi không trở về đúng kích thước cũ: 13 thành 6, rồi 6 thành 12. Vì vậy interpolate nhận thẳng size của skip thay vì scale_factor=2. Ảnh $12\times12$ của bài chia hết cho 2, nhưng cùng code vẫn nối đúng với ảnh có chiều lẻ.

class TinyUNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.encode = nn.Sequential(nn.Conv2d(1, 4, 3, padding=1), nn.ReLU())
        self.middle = nn.Sequential(nn.Conv2d(4, 8, 3, padding=1), nn.ReLU())
        self.decode = nn.Sequential(nn.Conv2d(12, 4, 3, padding=1), nn.ReLU())
        self.output = nn.Conv2d(4, 1, 1)

    def forward(self, images):
        skip = self.encode(images)
        low = self.middle(F.max_pool2d(skip, 2))
        up = F.interpolate(low, size=skip.shape[-2:], mode="nearest")
        joined = torch.cat([skip, up], dim=1)
        return self.output(self.decode(joined))

model = TinyUNet()
optimizer = torch.optim.Adam(model.parameters(), lr=0.02)
criterion = nn.BCEWithLogitsLoss()
print("logits:", tuple(model(train_x[:2]).shape))
for step in range(80):
    model.train()
    optimizer.zero_grad()
    loss = criterion(model(train_x), train_y)
    loss.backward()
    optimizer.step()

model.eval()
with torch.no_grad():
    valid_logits = model(valid_x)
    valid_loss = criterion(valid_logits, valid_y).item()
    masks = valid_logits.sigmoid() >= 0.5
scores = [mask_scores(masks[i, 0].numpy(), valid_y[i, 0].numpy())
          for i in range(len(valid_y))]
print("valid BCE:", round(valid_loss, 4))
print("mean IoU, Dice:", np.mean(scores, axis=0).round(4).tolist())

Kết quảlogits: (2, 1, 12, 12)
valid BCE: 0.0034
mean IoU, Dice: [0.9922, 0.996]

Dòng đầu xác nhận logits có đúng shape (2,1,12,12) của mask. BCE validation là 0.0034; mean IoU 0.9922 và mean Dice 0.996 lấy trung bình theo 8 ảnh validation. Loss là thứ được tối ưu, IoU/Dice là thứ được chấm, và cả hai ở đây đều đo trên validation.

Điểm gần 1 phản ánh dữ liệu sinh rất sạch; vùng mực thật mờ, chồng lên dòng kẻ hoặc có ánh sáng không đều khó hơn nhiều.

Dice cũng dùng làm loss được nếu thay tập nhị phân bằng xác suất $p_i=\sigma(z_i)$:

với $g_i\in\{0,1\}$ là nhãn và $\varepsilon$ nhỏ (thường bằng 1) để mask rỗng không gây chia cho 0. Bản mềm này có đạo hàm theo $p_i$; threshold trước khi tính thì mất đạo hàm. Pixel nền đoán đúng ($p_i\approx0$, $g_i=0$) gần như không góp gì vào tử lẫn mẫu, nên Dice loss không bị phần nền lớn lấn át như BCE lấy trung bình theo pixel.

Nhiều đội cộng BCE với Dice loss; tổ hợp nào tốt hơn do validation quyết định, vì tối ưu Dice loss không bảo đảm Dice sau threshold cao nhất.

Cùng TinyUNet này có trong script của chương để train trên ảnh và mask của mình:

• Dữ liệu cần bốn thư mục train/images, train/masks, valid/images, valid/masks; ảnh và mask cùng tên, mask chỉ mang 0/255 hoặc 0/1.
• Lệnh không có --run chỉ kiểm dữ liệu. Thêm --run cùng --from-scratch, hoặc cùng checkpoint kèm SHA-256, thì script resize ảnh bằng bilinear và mask bằng nearest, train trên CPU với batch 2, ghi checkpoint có mean Dice validation tốt nhất rồi nạp lại để suy luận thử một batch.

Khi so các cấu hình, ghi kèm split, kích thước ảnh và số epoch.

RLE: dùng lại kỹ năng xử lý dãy

Run-length encoding (RLE) lưu các đoạn foreground liên tiếp sau khi trải mask. Bài dùng một quy ước cụ thể: trải theo hàng (row-major), vị trí bắt đầu tính từ 1, mỗi cặp là start và length. Đó không phải chuẩn chung: COCO trải theo cột và ghi xen kẽ độ dài các run 0 và run 1; nhiều cuộc thi dùng cặp start–length nhưng đánh số pixel từ trên xuống dưới rồi mới sang phải, tức trải theo cột.

Mask $\begin{bmatrix}0&1&0\\1&1&0\end{bmatrix}$ trải theo hàng là 0,1,0,1,1,0, nên RLE là [2,1,4,2]. Trải theo cột cho 0,1,1,1,0,0, ra [2,3]. Cùng mask nhưng bytes đầu ra khác.

def encode_rle(mask):
    values = np.asarray(mask, dtype=np.uint8).ravel(order="C")
    padded = np.r_[0, values, 0]
    changes = np.flatnonzero(padded[1:] != padded[:-1]) + 1
    changes[1::2] -= changes[::2]
    return changes.tolist()

def decode_rle(runs, shape):
    assert len(runs) % 2 == 0
    flat = np.zeros(np.prod(shape), dtype=np.uint8)
    for start, length in zip(runs[::2], runs[1::2]):
        assert 1 <= start and length > 0 and start - 1 + length <= len(flat)
        flat[start - 1:start - 1 + length] = 1
    return flat.reshape(shape, order="C")

small_mask = np.array([[0, 1, 0], [1, 1, 0]], dtype=np.uint8)
runs = encode_rle(small_mask)
print(runs)
print(np.array_equal(small_mask, decode_rle(runs, small_mask.shape)))
print(encode_rle(np.zeros((2, 3), dtype=np.uint8)))
print(encode_rle(np.ones((2, 3), dtype=np.uint8)))

Kết quả[2, 1, 4, 2]
True
[]
[1, 6]

Round trip True kiểm encoder/decoder nhất quán. Nó chưa chứng minh format khớp đề, vì hai hàm có thể cùng dùng sai thứ tự. Đối chiếu file mẫu của đề và vẽ mask giải mã ra mới kiểm được điều đó.

Đổi kích thước mask

Mask là nhãn rời rạc, nên khi resize phải dùng nearest: bilinear tạo giá trị trung gian giữa hai lớp. Khi đưa dự đoán về kích thước ảnh gốc, phóng logits hoặc xác suất (bilinear được, vì chúng là số liên tục) rồi mới threshold; RLE nộp bài phải mã hóa mask ở đúng shape gốc mà scorer dùng.

Lệch một pixel giữa ảnh và mask, do crop hoặc padding không đồng bộ, gây lỗi lớn ở nét mảnh. Connected components chỉ tách được những vùng không nối nhau, nên không thay được instance segmentation khi vật thể chạm nhau.