Segmentasi gambar Medis 3D dengan tutorial transformator

Transformer adalah tren besar dalam visi komputer. Saya baru-baru ini memberikan ikhtisar tentang beberapa kemajuan luar biasa. Kali ini saya akan menggunakan implementasi ulang model berbasis transformator untuk segmentasi 3D. Secara khusus, saya akan menggunakan trafo UNETR yang terkenal dan mencoba melihat apakah kinerjanya setara dengan UNET klasik. Buku catatannya sudah tersedia.

UNETR adalah arsitektur transformator pertama yang berhasil untuk segmentasi citra medis 3D. Pada postingan blog kali ini saya akan mencoba mencocokkan hasil model UNET pada dataset BRATS yang berisi gambar otak 3D MRI. Berikut adalah ikhtisar tingkat tinggi UNETR yang akan kita latih dalam tutorial ini:

Sumber: UNETR: Transformers untuk Segmentasi Citra Medis 3D, Hatamizadeh dkk.

Untuk menguji implementasi saya, saya menggunakan tutorial yang ada pada kumpulan data segmentasi MRI 3D. Oleh karena itu, saya harus memberikan penghargaan kepada perpustakaan sumber terbuka Nvidia yang luar biasa bernama MONAI karena menyediakan tutorial awal yang saya modifikasi untuk tujuan pendidikan. Jika Anda menyukai pencitraan medis, pastikan untuk membaca perpustakaan mengagumkan ini dan tutorialnya.

Yuk lihat datanya dulu!

Pembaruan: Rilis buku! Pelajari tentang “Pembelajaran mendalam dalam produksi” untuk menyajikan model ML Anda kepada jutaan pengguna.

Kumpulan data BRATS

BRATS adalah kumpulan data pencitraan 3D skala besar multi-modal. Ini berisi 4 volume gambar MRI 3D yang diambil dengan modalitas dan pengaturan berbeda. Berikut adalah contoh kumpulan datanya. Penting untuk melihat bahwa hanya tumornya saja yang dijelaskan. Hal ini membuat hal-hal seperti segmentasi menjadi lebih sulit karena model harus dilokalisasi pada tumor.

Gambar teaser data resmi dari situs penyelesaian BRATS

Tambalan gambar menggambarkan kategori tumor sebagai berikut (dari kiri ke kanan):

1. Busung: Seluruh tumor (kuning) biasanya terlihat pada gambar MRI T2-FLAIR.

2. Inti padat non-peningkat: Inti tumor (merah) terlihat pada T2 MRI.

3. Itu meningkatkan tumor struktur (biru muda). Biasanya terlihat di T1Gd, disekitarnya inti nekrotik (hijau).

4. Segmentasi digabungkan untuk menghasilkan label akhir kumpulan data.

Dengan MONAImemuat kumpulan data dari kompetisi dasalomba pencitraan medis menjadi hal yang sepele.

Pemuatan data dengan MONAI dan transformasi

Dengan memanfaatkan DecathlonDataset kelas perpustakaan MONAI seseorang dapat memuat salah satu dari 10 kumpulan data yang tersedia dari situs web. Kami akan menggunakan Task01_BrainTumour dalam kasus kami.

cache_num = 8
from monai.apps import DecathlonDataset
train_ds = DecathlonDataset(
   root_dir=root_dir,
   task="Task01_BrainTumour",
   transform=train_transform,
   section="training",
   download=True,
   num_workers=4,
   cache_num=cache_num,
)
train_loader = DataLoader(train_ds, batch_size=2, shuffle=True, num_workers=2)
val_ds = DecathlonDataset(
   root_dir=root_dir,
   task="Task01_BrainTumour",
   transform=val_transform,
   section="validation",
   download=False,
   num_workers=4,
   cache_num=cache_num,
)
val_loader = DataLoader(val_ds, batch_size=2, shuffle=False, num_workers=2)

