An implementation of the research paper "Retina Blood Vessel Segmentation Using A U-Net Based Convolutional Neural Network"

Overview

Retina Blood Vessels Segmentation

This is an implementation of the research paper "Retina Blood Vessel Segmentation Using A U-Net Based Convolutional Neural Network" written by Wang Xiancheng, Li Weia, et al.

Check out the standalone demo notebook and run segRetino inferences here.

Open In Colab

Inspiration

Various eye diseases can be diagnosed through the characterization of the retinal blood vessels. The characterization can be extracted by using proper imaging techniques and data analysis methods. In case of eye examination, one of the important tasks is the retinal image segmentation.The paper presents a network and training strategy that relies on the data augmentation to use the available annotated samples more efficiently, to segment retinal blood vessels using a UNET convolutional neural network.

Dataset

We have used the Digital Retinal Images for Vessel Extraction (DRIVE) dataset for retinal vessel segmentation. It consists of a total of JPEG 40 color fundus images; including 7 abnormal pathology cases. Each image resolution is 584x565 pixels with eight bits per color channel (3 channels), resized to 512x512 for our model.

Guidelines to download, setup and use the dataset

The DRIVE dataset may be downloaded here as two files named training.zip and test.zip.

Please write the following commands on your terminal to extract the file in the proper directory.

  $ mkdir drive
  $ unzip </path/to/training.zip> -d </path/to/drive>
  $ unzip </path/to/test.zip> -d </path/to/drive>

The resulting directory structure should be:

/path/to/drive
    -> train
        -> image
            -> 21_training_0.tif
            -> 22_training_0.tif
               ...
        -> mask
            -> 21_training_0.gif
            -> 22_training_0.gif
    -> test
        -> image
            -> 01_test_0.tif
            -> 02_test_0.tif
               ...
        -> mask
            -> 01_test_0.gif
            -> 02_test_0.gif

Model Components

The UNET CNN architecture may be divided into the Encoder, Bottleneck and Decoder blocks, followed by a final segmentation output layer.

  • Encoder: There are 4 Encoder blocks, each consisting of a convolutional block followed by a Spatial Max Pooling layer.
  • Bottleneck: The Bottleneck consists of a single convolutional block.
  • Decoder: There are 4 Decoder blocks, each consisting of a deconvolution operation, followed by a convolutional block, along with skip connections.

Note: The convolutional block consists of 2 conv2d operations each followed by a BatchNorm2d, finally followed by a ReLU activation.

model_arch

Implementation Details

  • Image preprocessing included augmentations like HorizontalFlip, VerticalFlip, Rotate.
  • Dataloader object was created for both training and validation data
  • Training process was carried out for 50 epochs, using the Adam Optimizer with a Learning Rate 1e-4.
  • Validation was carried out using Dice Loss and Intersection over Union Loss.

Installation and Quick Start

To use the repo and run inferences, please follow the guidelines below

  • Cloning the Repository:

      $ git clone https://github.com/srijarkoroy/segRetino
    
  • Entering the directory:

      $ cd segRetino/
    
  • Setting up the Python Environment with dependencies:

      $ pip install -r requirements.txt
    
  • Running the file for inference:

      $ python3 test.py
    

Running the test file downloads the pretrained weights of the UNET Model that we have trained on the DRIVE Dataset. However if you want to re-train the model please mention the path to your dataset on you local machine after augmentations, inside the train.py file, as:

train_x = sorted(glob(<path/to/augmented/train/image/folder/>))
train_y = sorted(glob(<path/to/augmented/mask/image/folder/>))

valid_x = sorted(glob(<path/to/test/image/folder/>))
valid_y = sorted(glob(<path/to/test/mask/folder/>))

Once the path has been mentioned, the model may be trained by running the command:

  $ python3 train.py

Note: If images have not been augmented, please see the instructions for augmentation here.

The test file saves two images in the mentioned paths, a masked image showing only the blood vessels, and a blend image showing the blood vessels within the retina. If you don't want to save the blend image, consider running the following code snippet:

# Creating the SegRetino object initialized with the test image path
seg = SegRetino('<path/to/test/img>')

# Running inference
seg.inference(set_weight_dir = 'unet.pth', path = '<path/to/save/masked/image>', blend=False, blend_path = None)

Check out the standalone demo notebook and run segRetino inferences here.

Note: Is is advisable to use a GPU for running the inferences since performing segmentation on 512x512 images with a heavy UNET architecture is expensive.

Results from Implementation

Original Image Masked Image Blend Image

Contributors

Contribution

Contributions are always welcome! Please check out this doc for Contribution Guidelines.

You might also like...
The lightweight PyTorch wrapper for high-performance AI research. Scale your models, not the boilerplate.
The lightweight PyTorch wrapper for high-performance AI research. Scale your models, not the boilerplate.

