"Dạy AI làm bác sĩ" - Thực hành bài toán phân vùng ảnh y tế với mô hình Transformer
1. Giới thiệu bài toán
Ngày nay, Trí tuệ nhân tạo AI đang có rất nhiều ứng dụng trong các bài toán, trong đó có ứng dụng của AI trong xử lý ảnh y tế. Trên Kaggle có 1 challenge rất hay, mang tính ứng dụng cao đó là bài toán phân vùng khối u từ ảnh y tế (đường link: https://www.kaggle.com/c/bkai-igh-neopolyp/), tập dữ liệu của bài toán là BKAI-IGH NeoPolyp-Small do Trung tâm BKAI, Hanoi University of Science and Technology và Institute of Gastroenterology and Hepatology (IGH), Vietnam công bố. Dữ liệu bao gồm 1200 ảnh (1000 WLI images and 200 FICE images), tập training có 1000 ảnh, tập test có 200 ảnh và nhiệm vụ của chúng ta là dự đoán kết quả phân vùng cho 200 ảnh này và nộp kết quả lên Kaggle. Có 2 loại khối u cần phân vùng, 1 neoplastic tương ứng với màu đỏ - có thể tạm coi là khối u ác tính và loại 2 là non-neoplastic tương ứng màu xanh - có thể coi là khối u lành tính (All polyps are classified into neoplastic or non-neoplastic classes denoted by red and green colors, respectively).
Sau khi download data từ Kaggle về, ta thử xem dữ liệu
Như vậy nhiệm vụ của chúng ta rõ ràng rồi, đây là bài toán semantic segmentation cho ảnh y tế, đầu vào ảnh có khối u và nhiệm vụ của ta là tạo ra một mô hình AI phân vùng ra các khối u màu đỏ, màu xanh. Trình tự xử lý bài toán của chúng ta sẽ là:
- Đọc dữ liệu sử dụng Dataset của PyTorch kết hợp với các phép tăng cường dữ liệu từ thư viện Albumentations
- Một điểm dữ liệu đưa vào training của chúng ta sẽ là (X, y) = (ảnh đầu vào, ảnh kết quả phân vùng). Trong đó, ảnh đầu vào là ảnh y tế kích thước (H, W, 3); còn ảnh kết quả phân vùng sẽ là ma trận có kích thước (H, W) giống mô tả kết quả phân vùng, các giá trị trong ma trận kết quả kích thước (H, W) này nằm trong tập hợp {0, 1, 2} tương ứng với class background, class khối u màu xanh và class khối u màu đỏ.
- Khởi tạo mô hình và training
- Dự đoán kết quả và viết code sinh ra file nộp lên Kaggle theo chuẩn của BTC
- Lưu ý: Tập dữ liệu có 1000 train và 200 test nên mình sẽ train cả 1000 ảnh và lấy epoch cuối hoặc esemble của các epoch cuối nộp lên Kaggle mà không chia tập val
2. Thực hành
Đầu tiên sẽ là viết code để đọc data, phần này mình sử dụng Dataset của PyTorch cùng với thư viện Albumentations (code xử lý ảnh để đọc segmentation mask mình tham khảo từ https://github.com/GivralNguyen/BKAI-IGH-Neopolyp-Segmentation/blob/master/dataset.py). Lưu ý: khi sử dụng lại code của mình các bạn nhớ check lại path của ảnh cho chính xác với folder dữ liệu của các bạn.
trainsize = 384
class BKpolypDataset(torch.utils.data.Dataset):
def __init__(self, dir="path/to/data", transform=None):
self.img_path_lst = []
self.dir = dir
self.transform = transform
self.img_path_lst = glob.glob("{}/train/train/*".format(self.dir))
print(self.img_path_lst)
def __len__(self):
return len(self.img_path_lst)
def read_mask(self, mask_path):
image = cv2.imread(mask_path)
image = cv2.resize(image, (trainsize, trainsize))
image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
# lower boundary RED color range values; Hue (0 - 10)
lower1 = np.array([0, 100, 20])
upper1 = np.array([10, 255, 255])
# upper boundary RED color range values; Hue (160 - 180)
lower2 = np.array([160,100,20])
upper2 = np.array([179,255,255])
lower_mask = cv2.inRange(image, lower1, upper1)
upper_mask = cv2.inRange(image, lower2, upper2)
red_mask = lower_mask + upper_mask;
red_mask[red_mask != 0] = 2
# boundary RED color range values; Hue (36 - 70)
green_mask = cv2.inRange(image, (36, 25, 25), (70, 255,255))
green_mask[green_mask != 0] = 1
full_mask = cv2.bitwise_or(red_mask, green_mask)
full_mask = full_mask.astype(np.uint8)
return full_mask
def __getitem__(self, idx):
img_path = self.img_path_lst[idx]
image = cv2.imread(img_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = cv2.resize(image, (trainsize, trainsize))
label_path = img_path.replace("train", "train_gt")
label = self.read_mask(label_path)
if self.transform:
transformed = self.transform(image=image, mask=label)
image = transformed["image"]
label = transformed["mask"]
return image, label
Tiếp theo sẽ là các phép augmentation, mình sử dụng thư viện Albumentations để tạo ra các phép augmentation mà không cần phải lập trình xử lý ảnh. Các phép biến đổi mình sử dụng như sau:
train_transform = A.Compose([
A.HorizontalFlip(p=0.5),
A.VerticalFlip(p=0.5),
A.RandomGamma (gamma_limit=(70, 130), eps=None, always_apply=False, p=0.2),
A.RGBShift(p=0.3, r_shift_limit=10, g_shift_limit=10, b_shift_limit=10),
A.OneOf([A.Blur(), A.GaussianBlur(), A.GlassBlur(), A.MotionBlur(), A.GaussNoise(), A.Sharpen(), A.MedianBlur(), A.MultiplicativeNoise()]),
A.Cutout(p=0.2, max_h_size=35, max_w_size=35, fill_value=255),
A.RandomSnow(snow_point_lower=0.1, snow_point_upper=0.15, brightness_coeff=1.5, p=0.09),
A.RandomShadow(p=0.1),
A.ShiftScaleRotate(p=0.45, border_mode=cv2.BORDER_CONSTANT, shift_limit=0.15, scale_limit=0.15),
A.RandomCrop(trainsize, trainsize),
A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
ToTensorV2(),
])
val_transform = A.Compose([
A.Normalize(mean=(0.485, 0.456, 0.406),std=(0.229, 0.224, 0.225)),
ToTensorV2(),
])
Kết hợp các phép augmentation vào dataset vừa viết, ta được một điểm dữ liệu (X, y) khi đưa vào training như sau (VD: output tensor([0, 1]) có nghĩa là ảnh kết quả phân vùng chỉ chứa các điểm pixel thuộc class 0 - background và class 1 - khối u xanh)
Ok, như vậy chúng ta đã chuẩn bị xong dữ liệu rồi, bây giờ đến phần mô hình. Mô hình lựa chọn sẽ là mô hình SegFormer (các bạn xem cơ sở lý thuyết về SegFormer tại đây), ngoài SegFormer ra các bạn có thể thử các mô hình khác. Đầu tiên sẽ là cài đặt thư viện MMSegmentation và tải pretrained của SegFormer từ ImageNet.
Đối với Python Notebook, chạy các câu lệnh sau để cài đặt (nguồn: https://mmsegmentation.readthedocs.io/en/latest/get_started.html)
!pip install -U openmim
!mim install mmcv-full
!git clone https://github.com/open-mmlab/mmsegmentation.git
%cd mmsegmentation
!pip install -v -e .
# "-v" means verbose, or more output
# "-e" means installing a project in editable mode,
# thus any local modifications made to the code will take effect without reinstallation.
%cd ../
Sau đó sẽ là download pretrained trên ImageNet của backbone Mix Vision Transformer (nguồn: https://github.com/NVlabs/SegFormer)
!gdown 1EyaZVdbezIJsj8LviM7GaIBto46a1N-Z
!gdown 1L8NYh3LOSGf7xNm7TsZVXURbYYfeJVKh
!gdown 1m8fsG812o6KotF1NVo0YuiSfSn18TAOA
!gdown 1d3wU8KNjPL4EqMCIEO_rO-O3-REpG82T
!gdown 1BUtU42moYrOFbsMCE-LTTkUE-mrWnfG2
!gdown 1d7I50jVjtCddnhpf-lqj8-f13UyCzoW1
Tuy nhiên, các file pretrained này là dành cho SegFormer được code từ repo https://github.com/NVlabs/SegFormer (dựa trên MMSegmentation phiên bản cũ), chúng ta đang sử dụng MMSegmentation phiên bản mới và code SegFormer mới này sẽ khác về cách đặt tên các layer, do vậy cần phải chuyển lại key thì mới load được. MMSegmentation cung cấp tool để chuyển như sau (nguồn: https://github.com/open-mmlab/mmsegmentation/tree/master/configs/segformer)
!python mmsegmentation/tools/model_converters/mit2mmseg.py mit_b0.pth mit_b0_mmseg.pth
!python mmsegmentation/tools/model_converters/mit2mmseg.py mit_b1.pth mit_b1_mmseg.pth
!python mmsegmentation/tools/model_converters/mit2mmseg.py mit_b2.pth mit_b2_mmseg.pth
!python mmsegmentation/tools/model_converters/mit2mmseg.py mit_b3.pth mit_b3_mmseg.pth
!python mmsegmentation/tools/model_converters/mit2mmseg.py mit_b4.pth mit_b4_mmseg.pth
!python mmsegmentation/tools/model_converters/mit2mmseg.py mit_b5.pth mit_b5_mmseg.pth
Tiếp theo, mình sẽ khởi tạo mô hình SegFormer từ thư viện MMSegmentation, thư viện này khởi tạo mô hình từ config nên mình sẽ tạo config chuẩn của SegFormer B4 như sau (tham khảo cách tạo config: https://github.com/open-mmlab/mmsegmentation/blob/master/configs/base/models/segformer_mit-b0.py và https://github.com/open-mmlab/mmsegmentation/blob/master/configs/segformer/segformer_mit-b4_512x512_160k_ade20k.py)
model = dict(
type='EncoderDecoder',
pretrained=None,
backbone=dict(
type='MixVisionTransformer',
in_channels=3,
embed_dims=64,
num_stages=4,
num_layers=[3, 8, 27, 3],
num_heads=[1, 2, 5, 8],
patch_sizes=[7, 3, 3, 3],
sr_ratios=[8, 4, 2, 1],
out_indices=(0, 1, 2, 3),
mlp_ratio=4,
qkv_bias=True,
drop_rate=0.0,
attn_drop_rate=0.0,
drop_path_rate=0.1,
init_cfg=dict(type="Pretrained", checkpoint="mit_b4_mmseg.pth")),
decode_head=dict(
type='SegformerHead',
in_channels=[64, 128, 320, 512],
in_index=[0, 1, 2, 3],
channels=256,
dropout_ratio=0.1,
num_classes=3,
norm_cfg=dict(type='BN', requires_grad=True),
align_corners=False,
loss_decode=dict(
type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0)),
)
Có config rồi thì build mô hình từ config với thư viện MMSegmentation thôi:
import mmcv
from mmseg.models import build_segmentor
model = build_segmentor(model_cfg).to(device)
model.init_weights()
Thử forward 1 batch ảnh dữ liệu qua xem được chưa. Lưu ý: thư viện MMSegmentation hàm forward() thông thường của PyTorch đã bị ghi đè và yêu cầu truyền vào nhiều thuộc tính như meta data của ảnh, ... Còn hàm forward_dumy() của MMSegmentation thì hoạt động giống hàm forward thông thường của PyTorch.
Ok như vậy đã forward được thành công, mô hình trả về ảnh có cùng kích thước với ảnh đầu vào và số channel là 3 bằng với số class của bài toán.
Hàm loss mình sử dụng sẽ là focal loss, lý do là tập dữ liệu này đa phần là khối u màu đỏ, ít khối u màu xanh nên mình dùng focal loss để nhằm giảm thiểu tình trạng mất cân bằng dữ liệu này tới kết quả của mô hình. Giờ chúng ta sẽ training mô hình, mình cung cấp toàn bộ code cho toàn bộ quá trình này tại https://github.com/tungbt-k62/bkai-igh-neopolyp_SegFormer. Lưu ý: code này chạy được với nền tảng Google Colab, khi sử dụng với GPU riêng của các bạn thì chú ý check lại đường dẫn vào folder ảnh.
Quá trình training như sau:
Sau khi kết thúc training, theo đoạn code trên Github ta sẽ tiến hành đưa ra ảnh dự đoán cho toàn bộ tập test và xuất ra file csv (theo mẫu: https://github.com/sangdv/rle_encode/blob/main/mask2csv.py) để submit lên Kaggle. Chúng ta thử kiểm tra chất lượng "Bác sĩ AI" vừa được training qua một số ảnh test 
Chúng ta không có nhãn của tập dữ liệu test nên không thể đánh giá chính xác kết quả phân vùng của từng ảnh được, nhưng kiểm tra "bằng cơm" thì thấy "Bác sĩ AI" chúng ta vừa training có chất lượng tốt chứ nhỉ . Và khi ta submit file kết quả output.csv lên Kaggle, ta được kết quả khá tốt đó là 83.459 trên tập test dù train với size ảnh nhỏ hơn nhiều với size ảnh gốc.
3. Kết luận
Như vậy chúng ta đã cùng làm project nhỏ về bài toán có tính ứng dụng thực tiễn rất cao, đó là bài toán phân vùng khối u trong ảnh y tế. Mình cung cấp đoạn code baseline để các bạn có thể thử với nhiều mô hình khác trong thư viện MMSegmentation, cũng như các mô hình từ thư viện khác hoặc mô hình do chính các bạn tự lập trình.
Code các bạn lấy tại: https://github.com/tungbt-k62/bkai-igh-neopolyp_SegFormer
Nếu các bạn muốn tìm hiểu về PyTorch và bài toán Semantic Segmentation hoặc các bài toán về AI, hãy follow kênh Youtube của team mình nhé: https://www.youtube.com/@sunairesearch
Cảm ơn các bạn đã đọc bài & hẹn gặp lại!
All rights reserved
