Image Segmentation: FCN, U-Net, DeepLab, Mask R-CNN

Object detection vẽ một hình chữ nhật quanh con chó. Nhưng hệ thống y tế cần biết chính xác đường viền khối u, xe tự lái cần biết pixel nào là mặt đường, máy kiểm tra cần đo diện tích vết xước. Khi đó ta cần Image Segmentation (phân vùng ảnh): gán nhãn cho từng pixel. Bài này tương ứng phần "Other topics" (module 06) của bài giảng gốc, và là phần đọc thêm cho chặng Detection / Segmentation / Anomaly Detection.

Số trang (ví dụ "trang 81") trong bài chỉ tới slide của bài giảng Deep Learning – Computer Vision – NLP (Thang Nguyen, Viet Nguyen); bài này giải thích chi tiết, kiểm chứng và đính chính các slide đó.

Ba kiểu phân vùng ảnh

Phân vùng là phân loại từng pixel

Với \(C\) lớp, mô hình semantic segmentation nhận ảnh \(3\times H\times W\) và trả về tensor logits \(C\times H\times W\): mỗi pixel một vector điểm số. Nhãn dự đoán là argmax theo trục kênh. Hàm loss cơ bản chính là cross-entropy, lấy trung bình trên mọi pixel:

\[ L = -\frac{1}{HW}\sum_{i=1}^{H}\sum_{j=1}^{W} \log p_{ij}\big(y_{ij}\big) \]

Trong PyTorch, nn.CrossEntropyLoss nhận trực tiếp logits (N, C, H, W) và nhãn (N, H, W) kiểu long.

Cách ngây thơ: áp dụng sliding window, cắt một mảnh nhỏ quanh mỗi pixel rồi phân loại. Cực kỳ lãng phí vì các mảnh chồng lấn gần như toàn bộ. Ta cần một mạng xử lý cả ảnh một lần và trả ra kết quả cùng độ phân giải.

FCN: Fully Convolutional Network (2015)

Long, Shelhamer và Darrell nhận thấy lớp fully connected là thứ duy nhất ép CNN nhận đầu vào cố định và mất thông tin vị trí. FCN:

1. Thay các lớp FC của mạng phân loại (VGG) bằng conv 1×1: mạng thành "toàn tích chập", nhận ảnh kích thước bất kỳ và trả về bản đồ điểm số thay vì một vector.
2. Bản đồ này nhỏ hơn ảnh 32 lần (do pooling), nên phải upsample (phóng to) trở lại kích thước gốc.
3. Kết hợp đặc trưng từ các lớp nông hơn (độ phân giải cao, ít ngữ nghĩa) với lớp sâu (ngữ nghĩa mạnh, độ phân giải thấp) qua skip connection để viền sắc nét hơn (FCN-32s → 16s → 8s).

Từ đây hình thành kiến trúc encoder–decoder phổ biến: encoder thu nhỏ để hiểu ngữ cảnh ("cái gì"), decoder phóng to để khôi phục chi tiết vị trí ("ở đâu").

Các cách upsample

Kích thước đầu ra của transposed convolution:

\[ O = (N - 1)\cdot S - 2P + K + \text{output\_padding} \]

Hai cấu hình gấp đôi kích thước hay dùng: \(K = 2, S = 2, P = 0\) và \(K = 4, S = 2, P = 1\) (16 → 32). Transposed conv đôi khi tạo vân "bàn cờ" (checkerboard artifacts) khi kernel không chia hết cho stride; một cách khắc phục là upsample bilinear rồi mới conv 3×3.

U-Net (2015): chuẩn mực của phân vùng

Ronneberger và cộng sự thiết kế U-Net cho ảnh y sinh, nơi dữ liệu gán nhãn rất ít. Hai đặc điểm:

• Đối xứng hình chữ U: mỗi tầng encoder (2 conv 3×3 rồi max pool) có một tầng decoder tương ứng (upsample rồi 2 conv 3×3).
• Skip connection bằng phép nối (concatenate): feature map của encoder được nối thẳng vào decoder cùng độ phân giải, cung cấp chi tiết biên cạnh mà quá trình pooling đã làm mất.

Đến nay, U-Net và các biến thể vẫn là lựa chọn mặc định cho ảnh y tế, ảnh vệ tinh, phân vùng lỗi sản phẩm, và cũng là "xương sống" của các mô hình sinh ảnh diffusion.

import torch
import torch.nn as nn


def double_conv(cin, cout):
    return nn.Sequential(
        nn.Conv2d(cin, cout, 3, padding=1, bias=False), nn.BatchNorm2d(cout), nn.ReLU(inplace=True),
        nn.Conv2d(cout, cout, 3, padding=1, bias=False), nn.BatchNorm2d(cout), nn.ReLU(inplace=True),
    )


