Image Segmentation of Kidney Microscope Slides

The HuBMAP Kidney Segmentation challenge on Kaggle challenges an AI to segment regions of microvascular (blood vessels) from other tissue in a microscope slide of a healthy human kidney. Each pixel from a stained slide of the kidneys should be labeled as part of a vein or not. This page explores using a deep convolutional neural network to sole this problem. Otsu's segmentation is used as an example of a simpler, but ineffective approach, which justifies the machine learning approach.

Whole Slide Images (WSIs) from periodic acid-Schiff stained tissues slides were obtained from five healthy adults. These slides were split into tiles each of which is pixels, 8 bit RGB TIFF files. A CSV file contains the source whole stain image (WSI) and location of each tile within the WSI. Of the tiles, were labeled by expert pathologists. Three different structures were labelled, blood_vessel, glomerulus, and unsure. The labels are given per tile as a list polygons for each label. The model should predict a blood_vessel label for an unlabelled tile.

Semantic segmentation is the task of clustering parts of images together which belong to the same object class.

There exist several classical segmentation techniques, including thresh-holding, Otsu's method, active contour and mean shift. In recent years, deep convolutional neural networks (CNNs) have performed well in image segmentatio including segmentation of microscope slide.

There exist many CNNs designed for image segmentation, such as U-Net, DeepLab, and ErfNet. In some cases, an already trained network may be used and adapted by retraining through transfer learning. Transfer learning has been used successfully in biological image analysis.

In order to apply and evaluate our algorithms we need to convert the label data from polygons into mask files to match the training images. The polygons of each blood_vessel label are converted to a mask image. Any tiles which had no labels, including glomerulus and unsure, are excluded from training. If a tile has a label, but not a blood_vessel label, a blank mask is generated.

We need useful metrics which can be used to evaluate our algorithm's success. Müller et al. suggested a guideline for evaluation metrics for medical image segmentation tasks. These metrics are formally defined below, feel free to skip this section if you are already familiar with DSC, IoU, and Sensitivity and Specificity.

Let be matrices representing the predicted and ground-truth labels respectively. A in the matrix means the corresponding pixel is part of a vein. The true-positives, false-positives, false-negatives and true-negative counts are given by

where is the element-wise product.

Müller et al. suggests the dice similarity coefficient (DSC) and intersection over union (IoU or Jaccard index) because they are unbiased metrics (pixel accuracy is an example of a biased metric). Biased metrics can cause models which classify unbalanced classes to look more accurate than if balanced classes were classified. DSC and IoU can be defined as

Additionally, sensitivity and specificity are used to demonstrate functionality but not performance, defined as

Before jumping to more complex deep learning techniques, a simpler model should be considered. Otsu's segmentation is one of the simplest forms of image segmentation. The algorithm finds the optimal binary threshold to separate an image into two classes which have least variance. The algorithm minimizes

where and is the probability of a pixel belonging to the classes and , which are seperated by the threshold , and and are the variances of these respective classes.

Otsu's segmentation is a reasonable baseline for testing if machine learning is necessary to solve this problem. We can see by looking at the training data that a thresholding algorithm won't work as contextual information is required, we can see in table 1 that Otsu's is not much better than random noise. We need to use machine learning and table 1 will be a useful baseline for our neural network to beat.

U-Net is a fully convolutional neural network which was introduced by Ronneberger et al. where it was successful in segmenting neuronal tissue in electron microscope data. U-Net is widely used in medical image segmentation. U-Net's architecture is shown in figure 2, it contains no fully connected layers so can accept any resolution image as input.

For the following experiments a standard 90-10 random train-test split was used. U-Net's logits were passed through a softmax to normalize results into a probability distribution.

The model is initially trained without normalizing the tiles. Figure 3 compares DSC during training for normalized and raw tiles. The tiles are normalized with Albumentations with a mean and standard deviation pre-calculated from the entire training set (this technically is data leaking between training and validation sets).

Applying normalization allowed the epoch loss to decrease more quickly, this could be due to overfitting, but the validation shows an greater peak DSC over no normalization. Applying random morphological and color transforms didn't seem to have any benefit and slowed training.

The U-Net paper suggests a batch size of to maximise GPU usage and a correspondingly high momentum. Kandel et al. investigated different batch sizes when applied to CNN image classifiers and states that low batch sizes can lead the network to "bounce back and forth without achieving acceptable performance." The DSC in figure 3 shows this characteristic ping-ponging suggesting the batch size is too low. On the other hand, higher batch sizes can mean training takes too long to converge, the optimum batch size depends on many factors including the data-set, so experimentation is required.

Figure 4 shows that as batch size increases, IoU and DSC take longer to improve but show less variance, as predicted by Kandel et al. Higher batch sizes seem to reach a greater peak DSC, but IoU shows a greater peak for low batch sizes, this may be because more epochs are required to converge. A batch size of was chosen as optimal.

The loss function measures how far a predicted label is from the ground truth. Binary Cross Entropy (BCEwithLogitLoss) is a widely employed loss function in semantic segmentation because of its stability. There exist other common semantic segmentation loss functions, which include focal loss, and combo loss. With the help of bigironsphere's Kaggle notebook the loss functions which are not built-in can be implemented in PyTorch. See Jadon et al. for the mathematical definition of these loss functions, my focal implementation sets to the inverse class frequency and , my combo loss weighs BCE and dice equally where . Before applying focal and combo loss, U-Net's logit output is passed through a sigmoid.

Figure 5 compares these metrics over equivalent training setups. Combo loss consistently out performed BCE and focal loss, showing less variability and an greater peak metric. Focal loss performed very poorly, this may due improper tuning of the and hyper-parameters, as when , focal loss should be equivalent to BCE.

After tuning these hyper-parameters a final model was trained over epochs with a batch size of using a combo loss function. The DSC shown in figure 6 shows the model converged towards a DSC of . Results and an example of the model's predictions are shown in table 3.

This article examined different techniques and parameters when applying semantic segmentation techniques to the HuBMAP - Hacking the Human Vasculature competition. We found that normalization of input images allowed the loss to decrease more quickly. Applying random morphological or color transformation did not seem effective on this data-set. For the standard U-Net architecture, Combo loss outperformed the simpler binary cross entropy loss and focal loss. Greater batch sizes decreased loss variance at the cost of training time and should be tuned for each data-set.