Impor dan fungsi pendukung dapat ditemukan di notebook. Yang penting di sini adalah jalur transformasi, yang saya jamin bukanlah hal yang mudah dalam gambar 3D. MONAI menyediakan beberapa fungsi untuk membuat pipeline cepat untuk tujuan tutorial ini. Detail seperti orientasi gambar sengaja tidak disertakan dalam tutorial.

Secara singkat, kami akan mengambil sampel ulang gambar kami ke ukuran voxel 1,5, 1,5, dan 2,0 mm di setiap dimensi. Setelah itu, kita mengambil sub-volume 3D secara acak dengan ukuran 128, 128, 64. Hal ini tentu saja perlu diterapkan pada gambar masukan dan masker segmentasi.

Kemudian beberapa augmentasi diterapkan seperti membalik sumbu pertama secara acak, dan mengubah skala intensitas (jittering).

Kelas ConvertToMultiChannelBasedOnBratsClassesd membawa label ke format yang kita inginkan.

from monai.transforms import (
   Activations,
   AsChannelFirstd,
   AsDiscrete,
   CenterSpatialCropd,
   Compose,
   LoadImaged,
   MapTransform,
   NormalizeIntensityd,
   Orientationd,
   RandFlipd,
   RandScaleIntensityd,
   RandShiftIntensityd,
   RandSpatialCropd,
   Spacingd,
   ToTensord,
)
roi_size=(128, 128, 64)
pixdim=(1.5, 1.5, 2.0)
class ConvertToMultiChannelBasedOnBratsClassesd(MapTransform):
   """
   Convert labels to multi channels based on brats classes:
   label 1 is the peritumoral edema
   label 2 is the GD-enhancing tumor
   label 3 is the necrotic and non-enhancing tumor core
   The possible classes are TC (Tumor core), WT (Whole tumor)
   and ET (Enhancing tumor).
   """
   def __call__(self, data):
       d = dict(data)
       for key in self.keys:
           result = ()
           result.append(np.logical_or(d(key) == 2, d(key) == 3))
           result.append(
               np.logical_or(
                   np.logical_or(d(key) == 2, d(key) == 3), d(key) == 1
               )
           )
           result.append(d(key) == 2)
           d(key) = np.stack(result, axis=0).astype(np.float32)
       return d
train_transform = Compose(
   (
       LoadImaged(keys=("image", "label")),
       AsChannelFirstd(keys="image"),
       ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
       Spacingd(
           keys=("image", "label"),
           pixdim=pixdim,
           mode=("bilinear", "nearest"),
       ),
       Orientationd(keys=("image", "label"), axcodes="RAS"),
       RandSpatialCropd(
           keys=("image", "label"), roi_size=roi_size, random_size=False),
       RandFlipd(keys=("image", "label"), prob=0.5, spatial_axis=0),
       NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
       RandScaleIntensityd(keys="image", factors=0.1, prob=0.5),
       RandShiftIntensityd(keys="image", offsets=0.1, prob=0.5),
       ToTensord(keys=("image", "label")),
   )
)
val_transform = Compose(
   (
       LoadImaged(keys=("image", "label")),
       AsChannelFirstd(keys="image"),
       ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
       Spacingd(
           keys=("image", "label"),
           pixdim=pixdim,
           mode=("bilinear", "nearest"),
       ),
       Orientationd(keys=("image", "label"), axcodes="RAS"),
       CenterSpatialCropd(keys=("image", "label"), roi_size=roi_size),
       NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
       ToTensord(keys=("image", "label")),
   )
)

Selalu lebih baik untuk melihat alur beraksi, dengan memvisualisasikan beberapa bagian dari semua modalitas. Di bawah ini adalah contoh data kereta api kami:

Sumber: Gambar dari penulis berdasarkan buku catatan

Dapat diamati bahwa tumor tersebut adalah bukan saling eksklusif. Dalam hal ini kami memperkirakan peningkatan sel tumor dan nekrotik (peta segmentasi paling kanan) adalah yang paling sulit diprediksi.