class UNet(nn.Module):
    def __init__(self, in_ch=3, num_classes=2, base=32):
        super().__init__()
        c = [base, base * 2, base * 4, base * 8]
        self.enc1, self.enc2, self.enc3 = double_conv(in_ch, c[0]), double_conv(c[0], c[1]), double_conv(c[1], c[2])
        self.bottleneck = double_conv(c[2], c[3])
        self.pool = nn.MaxPool2d(2)
        self.up3, self.dec3 = nn.ConvTranspose2d(c[3], c[2], 2, stride=2), double_conv(c[3], c[2])
        self.up2, self.dec2 = nn.ConvTranspose2d(c[2], c[1], 2, stride=2), double_conv(c[2], c[1])
        self.up1, self.dec1 = nn.ConvTranspose2d(c[1], c[0], 2, stride=2), double_conv(c[1], c[0])
        self.head = nn.Conv2d(c[0], num_classes, 1)          # conv 1x1: mỗi pixel ra num_classes logits

    def forward(self, x):
        e1 = self.enc1(x)                     # (N, 32, H, W)
        e2 = self.enc2(self.pool(e1))         # (N, 64, H/2, W/2)
        e3 = self.enc3(self.pool(e2))         # (N, 128, H/4, W/4)
        b = self.bottleneck(self.pool(e3))    # (N, 256, H/8, W/8)
        d3 = self.dec3(torch.cat([self.up3(b), e3], dim=1))    # skip connection: nối kênh
        d2 = self.dec2(torch.cat([self.up2(d3), e2], dim=1))
        d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=1))
        return self.head(d1)                  # (N, num_classes, H, W)


model = UNet(num_classes=3)
logits = model(torch.rand(2, 3, 256, 256))
print(logits.shape)                           # torch.Size([2, 3, 256, 256])

masks = torch.randint(0, 3, (2, 256, 256))     # nhãn từng pixel
loss = nn.CrossEntropyLoss()(logits, masks)

Mô hình này có khoảng 1,93 triệu tham số; kích thước ảnh vào phải chia hết cho 8 (3 lần pooling).

DeepLab: giữ độ phân giải bằng atrous convolution

Thay vì thu nhỏ rồi phóng to, họ DeepLab (Google, 2015–2018) bớt pooling ở các tầng cuối và dùng dilated (atrous) convolution (xem Batch Normalization, Dropout và Receptive Field) để giữ receptive field lớn mà feature map không bị co lại quá nhiều (chỉ nhỏ hơn ảnh 8 hoặc 16 lần thay vì 32).

• ASPP (Atrous Spatial Pyramid Pooling): chạy song song nhiều conv atrous với dilation khác nhau (ví dụ 6, 12, 18) cộng global pooling, để nắm ngữ cảnh ở nhiều thang đo.
• DeepLabv3+ thêm một decoder nhẹ để làm sắc viền.

from torchvision.models.segmentation import deeplabv3_resnet50, DeepLabV3_ResNet50_Weights

weights = DeepLabV3_ResNet50_Weights.DEFAULT        # 21 lớp kiểu Pascal VOC (kể cả background)
model = deeplabv3_resnet50(weights=weights).eval()
batch = weights.transforms()(img).unsqueeze(0)       # img: tensor ảnh (C, H, W)
with torch.no_grad():
    out = model(batch)["out"]                         # (1, 21, H', W')
mask = out.argmax(dim=1)[0]                          # nhãn cho từng pixel
# Fine-tune cho dữ liệu riêng: model.classifier[4] = nn.Conv2d(256, num_classes, kernel_size=1)

Mask R-CNN: instance segmentation

He et al. (2017) mở rộng Faster R-CNN (xem Từ R-CNN đến Faster R-CNN) bằng cách thêm một nhánh thứ ba song song với nhánh phân loại và nhánh hồi quy box: một FCN nhỏ dự đoán mask nhị phân 28×28 cho mỗi RoI. Hai thay đổi quan trọng:

• RoI Align thay RoI Pooling: bỏ phép làm tròn, dùng nội suy song tuyến tính, giúp mask khớp đúng vị trí.
• Mask được dự đoán riêng cho từng lớp bằng sigmoid (không cạnh tranh giữa các lớp), lớp nào do nhánh phân loại quyết định.

from torchvision.models.detection import maskrcnn_resnet50_fpn_v2, MaskRCNN_ResNet50_FPN_V2_Weights

model = maskrcnn_resnet50_fpn_v2(weights=MaskRCNN_ResNet50_FPN_V2_Weights.DEFAULT).eval()
with torch.no_grad():
    pred = model([image])[0]          # image: float tensor (3, H, W), giá trị [0, 1]
print(pred.keys())                    # boxes, labels, scores, masks
binary_masks = pred["masks"][:, 0] > 0.5   # (số đối tượng, H, W)
# Khi huấn luyện, target cần thêm "masks": UInt8Tensor (N, H, W) bên cạnh "boxes" và "labels"

Đánh giá mô hình phân vùng

Ví dụ ảnh 4×4 với 3 lớp:

