Hiện nay bài toán phân đoạn ảnh (segmentation) đang phổ biến trong nhánh Xử Lý Ảnh đặc biệt là trong các bài toán về y khoa. Khi xét về mặt đạo đức và hiệu suất mang lại.

Có thể kể đến như LinkNet, UNet, PSPNet, Mask R-CNN, DeepLab trong đó bao gồm cả DeepLab v3 và DeepLab v3+,… các mạng mình vừa nêu chỉ là một trong số các mạng có khả năng phân đoạn tốt nhất hiện nay. {Mình sẽ viết hết các mạng trên nếu đủ khả năng :v}

Việc phân đoạn ảnh trong y khoa có rất nhiều ứng dụng. Ví dụ, trong chẩn đoán và điều trị ung thư, phân đoạn ảnh giúp xác định vị trí và biên giới của khối u, tế bào ung thư,… trong cơ thể để hỗ trợ việc xác định phạm vi bệnh và lập kế hoạch điều trị. Trong điều tra bệnh lý, phân đoạn ảnh có thể giúp xác định cấu trúc bất thường, ví dụ như khối u, vết thương hoặc mô tổn, để hỗ trợ việc chẩn đoán và điều trị.

Unet là một trong những mô hình mạng neural được sử dụng rộng rãi trong lĩnh vực y học và đặc biệt là trong phân đoạn ảnh y tế.

 

Ảnh 1: Minh họa phân đoạn ảnh để xác định vị trí vật thể

Ảnh 2: Màu trắng là vùng của tế bào ung thư

1. Mạng Học Sâu Unet

Unet là một kiến trúc mạng học sâu được giới thiệu vào năm 2015 bởi Olaf Ronneberger cùng các cộng sự của mình trong bài báo “U-Net: Convolutional Networks for Biomedical Image Segmentation” [1], sở dĩ được gọi là Unet là do kiến trúc này có hình chữ U. Bao gồm 2 phần chính là Encoder và Decoder.

Ảnh 3: Kiến trúc nguyên bản mạng Unet

2. Encoder trong kiến trúc Unet

Bằng cách sử dụng các lớp tích chập và max pooling, bộ mã hóa tạo ra một không gian đặc trưng với thông tin tập trung vào các đặc điểm quan trọng trong ảnh. Các đặc trưng này được mã hóa thành các ma trận có kích thước nhỏ hơn và chứa thông tin trừu tượng hơn về các cấu trúc và đối tượng trong ảnh.

Encoder trong Unet có nhiệm vụ là chuyển đổi ảnh đầu vào với kích thước input_width x input_height thành các đặc trưng (features) để biểu diễn thông tin của ảnh đó, Encoder sẽ sử dụng các lớp Convolutional để trích xuất đặc trưng kèm theo đó sẽ giảm kích thước đầu vào của ảnh theo từng lớp Downsampling.

Ở cuối mỗi lớp Downsampling thì sẽ là một lớp Maxpooling 2×2 để giảm kích thước ảnh.

Sau khi qua quá trình Encoder thì chúng ta thu được các đặc trưng với kích thước (new_width x new_height) < (input_width x input_height). Đồng thời chiều sâu của ảnh đầu vào cũng được tăng lên, các đặc trưng này sẽ được truyền đến Decoder để giúp mô hình phục hồi lại ảnh đầu ra từ ảnh đầu vào.

3. Decoder trong kiến trúc Unet

Các đặc trưng từ bộ mã hóa được chuyển đến bộ giải mã thông qua các kênh kết nối (skip connections). Kênh kết nối cho phép truyền thông tin chi tiết từ bộ mã hóa sang bộ giải mã, giúp tái tạo các chi tiết quan trọng trong quá trình phân đoạn. Cụ thể, kênh kết nối kết hợp các đặc trưng từ các tầng mã hóa và giải mã cùng với các đặc trưng tương ứng của các tầng giải mã để tạo ra bản đồ phân đoạn chi tiết.

Decoder thì ngược lại với Encoder, đoạn này sẽ dùng các đặc trưng từ phần Encoder sau đó sẽ trải qua các lớp Upsampling để học cách ghép các đặc trưng này lại với nhau và tái tạo lại ảnh đầu ra bằng với kích thước của ảnh đầu vào.

Decoder sử dụng các lớp Up-convolution 2×2 để tăng kích thước ảnh lên gấp đôi và sử dụng kết nối Skip-connections để tránh việc mất mát thông tin trong quá trình Encoder và Decoder

Kết quả là ta sẽ nhận được ảnh mask dự đoán trên pixel ảnh với kích thước (output_width x output_height) = (input_width x input_height)

