Thực hành U-Net: Phân đoạn ảnh hiển vi điện tử

Chuỗi U-Net — Bài 4. Ba bài trước tập trung vào kiến trúc, kích thước tensor và receptive field. Bài này chuyển sang một câu hỏi khó hơn: làm thế nào biến U-Net thành một hệ thống phân đoạn thực sự có thể huấn luyện, đánh giá và phân tích lỗi?

Bài toán cụ thể được chọn là phân đoạn cấu trúc thần kinh trên ảnh hiển vi điện tử truyền qua nối tiếp (serial section Transmission Electron Microscopy — ssTEM), gần với bài toán benchmark mà U-Net gốc được đánh giá. Đây là ví dụ phù hợp để thấy rằng chất lượng segmentation không chỉ phụ thuộc kiến trúc mạng mà còn phụ thuộc mạnh vào định nghĩa nhãn, loss, cách chia dữ liệu, augmentation và metric.

Tóm tắt nội dung

• Xây dựng pipeline U-Net cho dữ liệu ssTEM gồm ảnh grayscale và binary mask ở mức pixel.
• Phân tích tại sao binary cross-entropy và Dice loss tối ưu hai khía cạnh khác nhau, và vì sao kết hợp chúng thường là một baseline hợp lý.
• Augmentation hình học phải tác động đồng bộ lên ảnh và mask; interpolation của mask phải giữ nguyên nhãn rời rạc.
• Dice hoặc IoU cao chưa chắc bảo toàn topology: một lỗi nhỏ trên màng mỏng có thể làm hai cấu trúc thần kinh bị nối thành một vùng.

1. Bài toán: từ ảnh ssTEM đến bản đồ phân đoạn

ISBI 2012 tổ chức một challenge về phân đoạn cấu trúc thần kinh trong các stack ảnh hiển vi điện tử. Tập training được mô tả gồm 30 lát cắt ssTEM từ ventral nerve cord của ấu trùng Drosophila; nhãn được cung cấp dưới dạng binary segmentation. Đây cũng là benchmark được U-Net sử dụng trong công trình năm 2015. :contentReference[oaicite:0]{index=0}

Dữ liệu kiểu này khác đáng kể ảnh tự nhiên. Ảnh chỉ có một kênh cường độ, cấu trúc cần phân đoạn có hình dạng phức tạp và các màng ngăn cách hai vùng có thể rất mỏng. Một sai số chỉ vài pixel tại màng có khả năng làm hai cấu trúc vốn tách biệt trở thành một thành phần liên thông duy nhất.

Có thể biểu diễn bài toán binary segmentation bằng ánh xạ:

$$f:{\mathbb{R}}^{1\times H\times W}\to {\mathbb{R}}^{1\times H\times W}$$

Đầu vào x là ảnh grayscale, còn U-Net sinh một logit z_i cho mỗi pixel i. Xác suất foreground được tính bằng sigmoid:

$$p_i=\sigma \left(z_i\right)=\frac{1}{1+e^{-z_i}}$$

Trong PyTorch, nên để model trả về logits thay vì gọi sigmoid ngay trong forward. Các loss như BCEWithLogitsLoss kết hợp sigmoid và binary cross-entropy theo dạng số học ổn định hơn.

Hình 1. Kiến trúc U-Net gốc từ nhóm nghiên cứu University of Freiburg, tương ứng Figure 1 của công trình U-Net [1]. Contracting path thu nhận context, expanding path phục hồi localization và các kết nối ngang truyền feature có độ phân giải cao sang decoder. :contentReference[oaicite:1]{index=1}

2. Pipeline huấn luyện: phần khó không nằm riêng ở U-Net

Một pipeline segmentation hợp lý cần giữ rõ ranh giới giữa dữ liệu gốc, biến đổi dùng trong training, đầu ra xác suất và mask sau threshold. Trộn các bước này với nhau thường dẫn đến lỗi khó phát hiện, đặc biệt khi augmentation làm sai mask hoặc metric được tính trực tiếp trên logits.

