Retinal Cyst Segmentation Using Semi Supervised CNN’s

Mentors:                                                             

Hayyan Arshad                                                                                            

Vartika T Rao

Mohammad Aadil Shabbier

Mentees:

Amruth S D

Adithya Ubaradka

Abhishek Srinivas

Thrishank Reddy

Table of Contents:

1.  Title
2.  Aim
3.  Introduction
4.  Methodology
5.  Results
6.  Inferences
7.  References

Aim:

The aim of this project is to develop and implement a semi-supervised convolutional neural network (CNN) model for accurate and efficient segmentation of cysts in medical imaging data. By leveraging both labeled and unlabeled data, the goal is to enhance the performance of traditional CNN-based segmentation methods and improve the delineation of cyst boundaries in medical images.

Introduction:

Cyst segmentation plays a crucial role in medical image analysis, particularly in the diagnosis and treatment planning of various diseases. Accurate delineation of cyst boundaries is essential for effective clinical decision-making. Convolutional Neural Networks (CNNs) have demonstrated remarkable success in automated medical image segmentation tasks. However, the performance of CNN-based segmentation models heavily relies on the availability of annotated data, which can be scarce and costly to obtain.

To address this challenge, this project proposes the utilization of semi-supervised learning techniques, specifically FixMatch or MixMatch, in the development of a robust cyst segmentation model. Semi-supervised learning leverages both labeled and unlabeled data to improve model performance, making it particularly suitable for scenarios where labeled data is limited but unlabeled data is abundant.

Methodology and Architecture used:

UNET:

Initially, we started by working on the Retouch dataset using the UNET architecture (This was a segmentation task using supervised learning). The dataset (Retouch), had 4 classes of cysts and every pixel was classified as belonging to one of those 4 classes. The UNET architecture comprises an encoder and the decoder that have skip connections in between allowing the UNET to capture both global context as well as local features. The UNET architecture is shown below:

Mean Teacher:

Next, we started working on the classification of the CIFAR 10 dataset using Mean Teacher Architecture(a semi supervised learning algorithm). 

The CIFAR - 10 dataset comprises of 50000 images of 32 * 32 resolutions belonging to one of 10 classes. The Mean Teacher algorithm uses 2 separate models (albeit having the same architecture) - one being the student model and the other being the teacher model. The teacher model essentially acts as a guide for the student model. In this Mean Teacher model, the teacher uses an Exponential Moving Average of the student’s weights, allowing a quick feedback loop between the teacher and the student models. The loss function that is used in the Mean Teacher model is a combination of two : - (The labeled image loss consisting of the Cross Entropy Loss of the labeled images added with the Mean Square Error Loss and the unlabeled image loss consisting of the Mean Square Loss between the teacher and the student predictions). The following is the architecture used by the Mean Teacher model:

FixMatch:

Next, we worked on the FixMatch architecture (again, a supervised learning algorithm). The FixMatch learning algorithm combines 2 techniques: consistency regularization and pseudo labeling. The FixMatch algorithm came about because of the domain specific data augmentation involved in other SSL models. In the unlabeled sample of the FixMatch, both weak and strong augmentation is done on the input image. The high confidence predictions of the weakly augmented unlabeled samples are taken as the pseudo labels. These are used as the labels for the unlabeled samples and an unsupervised loss (Cross Entropy) is calculated. There is also an additional Cross Entropy Loss that arises from the labeled images that are weakly augmented. The architecture depicting the FixMatch algorithm is as follows:

MixMatch:

Next, we also used the Mix Match semi supervised learning algorithm. In this algorithm, In MixMatch, the algorithm starts by applying data augmentation to both labeled and unlabeled samples. Then, it introduces a mixing mechanism where labeled samples are paired with their augmented versions, and the model's predictions on these mixed samples are combined using a mixing coefficient. MixMatch employs a consistency regularization strategy by enforcing consistency between the model's predictions on the augmented unlabeled samples and their corresponding original versions. This consistency loss penalizes discrepancies between predictions. Additionally, MixMatch introduces a sharpening mechanism to improve the quality of pseudo-labels assigned to unlabeled samples. This sharpening step adjusts the confidence of the model's predictions, making them more confident and reliable. The loss function of the MixMatch is given below along with its architecture:

Results:

UNET

First, we implemented UNET segmentation on the Retouchh dataset using only supervised learning. There were a total of 2685 images and masks. Of these we used 2148 for training, 161 for testing and 376 for validation. There were 4 classes every pixel in the image belonged to (3 types of cysts classes and one background class). 

This is the results for the best validation data when the model was saved.

Mean Teacher

We implemented the Mean Teacher architecture on the CIFAR - 10 dataset, The dataset comprises 10 classes and contains 50000 images. These classes are airplanes, cars, birds, cats, deer, dogs, frogs, horses, ships, and trucks. We split the data into 10000 labeled images and 40000 unlabeled images. We used the Resnet Architecture for the teacher and the student model. We used Gaussian noise on the input data as well. We also used 10000 images from the CIFAR 10 dataset to test our model. On training, we were able to achieve a loss around 1.17 when we ran the model for 100 epochs.

On the test set, we were able to achieve an accuracy of around 63 percent. 

Inferences:

In our study, we used semi-supervised convolutional neural networks (CNNs) to help spot retinal cysts, which are important in diagnosing eye problems. We tested our method on a bunch of different eye images and found that it's really good at finding cysts, even when we didn't have a lot of labeled data. By mixing in some unlabeled data and using semi-supervised learning  our model got even better at handling different kinds of eye images due to better regularization. Although our model ran on a few epochs we got good results which would definitely increase on further training. So, it seems like using semi-supervised CNNs could help doctors diagnose eye issues faster and more accurately. This could mean better tools and treatments for patients with eye problems.

References:

1. UNET_paper
2. Mean_Teacher_Paper
3. Fix_Match_Paper
4. Mix_Match_Paper
5. Fix_match_Reference2
6. Mix_Match_Reference2