*Các bạn khi setup mạng nhớ chú ý đến đường concatenate vì nó là điểm quan trọng giúp mô hình có khả năng tái tạo lại các chi tiết trong quá trình phân đoạn ảnh. {Mình đã có thử setup không concatenate, mặc dù model-size có giảm nhưng kết quả cho ra khá tệ}

Ảnh 4: Lưu ý đường concatenate khi setup mạng

4. Ứng dụng Unet để phân đoạn ảnh siêu âm xác định tế bào ung thư vú

Hiện nay UNet cũng đã được đóng gói thành thư viện để có thể dễ sử dụng hơn. (tham khảo: https://github.com/qubvel/segmentation_models)

Trong bài viết này mình sẽ setup lại mạng nhưng sẽ cắt giảm một số lớp Convolutional. Bộ dữ liệu được sử dụng là hình ảnh chụp siêu âm ngực của hơn 600 bệnh nhân nữ với độ tuổi từ 25 đến 75 tuổi có chứa tế bào ung thu vú ở mức độ nhẹ, được chia sẻ từ nghiên cứu của Trường Đại học Đại học Cairo – Ai Cập.[2]

Trên bộ dữ liệu siêu âm vú được phân loại thành ba loại: hình ảnh bình thường, lành tính và ác tính, nhưng mình sẽ chỉ lấy ra ngẫu nhiên trong nhãn đó 437 ảnh siêu âm tương ứng với 437 ảnh mask, sẽ được chia thành 351 để train model và 86 ảnh test.

Ảnh 5: Cấu trúc lưu trữ dataset

Trực quan ảnh ra xem các ảnh đã khớp với nhãn hay chưa

Ảnh 6: Trực quan dữ liệu

Cấu hình máy:

• Ngôn ngữ: Python 3.10.11
• TensorFlow: 12.0
• Sản phẩm được viết và thực thi tại Google Colaboratory.
• Lưu trữ dưới dạng file Jupyter Notebook (ipynb)
• Chạy trên thiết bị: Asus Tuf A15, CPU AMD Ryzen 5 4600H, Ram 8GB, Card rời GTX 1650Ti
• Bộ dữ liệu được lưu trữ tại Drive.

Setup mạng: mình có điều chỉnh lại ở chiều sâu của các lớp, ảnh đầu vào chỉ cho bộ lọc là 32 thay vì là 64 như của tác giả, tương tự các lớp tích chập tiếp theo số bộ lọc sẽ được lũy thừa lên và đến lớp cuối cùng sẽ là 512 (xem ảnh 7, 8). Kích thước ảnh đầu vào sẽ là 128x128x1.

Ảnh 7: Encoder

Ảnh 8: Decoder

các tham số huấn luyện như sau:

Ảnh 9: Setup tham số huấn luyện

5. Kết quả huấn luyện Unet

Với số Epoch setup là 300 và EarlyStopping là 80 epoch thì model đã dừng sớm ở epoch thứ 120 với tổng thời gian train là: 2.18 phút khi train trên GPU do Colab cung cấp
Kết quả cho ra hàm loss của tập train có xu hướng giảm dần đều, nhưng làm loss tập validation lại có xu hướng tăng dần, nên model đã dừng sớm để tiết kiệm chi phí tính toán.

Ảnh 10: Hàm loss của tập train và validation

Dự đoán ảnh mask trên tập test thấy rằng model có phần nào đó dự đoán được đặc trưng, nhưng xung quanh phần dự đoán các giá trị vẫn còn tương đối bị nhiễu nên chưa cho ra được hình ảnh “mịn” nhất.

Ảnh 11: Kết quả dự đoán trên tập test

Có thể thấy rằng Model Unet có thể dự đoán tốt được ảnh mask khi có thể xác định được vùng của tế bào ung thư xuất hiện, nhưng điểm hạn chế là kích thước của “vùng” này chưa thật sự khớp, khi vẫn còn tồn tại các giá trị nhiễu và các nét của vùng chưa được “mịn”
Điều này dẫn đến việc khi ta thêm một ảnh mới vào để dự đoán, khiến hiệu suất dự đoán trên dữ liệu mới sẽ bị suy giảm.
Khả năng dự đoán mask của bộ dữ liệu cũng tương đối khả quan, nhưng chỉ có thể áp dụng thực tế đối với hình ảnh siêu âm mà các tế bào ung thư không bị chồng lên nhau.

Tài liệu tham khảo:

[1] O. Ronneberger, P. Fischer, and T. Brox, ‘U-Net: Convolutional Networks for Biomedical Image Segmentation’, arXiv [cs.CV]. 2015

[2] W. Al-Dhabyani, M. Gomaa, H. Khaled, and A. Fahmy, ‘Dataset of breast ultrasound images’, Data in Brief, vol. 28, p. 104863, 2020.