Đặc biệt với serial sections, việc chia dữ liệu cũng cần thận trọng. Hai lát cắt liên tiếp có thể rất giống nhau về cấu trúc giải phẫu. Nếu lấy slice 10 làm training và slice 11 làm validation, kết quả validation có thể lạc quan hơn khả năng tổng quát hóa thực sự.

Nguyên tắc thực nghiệm: nếu dữ liệu gồm nhiều bệnh nhân, specimen hoặc volume, hãy chia train/validation/test ở cấp subject hoặc volume, không phải random từng patch. Với một stack nhỏ dùng cho bài thực hành, có thể giữ một block lát cắt liên tiếp làm validation, nhưng không nên coi kết quả đó là benchmark chính thức.

Challenge ISBI cung cấp training và test từ các volume riêng biệt, phản ánh đúng hơn cách đánh giá thuật toán trên dữ liệu chưa thấy. :contentReference[oaicite:2]{index=2}

3. Loss: pixel đúng chưa chắc cấu trúc đúng

Binary cross-entropy

Với ground truth y_i ∈ {0,1} và xác suất dự đoán p_i, binary cross-entropy trên N pixel là:

$${\mathcal{L}}_{\mathrm{BCE}}=-\frac{1}{N}\sum _{i=1}^N{y}_i\log \left(p_i\right)+(1-y_i)\log (1-p_i)$$

BCE xem mỗi pixel như một quyết định nhị phân. Điều này tạo gradient rõ ở mức pixel nhưng không trực tiếp tối ưu độ chồng lấp của toàn bộ vùng segmentation.

U-Net gốc thực tế không sử dụng Dice loss. Ronneberger và cộng sự sử dụng pixel-wise softmax kết hợp weighted cross-entropy, trong đó weight map được thiết kế để tăng mức phạt ở vùng biên giữa các đối tượng chạm nhau [1]. :contentReference[oaicite:3]{index=3}

Soft Dice loss

Dice coefficient cho binary segmentation có thể viết dưới dạng:

$$\mathrm{Dice}=\frac{2\mathrm{TP}}{2\mathrm{TP}+\mathrm{FP}+\mathrm{FN}}$$

Để dùng trong backpropagation, prediction nhị phân được thay bằng xác suất mềm p:

$$D_{\mathrm{soft}}=\frac{2\sum _i^N{p}_i{y}_i+\epsilon }{\sum _i^N{p}_i+\sum _i^N{y}_i+\epsilon }$$

Loss tương ứng là:

$${\mathcal{L}}_{\mathrm{Dice}}=1-D_{\mathrm{soft}}$$

Một objective dựa trên Dice được Milletari, Navab và Ahmadi sử dụng trong V-Net nhằm xử lý các tình huống có mất cân bằng mạnh giữa foreground và background [3]. :contentReference[oaicite:4]{index=4}

Tại sao kết hợp BCE và Dice?

Một baseline thực dụng là:

$$\mathcal{L}=\lambda {\mathcal{L}}_{\mathrm{BCE}}+(1-\lambda ){\mathcal{L}}_{\mathrm{Dice}}$$

BCE duy trì supervision tại từng pixel, còn Dice quan tâm trực tiếp hơn đến độ chồng lấp của toàn vùng. Giá trị λ không có một lựa chọn tối ưu cho mọi dataset; λ = 0.5 chỉ nên coi là điểm bắt đầu để thực nghiệm, không phải một hằng số chuẩn.

Một vấn đề sâu hơn: nếu màng thần kinh chỉ rộng vài pixel, một mô hình có thể đạt Dice vùng rất cao nhưng bỏ sót một đoạn màng ngắn. Về overlap, lỗi đó nhỏ; về topology, nó có thể làm hai neuron bị merge. Vì vậy loss theo vùng và metric theo vùng không mô tả đầy đủ mọi lỗi có ý nghĩa sinh học.

4. Chuẩn bị dữ liệu đúng: augmentation phải đồng bộ ảnh và mask

U-Net gốc nhấn mạnh data augmentation để khai thác hiệu quả số lượng ảnh gán nhãn hạn chế, bao gồm cả các biến dạng đàn hồi nhằm mô phỏng biến thiên hình học trong ảnh hiển vi [1]. :contentReference[oaicite:5]{index=5}