• Pixel accuracy = 13/16 = 0,8125. Dễ hiểu nhưng đánh lừa khi mất cân bằng: nếu 98% ảnh là nền, mô hình đoán toàn nền vẫn đạt 98%.
• IoU từng lớp = giao / hợp (như IoU của box, nhưng tính trên tập pixel). Lớp 0: 5/7 ≈ 0,714; lớp 1: 4/6 ≈ 0,667; lớp 2: 4/6 ≈ 0,667. mIoU (trung bình các lớp) ≈ 0,683: thước đo chuẩn của semantic segmentation.
• Dice coefficient (F1 trên pixel): \(\text{Dice} = \dfrac{2|A\cap B|}{|A| + |B|}\). Lớp 0: \(2\cdot5/(6+6) \approx 0{,}833\); trung bình ≈ 0,811. Phổ biến trong ảnh y tế. Liên hệ: \(\text{Dice} = \dfrac{2\,\text{IoU}}{1 + \text{IoU}}\).
• Instance segmentation dùng mask AP kiểu COCO (giống box AP nhưng IoU tính trên mask).

def miou_dice(pred, target, num_classes):
    ious, dices = [], []
    for c in range(num_classes):
        p, t = pred == c, target == c
        inter = (p & t).sum().item()
        union = (p | t).sum().item()
        if union == 0:              # lớp không xuất hiện ở cả hai: bỏ qua
            continue
        ious.append(inter / union)
        dices.append(2 * inter / (p.sum().item() + t.sum().item()))
    return sum(ious) / len(ious), sum(dices) / len(dices)


target = torch.tensor([[0, 0, 1, 1], [0, 0, 1, 1], [0, 2, 2, 1], [0, 2, 2, 2]])
pred   = torch.tensor([[0, 0, 1, 1], [0, 1, 1, 1], [0, 2, 2, 2], [0, 0, 2, 2]])
print(miou_dice(pred, target, 3))   # (0.6825, 0.8111)

Loss cho dữ liệu mất cân bằng

Khi vùng cần tìm rất nhỏ (vết nứt mảnh, khối u nhỏ), cross-entropy bị số đông pixel nền lấn át. Các lựa chọn:

• Cross-entropy có trọng số lớp: nn.CrossEntropyLoss(weight=...).
• Dice loss \(= 1 - \text{Dice}\) (tính trên xác suất mềm): tối ưu trực tiếp độ chồng lấn, không bị chi phối bởi pixel nền.
• Focal loss: giảm trọng số các pixel đã dễ đoán đúng.
• Thực tế hay dùng tổng CE + Dice.

Dữ liệu và công cụ

• Gán nhãn bằng đa giác (polygon) với CVAT, Label Studio, Labelme; lưu mask dạng ảnh PNG mỗi pixel là chỉ số lớp, hoặc định dạng COCO (polygon/RLE).
• Augmentation hình học phải áp dụng đồng thời và giống hệt cho ảnh và mask; khi resize mask dùng nội suy nearest (không dùng bilinear, sẽ sinh ra nhãn "lai" không tồn tại).
• Thư viện segmentation_models_pytorch cung cấp sẵn U-Net, FPN, DeepLabV3+ với hàng trăm backbone pretrained.
• Họ mô hình nền tảng Segment Anything (SAM, SAM 2… của Meta) phân vùng mọi vật thể theo gợi ý (điểm, box), rất hữu ích để gán nhãn nhanh rồi huấn luyện mô hình nhỏ chuyên biệt.

Tóm tắt

• Semantic: lớp cho từng pixel; instance: mask riêng cho từng đối tượng; panoptic: cả hai.
• FCN: toàn tích chập + upsample + skip; mở đầu kiến trúc encoder–decoder.
• U-Net: encoder–decoder đối xứng, skip connection nối kênh; tốt với ít dữ liệu.
• DeepLab: atrous conv + ASPP giữ độ phân giải và ngữ cảnh rộng.
• Mask R-CNN: Faster R-CNN + nhánh mask + RoI Align.
• Đánh giá bằng mIoU, Dice; cẩn thận với pixel accuracy khi dữ liệu mất cân bằng.

Thực hành: U-Net cho phân đoạn ảnh và Phân đoạn vết lỗi với YOLO26-seg. Đọc thêm: thay vì hiểu ảnh, hãy để máy tính tạo ra ảnh với GAN – mạng đối sinh.

Tài liệu tham khảo

• Long, Shelhamer, Darrell (2015). Fully Convolutional Networks for Semantic Segmentation.
• Ronneberger, Fischer, Brox (2015). U-Net: Convolutional Networks for Biomedical Image Segmentation.
• Chen et al. (2017–2018). Rethinking Atrous Convolution for Semantic Image Segmentation (DeepLabv3), DeepLabv3+.
• He et al. (2017). Mask R-CNN; Kirillov et al. (2019). Panoptic Segmentation; Kirillov et al. (2023). Segment Anything.
• CS231n Lecture: Detection and Segmentation; Dumoulin & Visin (2016). A guide to convolution arithmetic (transposed convolution).
• Tài liệu torchvision: Semantic segmentation và Instance segmentation models.