The lightweight PyTorch wrapper for high-performance AI research. Scale your models, not the boilerplate. Website • Key Features • How To Use • Docs •

A research toolkit for particle swarm optimization in Python
A research toolkit for particle swarm optimization in Python

PySwarms is an extensible research toolkit for particle swarm optimization (PSO) in Python. It is intended for swarm intelligence researchers, practit

Plato: A New Framework for Federated Learning Research

a new software framework to facilitate scalable federated learning research.

Research shows Google collects 20x more data from Android than Apple collects from iOS. Block this non-consensual telemetry using pihole blocklists.

pihole-antitelemetry Research shows Google collects 20x more data from Android than Apple collects from iOS. Block both using these pihole lists. Proj

A Research-oriented Federated Learning Library and Benchmark Platform for Graph Neural Networks. Accepted to ICLR'2021 - DPML and MLSys'21 - GNNSys workshops.

FedGraphNN: A Federated Learning System and Benchmark for Graph Neural Networks A Research-oriented Federated Learning Library and Benchmark Platform

This repository contains the implementations related to the experiments of a set of publicly available datasets that are used in the time series forecasting research space.

TSForecasting This repository contains the implementations related to the experiments of a set of publicly available datasets that are used in the tim

This is the research repository for Vid2Doppler: Synthesizing Doppler Radar Data from Videos for Training Privacy-Preserving Activity Recognition.
This is the research repository for Vid2Doppler: Synthesizing Doppler Radar Data from Videos for Training Privacy-Preserving Activity Recognition.

Vid2Doppler: Synthesizing Doppler Radar Data from Videos for Training Privacy-Preserving Activity Recognition This is the research repository for Vid2

Automatic voice-synthetised summaries of latest research papers on arXiv

PaperWhisperer PaperWhisperer is a Python application that keeps you up-to-date with research papers. How? It retrieves the latest articles from arXiv

A Dataset of Python Challenges for AI Research

Python Programming Puzzles (P3) This repo contains a dataset of python programming puzzles which can be used to teach and evaluate an AI's programming

Owner
Srijarko Roy
AI Enthusiast!
Srijarko Roy
A modular framework for vision & language multimodal research from Facebook AI Research (FAIR)

MMF is a modular framework for vision and language multimodal research from Facebook AI Research. MMF contains reference implementations of state-of-t

Facebook Research 5.1k Jan 4, 2023
The official implementation of the research paper "DAG Amendment for Inverse Control of Parametric Shapes"

DAG Amendment for Inverse Control of Parametric Shapes This repository is the official Blender implementation of the paper "DAG Amendment for Inverse

Elie Michel 157 Dec 26, 2022
A sample pytorch Implementation of ACL 2021 research paper "Learning Span-Level Interactions for Aspect Sentiment Triplet Extraction".

Span-ASTE-Pytorch This repository is a pytorch version that implements Ali's ACL 2021 research paper Learning Span-Level Interactions for Aspect Senti

来自丹麦的天籁 10 Dec 6, 2022
Pytorch port of Google Research's LEAF Audio paper

leaf-audio-pytorch Pytorch port of Google Research's LEAF Audio paper published at ICLR 2021. This port is not completely finished, but the Leaf() fro

Dennis Fedorishin 80 Oct 31, 2022
Research code for CVPR 2021 paper "End-to-End Human Pose and Mesh Reconstruction with Transformers"

MeshTransformer ✨ This is our research code of End-to-End Human Pose and Mesh Reconstruction with Transformers. MEsh TRansfOrmer is a simple yet effec

Microsoft 473 Dec 31, 2022
Multiple paper open-source codes of the Microsoft Research Asia DKI group

?? Paper Code Collection (MSRA DKI Group) This repo hosts multiple open-source codes of the Microsoft Research Asia DKI Group. You could find the corr

Microsoft 249 Jan 8, 2023
Research Artifact of USENIX Security 2022 Paper: Automated Side Channel Analysis of Media Software with Manifold Learning

Manifold-SCA Research Artifact of USENIX Security 2022 Paper: Automated Side Channel Analysis of Media Software with Manifold Learning The repo is org

Yuanyuan Yuan 172 Dec 29, 2022
Reference implementation of code generation projects from Facebook AI Research. General toolkit to apply machine learning to code, from dataset creation to model training and evaluation. Comes with pretrained models.

This repository is a toolkit to do machine learning for programming languages. It implements tokenization, dataset preprocessing, model training and m

Facebook Research 408 Jan 1, 2023
The lightweight PyTorch wrapper for high-performance AI research. Scale your models, not the boilerplate.

The lightweight PyTorch wrapper for high-performance AI research. Scale your models, not the boilerplate. Website • Key Features • How To Use • Docs •

Pytorch Lightning 21.1k Jan 1, 2023