Trong một baseline hiện đại, có thể bắt đầu bằng flip và xoay 90°. Điều quan trọng là mọi biến đổi hình học phải áp dụng chính xác cùng tham số cho image và mask. Ngược lại, biến đổi cường độ như brightness, contrast hoặc noise chỉ tác động lên ảnh.

Không nội suy mask bằng bilinear. Mask chứa nhãn rời rạc. Khi resize hoặc rotate với góc bất kỳ, nên dùng nearest-neighbor cho mask; bilinear có thể tạo các giá trị 0.23, 0.61… tại biên và làm thay đổi ý nghĩa ground truth.

Giả sử dữ liệu được tổ chức thành hai thư mục images/ và masks/, đoạn mã sau tạo dataset cho ảnh grayscale:

dataset.py

from pathlib import Path
import random
import numpy as np
import torch
from PIL import Image
from torch.utils.data import Dataset

class EMSegmentationDataset(Dataset): def **init**(self, image_dir, mask_dir, augment=False):
self.images = sorted(Path(image_dir).glob("*.tif"))
self.masks = sorted(Path(mask_dir).glob("*.tif"))
self.augment = augment

```
    assert len(self.images) == len(self.masks)

def __len__(self):
    return len(self.images)

def __getitem__(self, idx):
    image = np.asarray(
        Image.open(self.images[idx]).convert("L"),
        dtype=np.float32
    ) / 255.0

    mask = np.asarray(
        Image.open(self.masks[idx]).convert("L"),
        dtype=np.float32
    ) / 255.0

    image = torch.from_numpy(image).unsqueeze(0)
    mask = torch.from_numpy(mask).unsqueeze(0)

    # Đưa mask về đúng {0, 1}
    mask = (mask > 0.5).float()

    if self.augment:
        if random.random() < 0.5:
            image = torch.flip(image, dims=[2])
            mask = torch.flip(mask, dims=[2])

        if random.random() < 0.5:
            image = torch.flip(image, dims=[1])
            mask = torch.flip(mask, dims=[1])

        k = random.randint(0, 3)
        image = torch.rot90(image, k, dims=[1, 2])
        mask = torch.rot90(mask, k, dims=[1, 2])

    # Chuẩn hóa riêng ảnh, không chuẩn hóa mask
    image = (image - image.mean()) / (image.std() + 1e-6)

    return image, mask

```

Việc dùng xoay theo bội số 90° ở đây có một lợi ích thực dụng: không cần interpolation nên cả ảnh và mask giữ được lưới pixel chính xác. Với elastic deformation như U-Net gốc, pipeline phải xử lý interpolation của image và mask riêng biệt.

5. Cài đặt BCE–Dice loss và metric đúng cách

Đoạn code sau tính loss trên xác suất mềm nhưng tính metric trên mask đã threshold. Hai bước này không nên nhầm lẫn.

loss_metrics.py

import torch
import torch.nn.functional as F

def soft_dice_loss(logits, targets, eps=1e-6):
probs = torch.sigmoid(logits)

```
dims = (1, 2, 3)
intersection = (probs * targets).sum(dims)
denominator = probs.sum(dims) + targets.sum(dims)

dice = (
    2.0 * intersection + eps
) / (denominator + eps)

return 1.0 - dice.mean()
```

def bce_dice_loss(logits, targets, alpha=0.5):
bce = F.binary_cross_entropy_with_logits(logits, targets)
dice = soft_dice_loss(logits, targets)

```
return alpha * bce + (1.0 - alpha) * dice
```

def segmentation_metrics(
logits, targets, threshold=0.5, eps=1e-6
):
probs = torch.sigmoid(logits)
preds = (probs >= threshold).float()

```
dims = (1, 2, 3)

tp = (preds * targets).sum(dims)
fp = (preds * (1 - targets)).sum(dims)
fn = ((1 - preds) * targets).sum(dims)

dice = (
    2 * tp + eps
) / (2 * tp + fp + fn + eps)

iou = (
    tp + eps
) / (tp + fp + fn + eps)

return {
    "dice": dice.mean().item(),
    "iou": iou.mean().item(),
}

```

