Pytorch Image Segmentation Tutorial For Beginners — I
AI in Healthcare
Artificial intelligence (AI) is used in healthcare for prognosis, diagnosis, and treatment. AI reduces the healthcare cost and performance time, makes revenue cycle faster and more accurate, enhances the quality of healthcare processes.
AI to Enhance Radiology
Machine Learning and Artificial Intelligence are tremendous developments, especially with the new capabilities for image analysis that is making a real difference in speed, cost, and accuracy. It is going to power a lot of really exciting applications and solutions for radiologists to really help them do their job better, more easily, more reliably, and more usefully for their patients.
Image segmentation
In image segmentation the basic idea is we want to separate objects, we want to call different objects with different names depending on some properties of objects.
Making pixelwise binary classification of images is called “Semantic Segmentation”.
If we are trying to recognize many objects in an image we are performing “Instance Segmentation”. Instance Segmentation is a multiclass segmentation. For example, in self-driving cars, objects are classified as car, road, tree, house, sky, pedestrian, etc.
In this tutorial, we are doing semantic segmentation of Brain Tumor MRI images by making masks to them.
About The Project
We are making masks for brain tumor MRI images.
Brain tumor segmentation is an important task in healthcare AI. Early diagnosis of brain tumors has an important role in improving treatment possibilities and increases the survival rate of the patients. Manual segmentation of the brain tumors for cancer diagnosis, from large amount of MRI images generated in clinical routine, is a difficult and time consuming task. There is a need for automatic tumor image segmentation.
Pytorch
In this tutorial, I explained how to make an image segmentation mask in Pytorch. I gave all the steps to make it easier for beginners.
Models Genesis
In this project, I used Models Genesis. The difference of Models Genesis is to train a U-Net model using health data. I downloaded the U-Net model class and weights from the Github page of Models Genesis.
Dataset
We can download the dataset from figshare.
This brain tumor dataset containing 3064 T1 MRIs from 233 patients with three kinds of brain tumor: Meningioma (708 slices), Glioma (1426 slices), Pituitary tumor (930 slices).
This data is organized in matlab data format (.mat file). Each file stores a struct containing the following fields for an image:
cjdata.label: 1 for meningioma, 2 for glioma, 3 for pituitary tumor
cjdata.PID: patient ID
cjdata.image: image data
cjdata.tumorBorder: a vector storing the coordinates of discrete points on tumor border. We can use it to generate a binary image of the tumor mask.
cjdata.tumorMask: a binary image with 1s indicating tumor region
I used only the cjdata.label and the cjdata.image.
Load The Images
We need to load 3064 images into the notebook. Images are in .mat format and names of the files are numbers. (i.e. 1.mat, 2.mat…)
Initialize two lists to store the images and masks.
image_data = []mask_data = []
Define the filenames.
filenames = range(1,3065)Load all the images in a for loop
for name in filenames:file = h5py.File(‘/content/drive/My Drive/Brain-Tumor-Segmentation-Project/brainTumorData/’+str(name)+’.mat’, ‘r’).get(‘cjdata’)
Get the files
input = file.get(‘image’)[()]mask = file.get(‘tumorMask’)[()]
Change dimensions
Reshape, expand dimensions and append the images and masks to the lists.
input= np.reshape(input,(64,64,64))input = np.expand_dims(input,axis=0)image_data.append(input)mask = np.reshape(mask,(64,64,64))mask = np.expand_dims(mask,axis=0)mask_data.append(mask)
Note: There are 15 images whose sizes are not suitable for the model. We omit these images.
Split The Data
In deep learning, we need to make 3 splits: Train, test, and validation. So we use train_test_split two times. We can change the split ratio by changing the test size parameter. In this split ratio, we have train size = 2439, test size = 305, validation size= 305
from sklearn.model_selection import train_test_splitimage_train_data, image_test_data, mask_train_data, mask_test_data = train_test_split(image_data, mask_data, test_size = 0.1, random_state=123)image_train_data, image_val_data, mask_train_data, mask_val_data = train_test_split(image_train_data,mask_train_data, test_size = 0.111, random_state=123)
Data Loaders
Transform To Torch Tensor
We need to transform train, test, and validation data to tensors.
image_train_data = torch.Tensor(image_train_data)
mask_train_data = torch.Tensor(mask_train_data)Create Dataset Object
We create dataset objects for train, test, and validation sets.
train_dataset = TensorDataset(image_train_data, mask_train_data)Create Data Loader Object
We create data loader objects for train, test and validation dataset objects.
We should set shuffle = True for the loader to mix the dataset and give all classes to the training model.
train_dataloader = DataLoader(train_dataset, batch_size = 4, num_workers = 2, shuffle = True)Transfer Learning Model Structure
Imagenet is state of art in transfer learning. However, we use the structure of a U-net model that is trained in health data. We download the U-net class from Models Genesis Github page.
We initialize our model with the U-net class.
model = UNet3D()Pre-trained Weights
We download the weights of the model we use for transfer learning. You can download the weights from this request page.
Path
After getting the weights we define the path to the Genesis_Chest_CT.pt file.
weight_dir = ‘/content/drive/My Drive/Brain-Tumor-Segmentation-Project/Genesis_Chest_CT.pt’Weights
Load the weights of Model Genesis from Genesis_Chest_CT.pt file to the notebook.
checkpoint = torch.load(weight_dir,map_location=torch.device(‘cpu’))state_dict = checkpoint[‘state_dict’]
Initialize a dictionary in the notebook to store the weights.
unParalled_state_dict = {}Store weights in unparalleled_state_dict.
for key in state_dict.keys():unParalled_state_dict[key.replace(“module.”, “”)] = state_dict[key]
Load the new dictionary to the model.
model.load_state_dict(unParalled_state_dict)Parameters
We define the training parameters of the network before we start training.
Criterion
We define a loss function for the model. We can use Binary Cross-Entropy(BCE) loss but we use a combination of BCE and DICE losses. For detailed information about image segmentation metrics, read this post. You can reach the code for bce_dice_loss from that post.
criterion = bce_dice_lossOptimizer
We define an optimizer for the model. In this project, we use AdamW optimizer to prevent overfitting. AdamW implements Adam algorithm with weight decay fix.
from torch.optim import AdamWoptimizer = AdamW(model.parameters(), 0.1)
Learning Rate Scheduler
When we use a learning rate scheduler, we change the learning rate while training. We can achieve increased performance and faster training by adapting the learning rate for our stochastic gradient descent optimization procedure.
from torch.optim import lr_scheduler
scheduler = lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)Conclusion
Now we are ready to start training our model. In the next tutorial, we are going to train our model and interpret the results.