Jalur data dan transformasi kini sudah siap. Mari kita lihat lebih dekat arsitektur modelnya.

Pelajari lebih lanjut tentang AI yang diterapkan dalam aplikasi pencitraan medis dari kursus AI untuk Kedokteran yang terstruktur dengan baik yang ditawarkan oleh Coursera.

Arsitektur UNETR

Berikut adalah arsitektur model yang menggabungkan transformator ke dalam arsitektur UNET yang terkenal:

Sumber: UNETR: Transformers untuk Segmentasi Citra Medis 3D, Hatamizadeh dkk.

Menariknya, saya mulai menerapkan model ini seperti pada gambar kertas di atas. Belakangan saya tahu bahwa hal itu sudah diterapkan di MONAI. Setelah memeriksa kode mereka, saya menemukan detail penting hilang. Kesimpulan: jangan percaya pada gambar arsitektur, gambar tersebut tidak memuat seluruh cerita tentang cara mengimplementasikan makalah tersebut. Untuk melihat kode implementasi, lihat implementasi saya di perpustakaan self-attention-cv.

Sekarang saya akhirnya dapat menggunakan implementasi UNETR saya. Saya telah membuat perpustakaan kecil yang mengimplementasikan beberapa blok perhatian mandiri untuk visi komputer dan mengemasnya dalam paket yang dapat diinstal pip. Jadi sekarang saya hanya perlu menginstal paket pip saya yang berisi model dan voila:

$ pip install self-attention-cv==1.2.3

Untuk menginisialisasi model kita perlu menyediakan ukuran volume, modalitas input pencitraan, jumlah label (output_dim) dan beberapa hal mengenai vision trafo. Contohnya termasuk menyematkan dimensi patch, ukuran patch, jumlah head, tipe normalisasi, dll.

from self_attention_cv import UNETR
device = torch.device("cuda:0")
num_heads = 10
embed_dim= 512
model = UNETR(img_shape=tuple(roi_size), input_dim=4, output_dim=3,
                embed_dim=embed_dim, patch_size=16, num_heads=num_heads,
                ext_layers=(3, 6, 9, 12), norm='instance',
                base_filters=16,
                dim_linear_block=2048).to(device)

Saya masih tidak yakin mengapa normalisasi instance bekerja sangat baik dengan UNET dan kumpulan data multi-model, tetapi ternyata berhasil! Intinya adalah kita memiliki 49,7 juta model parameter yang siap untuk dilatih.

Kami akan menggunakan DICE loss dikombinasikan dengan cross-entropy, dan buat loop pelatihan sederhana:

import torch.nn as nn
from monai.losses import DiceLoss, DiceCELoss
loss_function = DiceCELoss(to_onehot_y=False, sigmoid=True)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-5)
max_epochs = 180
val_interval = 5
best_metric = -1
best_metric_epoch = -1
epoch_loss_values = ()
for epoch in range(max_epochs):
   print(f"epoch Halo/para")
   model.train()
   epoch_loss = 0
   step = 0
   for batch_data in train_loader:
       step += 1
       inputs, labels = (
           batch_data("image").to(device),
           batch_data("label").to(device),
       )
       optimizer.zero_grad()
       outputs = model(inputs)
       loss = loss_function(outputs, labels)
       loss.backward()
       optimizer.step()
       epoch_loss += loss.item()
   epoch_loss /= step
   epoch_loss_values.append(epoch_loss)
   print(f"epoch pencinta average loss: slots")

Perbandingan dasar: UNET

Meskipun demikian, pertanyaan terbesarnya di sini adalah seberapa baik kinerja model ini. Oleh karena itu, kita memerlukan landasan yang kuat! Apa yang lebih baik dari UNET yang terkonfigurasi dengan baik yang digunakan dalam tutorial awal?