Dice và Intersection over Union (IoU) đều đo overlap. Với cùng một binary mask:

$$\mathrm{IoU}=\frac{\mathrm{TP}}{\mathrm{TP}+\mathrm{FP}+\mathrm{FN}}$$

và hai metric liên hệ bởi:

$$\mathrm{Dice}=\frac{2\mathrm{IoU}}{1+\mathrm{IoU}}$$

Do đó Dice và IoU không phải hai góc nhìn hoàn toàn độc lập về cùng một binary segmentation. Nếu cần phân tích sâu hơn, nên bổ sung metric về boundary, precision/recall hoặc topology thay vì chỉ báo cáo thêm nhiều biến thể của overlap.

Threshold 0.5 không phải định luật

Sigmoid tạo xác suất trong khoảng 0–1, nhưng việc chọn 0.5 để chuyển sang binary mask chỉ là mặc định phổ biến. Nếu validation cho thấy threshold 0.42 tạo trade-off tốt hơn giữa false positive và false negative, có thể dùng 0.42.

Threshold phải được chọn trên validation set. Tìm threshold tối ưu trực tiếp trên test set làm rò rỉ thông tin test vào quá trình thiết kế mô hình và khiến kết quả cuối không còn là đánh giá độc lập.

6. Training loop và cách đọc kết quả

Giả sử class UNet đã được xây dựng như ở Bài 1, với một input channel và một output channel. Training loop cơ bản có thể viết như sau:

train.py

import torch

device = torch.device( "cuda" if torch.cuda.is_available() else "cpu"
)

model = UNet(
in_channels=1,
num_classes=1
).to(device)

optimizer = torch.optim.AdamW(
model.parameters(),
lr=1e-3,
weight_decay=1e-4
)

for epoch in range(30):
model.train()
train_loss = 0.0

```
for images, masks in train_loader:
    images = images.to(device)
    masks = masks.to(device)

    optimizer.zero_grad()

    logits = model(images)
    loss = bce_dice_loss(logits, masks)

    loss.backward()
    optimizer.step()

    train_loss += loss.item() * images.size(0)

model.eval()
val_dice = 0.0
val_iou = 0.0
n_batches = 0

with torch.no_grad():
    for images, masks in val_loader:
        images = images.to(device)
        masks = masks.to(device)

        logits = model(images)
        metrics = segmentation_metrics(logits, masks)

        val_dice += metrics["dice"]
        val_iou += metrics["iou"]
        n_batches += 1

val_dice /= n_batches
val_iou /= n_batches

print(
    f"Epoch {epoch + 1:02d} | "
    f"Dice={val_dice:.4f} | IoU={val_iou:.4f}"
)

```

Các giá trị learning rate, số epoch hoặc weight decay trong ví dụ trên chỉ là cấu hình khởi đầu. Không nên biến chúng thành “tham số chuẩn của U-Net”. Learning rate cần được kiểm tra qua validation; số epoch nên kết hợp early stopping hoặc lưu checkpoint tốt nhất thay vì mặc định model ở epoch cuối là model tốt nhất.

Đọc learning curve như một thí nghiệm

Hiện tượng	Khả năng xảy ra	Kiểm tra tiếp theo
Train loss giảm, validation Dice giảm	Overfitting	Augmentation, regularization, early stopping, kiểm tra split
Loss gần như không giảm	Learning rate, normalization hoặc bug dữ liệu	Overfit thử 1–2 ảnh trước
Dice cao nhưng nhiều vùng bị merge	Metric vùng bỏ qua lỗi topology nhỏ	Kiểm tra boundary/topology metric và visualization
Prediction toàn foreground hoặc background	Class imbalance, label sai hoặc loss sai	Kiểm tra phân bố mask, sigmoid và mapping nhãn

Một test rất hiệu quả trước khi huấn luyện toàn bộ dữ liệu là yêu cầu model overfit một hoặc hai ảnh. Nếu U-Net đủ năng lực nhưng không thể học gần khớp một mẫu nhỏ, thường có lỗi ở preprocessing, mask, loss, shape hoặc optimizer. Đây là phép kiểm tra pipeline, không phải phép đánh giá khả năng tổng quát hóa.

