/sseg-oct

Semantic segmentation systems for the detection of retinal pathological fluid from OCT images

Primary LanguagePython

Semantic segmentation systems for the detection of retinal pathological fluid from OCT images

This project was developed as a Computer Vision assignment in our last college year. Our professors provided us a dataset of 50 OCT images with their respective segmentation masks to detect pathological fluid from images of the retina. In the description of this assignment, they suggested us using U-Net architecture as our baseline and reviewing the SoTA in semantic segmentation to improve its results. We were expected to familiarize with the most used Deep Learning models of the day and experiment with diverse alternatives and techniques to research the performance of modern CV proposals in this problem. Then, we were asked to submit a paper with a detailed description of the tested approaches and the obtained results (check it here).

Segmentation models

Results

F-Score results:

Default benchmark Adversarial benchmark
No aug. Data aug. No aug. Data aug.
Baseline U-Net 0.60.15 0.650.16
U-Net ® 0.70.05 0.60.09 0.870.03 0.860.06
LinkNet ® 0.60.1 0.570.09 0.770.25 0.830.12
PSPNet ® 0.670.04 0.670.05 0.820.03 0.840.04
PAN ® 0.70.05 0.660.09
Attention U-Net 0.740.06 0.810.08
Deform U-Net 0.710.06 0.730.09

Intersection over Union (IoU) results:

Default benchmark Adversarial benchmark
No aug. Data aug. No aug. Data aug.
Baseline U-Net 0.440.13 0.560.11
U-Net ® 0.540.06 0.480.1 0.780.04 0.760.08
LinkNet ® 0.430.11 0.40.09 0.670.22 0.720.14
PSPNet ® 0.50.04 0.510.06 0.690.04 0.730.05
PAN ® 0.530.05 0.50.03
Attention U-Net 0.590.07 0.720.08
Deform U-Net 0.550.06 0.580.07

Pretrained models are marked with ®.

It is possible to check the number of parameters of the models:

python3 counter.py

Reproduce results

All code needed to reproduce our results is available in segmenter/ folder (if you require access to the original dataset, please contact me). The script test.py admits the following parameters to select the segmentation model:

python3 test.py <model> <mode> -v -aug -adv  
  • model: Specifies the model to train or test (choices are base, unet, linknet, pspnet, pan, attnunet, deformunet).
  • mode: Specifies the execution mode (kfold or train).
    • kfold accepts argument --k to run k fold cross validation. Final metrics are stored in the folder --model-path.
    • train executes a single train-validation split and the training procedure of the model selected and validating with the 10% of the dataset.
  • -v: Flag to show the training trace.
  • -aug: Flag to use data augmentation.
  • -adv: Flag to train the model in the adversarial benchmark.

Examples:

Execute the baseline model without data augmentation with batch size of 10 images and save results in ../results/base.

python3 test.py base kfold -v --batch_size=10 --model_path=../results/ base/  --route=../OCT-dataset/

Execute pretrained U-Net with data augmentation:

python3 test.py unet kfold -v --batch_size=10 -aug  

Execute LinkNet with data augmentation and adversarial learning:

python3 test.py unet train -v --batch_size=10 -aug  -adv

Note that the argument --route specifies the path where the OCT dataset is stored. The structure of this folder is expected to store two subfolders: images/ and masks/, which contain respectively the tomography retinal images used as input to the model and the pathological fluid masks used as output.

Team builders 👷