Saya juga membandingkan implementasi saya dengan implementasi UNETR MONAI. Mengapa? Karena tidak ada artinya jika saya menyamai kinerja baseline UNET dan masih kalah dengan implementasi resmi. Lagi pula, saya mengubah kode saya untuk mencerminkan perubahan arsitektur dari kode resmi. Dan memang saya melihat peningkatan kinerja yang sangat besar dibandingkan dengan implementasi sederhana dari gambar di makalah.

from monai.networks.nets import UNet
model = UNet(
   dimensions=3,
   in_channels=4,
   out_channels=3,
   channels=(16, 32, 64, 128, 256),
   strides=(2, 2, 2, 2),
   num_res_units=2,
).to(device)

Mari kita lihat nomornya terlebih dahulu:

Model	zaman	Berarti koefisien DICE.
MIMPI (dasar)	170	76,6%
UNETR (perhatian diri-cv)	180	76,9%
UNETR(MONAI)	180	76,1%

Untuk melacak pelatihan, kami mengukur kerugian pelatihan dari hilangnya dadu dan entropi silang. Kami juga melaporkan koefisien dadu untuk 3 label (saluran), yaitu Tumor Core (TC), Whole Tumor (WT), dan Enhancing Tumor (EC).

Di bawah ini Anda dapat melihat metrik ini saat berlatih:

Sumber: Gambar dari penulis berdasarkan buku catatan

Terakhir, kita dapat melihat hasilnya dengan membandingkan peta segmentasi keluaran dibandingkan dengan kebenaran dasar:

Sumber: Gambar dari penulis berdasarkan buku catatan

Saluran pada area nekrotik dihilangkan karena irisan khusus ini hampir tidak memiliki label ini. Ilustrasi ini hanyalah bagian tengah dari peta segmentasi 3D, jadi tentu saja bukan gambaran keseluruhan. Namun, hal ini memberi Anda gambaran tentang bagaimana model yang dilatih memberikan versi yang lebih sederhana dari label aslinya, yang dijelaskan oleh ahli radiologi ahli. Karena seperti biasa, jaringan saraf menyukai ruang pengoptimalan yang mulus.

Kesimpulan dan kekhawatiran

Saya belum yakin dengan kinerja transformator dalam pencitraan medis 3D. Saya yakin metode yang lebih maju dan kontribusi lainnya akan menyusul. Namun saya akui bahwa ini adalah karya menarik pertama yang menantang arsitektur UNET yang terkonfigurasi dengan baik, yang merupakan pilihan utama dalam tugas-tugas ini.

Dari analisis di atas, saya merasa penting untuk menyoroti juga bahwa aspek terpenting untuk mendapatkan kinerja yang baik, di sini Koefisien Dice, adalah prapemrosesan data dan jalur transformasi. Itulah sebabnya saya melihat terbatasnya inovasi dalam dunia pencitraan medis dalam hal pemodelan pembelajaran mesin, dan pekerjaan yang lebih menjanjikan dalam pengoptimalan pemrosesan data. Itu saja tidak menimbulkan masalah sama sekali, tapi itu membuat saya sangat curiga ketika sebuah makalah baru keluar dan mengklaim arsitektur baru. Karena perbandingannya sering kali tidak adil di domain khusus yang pernah saya kerjakan seperti pencitraan medis.

Seperti biasa, terima kasih atas minat Anda pada AI dan nantikan informasi lebih lanjut. Kami bangga berbagi dengan Anda buku kami tentang “Pembelajaran mendalam dalam produksi”, yang mengajarkan Anda cara menempatkan model Anda dalam produksi dan meningkatkannya. Dukungan komunitas (seperti berbagi di media sosial) selalu dihargai.

* Pengungkapan: Harap dicatat bahwa beberapa tautan di atas mungkin merupakan tautan afiliasi, dan tanpa biaya tambahan kepada Anda, kami akan mendapat komisi jika Anda memutuskan untuk melakukan pembelian setelah mengklik.