7. Dice cao vẫn có thể là segmentation kém

Đây là phần quan trọng nhất khi chuyển từ một bài demo sang một thí nghiệm có ý nghĩa. Giả sử ground truth chứa hai vùng thần kinh tách nhau bởi một màng chỉ rộng ba pixel. Mô hình dự đoán đúng gần như toàn bộ diện tích nhưng làm mất một đoạn màng dài 15 pixel.

Số pixel sai có thể rất nhỏ so với toàn ảnh, vì vậy Dice và IoU chỉ giảm nhẹ. Tuy nhiên sau connected-component analysis, hai vùng đã merge thành một đối tượng. Với bài toán phân tích cấu trúc thần kinh, lỗi đó có thể nghiêm trọng hơn hàng trăm pixel sai nằm sâu bên trong một vùng.

Chính vì lý do này, challenge ISBI 2012 không chỉ quan tâm pixel classification mà còn sử dụng các metric nhấn mạnh lỗi segmentation và topology, bao gồm warping error và foreground-restricted Rand measure. :contentReference[oaicite:6]{index=6}

Điều này dẫn tới một nguyên tắc quan trọng:

Metric phải phản ánh lỗi mà bài toán thực sự quan tâm.
Nếu mục tiêu chỉ là diện tích vùng, Dice có thể phù hợp. Nếu biên, connectivity hoặc số lượng đối tượng quan trọng, cần thêm boundary hoặc topology-aware evaluation.

Paper U-Net gốc cũng nhận ra vấn đề tách các đối tượng chạm nhau và thiết kế weight map cho cross-entropy để tăng trọng số tại vùng biên giữa hai đối tượng [1]. Đây không phải chi tiết phụ: nó cho thấy loss function nên được thiết kế dựa trên loại lỗi của bài toán chứ không chỉ dựa trên sự tiện lợi khi cài đặt. :contentReference[oaicite:7]{index=7}

Weighted loss của U-Net gốc

Weight map được paper định nghĩa theo khoảng cách đến biên của hai đối tượng gần nhất:

$$w\left(x\right)=w_c\left(x\right)+w_0\exp (-\frac{{\left(d_1\right(x)+d_2(x\left)\right)}^2}{2{\sigma }^2})$$

Trong đó d<sub>1</sub>(x) là khoảng cách từ pixel x đến biên của đối tượng gần nhất và d<sub>2</sub>(x) là khoảng cách đến biên của đối tượng gần thứ hai. Khi pixel nằm ở khe hẹp giữa hai đối tượng, cả hai khoảng cách đều nhỏ nên exponential term lớn và lỗi tại đó được phạt mạnh hơn.

Đây là ví dụ điển hình cho việc đưa prior về hình học của bài toán vào objective. Dice loss xử lý overlap; weighted cross-entropy có thể nhấn mạnh vùng biên; Tversky loss cho phép điều chỉnh bất đối xứng giữa false positive và false negative trong các bài toán mất cân bằng mạnh [4]. :contentReference[oaicite:8]{index=8}

8. Một protocol thực nghiệm đáng tin cậy hơn

Nếu mục tiêu là làm nghiên cứu thay vì chỉ chạy được model, nên cố định protocol trước khi tuning. Một cấu hình hợp lý cho bài toán này có thể gồm:

• Input: ảnh grayscale, normalize bằng thống kê train hoặc một quy tắc cố định.
• Split: tách theo volume/specimen khi có thể; không để các patch từ cùng ảnh xuất hiện đồng thời ở train và validation.
• Baseline: U-Net thuần với cùng preprocessing trước khi thử Attention U-Net, U-Net++ hoặc Transformer.
• Loss: BCE, Dice và BCE+Dice như ba experiment riêng thay vì chọn loss tốt nhất rồi chỉ báo cáo nó.
• Metric: Dice/IoU cho overlap, thêm metric phù hợp topology hoặc boundary nếu bài toán quan tâm connectivity.
• Threshold: chọn trên validation, sau đó khóa threshold trước test.
• Qualitative analysis: hiển thị input, ground truth, probability map và binary prediction cho cả ca tốt lẫn ca thất bại.

Một ablation nhỏ nhưng có giá trị hơn việc thay đổi hàng loạt thành phần cùng lúc:

Experiment	Loss	Augmentation	Mục đích
E1	BCE	Không	Baseline tối thiểu
E2	BCE + Dice	Không	Ảnh hưởng của loss
E3	BCE + Dice	Flip + rotate	Ảnh hưởng augmentation
E4	Boundary-weighted	Flip + rotate	Kiểm tra lỗi biên/topology

Không nên điền sẵn các con số “kỳ vọng” vào bảng trước khi chạy thí nghiệm. Giá trị của ablation nằm ở việc giữ các thành phần khác cố định để xác định thay đổi nào thực sự tạo ra khác biệt.

Kiểm tra cuối cùng trước khi tin kết quả: quan sát trực tiếp prediction. Nếu metric tăng nhưng mask xuất hiện lỗ, merge sai đối tượng hoặc mất các cấu trúc mỏng, metric đang không phản ánh đầy đủ objective của bài toán.

9. Kết luận

Trong một bài toán thực tế, U-Net chỉ là một thành phần của hệ thống segmentation. Quyết định nhãn nào là foreground, chia dữ liệu ra sao, augmentation có bảo toàn mask hay không, loss tối ưu loại lỗi gì và metric đo điều gì đều có thể ảnh hưởng đến kết luận nhiều không kém việc thay đổi kiến trúc.

Với ảnh hiển vi điện tử, BCE–Dice là baseline tốt để bắt đầu nhưng chưa chắc là objective cuối cùng. Khi các màng mỏng quyết định connectivity của cấu trúc, cần nhìn xa hơn overlap và phân tích trực tiếp boundary cũng như topology.

Đây cũng là bước chuyển quan trọng từ việc “dùng U-Net” sang thiết kế một thí nghiệm segmentation có thể bảo vệ về mặt khoa học. Bài tiếp theo có thể đi sâu vào Dice Loss, Focal Loss, Tversky Loss và Boundary Loss, phân tích gradient và chỉ rõ từng loss phù hợp với dạng mất cân bằng nào.

Tài liệu tham khảo

[1] O. Ronneberger, P. Fischer, and T. Brox, “U-Net: Convolutional Networks for Biomedical Image Segmentation,” in Medical Image Computing and Computer-Assisted Intervention – MICCAI 2015, Lecture Notes in Computer Science, vol. 9351, Munich, Germany, 2015, pp. 234–241, doi: 10.1007/978-3-319-24574-4_28.

[2] A. Cardona, S. Saalfeld, S. Preibisch, B. Schmid, A. Cheng, J. Pulokas, P. Tomancak, and V. Hartenstein, “An Integrated Micro- and Macroarchitectural Analysis of the Drosophila Brain by Computer-Assisted Serial Section Electron Microscopy,” PLoS Biology, vol. 8, no. 10, Art. no. e1000502, Oct. 2010, doi: 10.1371/journal.pbio.1000502. :contentReference[oaicite:9]{index=9}

[3] F. Milletari, N. Navab, and S.-A. Ahmadi, “V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation,” in Proc. 4th International Conference on 3D Vision (3DV), Stanford, CA, USA, 2016, pp. 565–571, doi: 10.1109/3DV.2016.79. :contentReference[oaicite:10]{index=10}

[4] S. S. M. Salehi, D. Erdogmus, and A. Gholipour, “Tversky Loss Function for Image Segmentation Using 3D Fully Convolutional Deep Networks,” in Machine Learning in Medical Imaging, Lecture Notes in Computer Science, vol. 10541, 2017, pp. 379–387.

[5] V. Jain et al., “Boundary Learning by Optimization with Topological Constraints,” in Proc. IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR), San Francisco, CA, USA, 2010, pp. 2488–2495, doi: 10.1109/CVPR.2010.5539950. :contentReference[oaicite:11]{index=11}