PyTorch Implementation of Small Lesion Segmentation in Brain MRIs with Subpixel Embedding (ORAL, MICCAIW 2021)

Overview

Small Lesion Segmentation in Brain MRIs with Subpixel Embedding

PyTorch implementation of Small Lesion Segmentation in Brain MRIs with Subpixel Embedding

To appear in MICCAI Brain Lesion Workshop 2021 (ORAL)

[publication] [arxiv] [poster] [talk]

Model have been tested on Ubuntu 16.04, 20.04 using Python 3.6, 3.7, PyTorch 1.7.0, 1.7.1

Authors: Alex Wong, Allison Chen, Yangchao Wu, Safa Cicek, Alexandre Tiard

If this work is useful to you, please cite our paper (to be updated from preprint to MICCAI):

@article{wong2021small,
  title={Small Lesion Segmentation in Brain MRIs with Subpixel Embedding},
  author={Wong, Alex and Chen, Allison and Wu, Yangchao and Cicek, Safa and Tiard, Alexandre and Hong, Byung-Woo and Soatto, Stefano},
  journal={arXiv preprint arXiv:2109.08791},
  year={2021}
}

Table of Contents

  1. About ischemic strokes
  2. About stroke lesion segmentation
  3. About Subpixel Network (SPiN)
  4. Setting up
  5. Downloading pretrained models
  6. Running SPiN
  7. Training SPiN
  8. Running and training related works
  9. License and disclaimer

About ischemic strokes

Ischemic strokes occur when a lack of blood flow prevents brain tissue from receiving adequate oxygen and nutrients and affect approximately 795,000 people annually. It’s currently the 5th leading cause of death in the United States and even upon recovery, many patients suffer from disability and even paralysis.

The severity of which is dependent on the size, location, and overlap of the lesions (caused by the stroke) with brain structures., but can be reduced with early diagnosis. Thus, preserving cognitive and motor function is dependent on identifying stroke lesions quickly and precisely, but doing so manually requires expert knowledge and is prohibitively time consuming. Here are some examples of manually traced lesions.

About stroke lesion segmentation

We focus on automatically identifying or segmenting ischemic stroke lesions from anatomical MRI images. Existing methods leverage deep convolutional neural networks that take as input an MRI image and output a heat or range map of confidence scores corresponding to lesion tissues. These lesions are often characterized by high variability in location, shape, and size. The latter two are problematic for conventional convolutional neural networks where the precision of irregularly shaped lesion boundaries and recall of small lesions are critical measures of success -- existing methods are able to capture the general shape of the large and medium sized lesions, but missed small ones and the intricacies in the boundary between lesion and normal tissue (see figure below).

Because deep neural networks are comprise of fixed sized filters with limited receptive fields that filters operate locally, in order to obtain a global representation of the input, one must employ spatial downsampling by the means of max pooling or strided convolutions. This is often performed aggressively so that the design does not limit applicability. Hence, the bottleneck or latent vector is typically much smaller in spatial dimensions than the input, which means that the details of fine local structures are lost in the process.

About Subpixel Network (SPiN)

To address this, we propose to learn an embedding that maps the input MRI image to high dimensional feature maps at double the input resolution where a single element or pixel in the image is represented by four subpixel feature vectors. Our subpixel embedding is comprised of a feature extraction phase using ResNet blocks and a spatial expansion phase that is achieved by rearranging elements from the channel dimension to the height and width dimensions using a depth to space operator.

Here we show the feature maps outputted by our subpixel embedding -- the high resolution feature maps contains anatomical structures that resembles the low resolution input MRI image.

This is particularly interesting because there is no explicit constraints on the structure of the learned representation. We note that we only minimize cross entropy, but our subpixel embedding naturally learns to preserve the details of the input as high resolution feature maps.

Unlike the typical architecture, we do not perform any lossy downsampling on this representation; hence, the embedding can preserve fine details of local structures, but in exchange, it lacks global context. But when combined with the standard encoder-decoder architecture, our subpixel embedding complements the decoder by guiding it with fine-grained local information that is lost during the encoding process so that the network can resolve detailed structures in the output segmentation. We refer to this architecture as a subpixel network or SPiN for short.

Unlike previous works, SPiN predicts an initial segmentation map at twice the input resolution where a single pixel in the input is predicted by a two by two neighborhood of subpixels in our output. The final segmentation at the original input resolution is obtained by combining the four subpixel predictions (akin to an ensemble) as a weighted sum where the contribution of each subpixel in a local neighborhood is predicted by a learnable downsampler.

Our learnable downsampler performs space-to-depth on the initial high resolution segmentation map and on each of channel of latent vector to rearrange each 2 by 2 neighborhood into a 4 element vector. This is followed by several convolutions and a softmax on the latent vector to learn the contribution or weight of each subpixel. The final segmentation at the input resolution is obtained by an element-wise dot product between the weights and the high resolution segmentation map.

Unlike hand crafted downsampling methods such as bilinear or nearest-neighbor interpolation where the weights are predetermined and independent of the input, our learnable downsampler is directly conditioned on the latent vector of the input and its initial segmentation. Hence it is able to predict weights that better respect lesion boundaries to yield detailed segmentation. Here we show qualitative comparisons on the ATLAS dataset.

Setting up your virtual environment

We will create a virtual environment with the necessary dependencies

virtualenv -p /usr/bin/python3.7 spin-py37env
source spin-py37env/bin/activate
pip install opencv-python scipy scikit-learn scikit-image matplotlib nibabel gdown numpy Pillow
pip install torch==1.7.1 torchvision==0.8.2 tensorboard==2.3.0

Setting up your datasets

For datasets, we will use the Anatomical Tracings of Lesion After Stroke (ATLAS) public dataset. You will need to sign an usage agreement in order to download the dataset. Please visit their dataset page, follow the instructions (sign a form and email it to the maintainers) to obtain the data, and save them on your machine.

Our repository provides the training and testing (referred to as traintest) data split in the data_split folder. The setup script will create numpy (.npy) files from .nii files provided in the raw dataset. It will also create training and validation directories containing text (.txt) files with relative paths to all data samples.

Assuming you are in the root of your repository, follow these command to setup the data:

# Create a data folder
mkdir data
ln -s /path/to/atlas data/

# Activate the virtual environment
source spin-py37env/bin/activate

# Set up the dataset
python setup/setup_dataset_atlas.py

Additionally, please run the following for the small lesion subset discussed in our paper:

python setup/setup_dataset_atlas_small_lesion.py

Downloading our pretrained models

To use pretrained models trained on ATLAS, please use gdown to download the zip file from our Google Drive.

# Download with gdown
gdown https://drive.google.com/uc?id=1ZSeLeNT-LX1n4QR-r8fVzHk3KmufSps_

# Unzip to directory
unzip pretrained_models.zip

This will output a pretrained_models directory containing the checkpoint for our model. For the ease of future comparisons between different works, we also provide checkpoints for various methods compared in our paper. All models are trained on the proposed training and testing split on ATLAS. The directory will have the following structure:

pretrained_models
|---- spin
      |---- spin_atlas_traintest.pth
|---- u-net
      |---- unet_atlas_traintest.pth
.
.
.

|---- kiu-net
      |---- kiunet_atlas_traintest.pth

Note: gdown fails intermittently and complains about permission. If that happens, you may also download the models via:

https://drive.google.com/file/d/1ZSeLeNT-LX1n4QR-r8fVzHk3KmufSps_/view?usp=sharing

Running SPiN

To run SPiN on the provided data splits (training, testing, small lesion subset), you may use our evaluation bash scripts located in the bash/spin directory.

To evaluate our pretrained model on our training and testing split:

bash bash/spin/evaluate/evaluate_spin_traintest.sh

The checkpoint should produce numbers consistent with Table 2 from our main paper:

Model DSC IOU Precision Recall
U-Net 0.584 0.432 0.674 0.558
D-UNet 0.548 0.404 0.652 0.521
CLCI-Net 0.599 0.469 0.741 0.536
KiU-Net 0.524 0.387 0.703 0.459
X-Net 0.639 0.495 0.746 0.588
SPiN (Ours) 0.703 0.556 0.806 0.654

To evaluate our pretrained model on our small lesion subset:

bash bash/spin/evaluate/evaluate_spin_small_lesion.sh

The checkpoint should produce numbers consistent with Table 3 from our main paper:

Model DSC IOU Precision Recall
U-Net 0.368 0.225 0.440 0.316
D-UNet 0.265 0.180 0.377 0.264
CLCI-Net 0.246 0.178 0.662 0.215
KiU-Net 0.246 0.255 0.466 0.206
X-Net 0.306 0.213 0.546 0.268
SPiN (Ours) 0.424 0.269 0.546 0.347

Note: You may replace the restore_path argument in these bash scripts to evaluate your own checkpoints.

By default, our evaluation scripts activate the flag --do_visualize_predictions. In addition to logging the evaluation summary in the checkpoint path directory, this will create a subdirectory called atlas_test_scans with the ground truths and predictions for each 2D image for visualization purposes. For quicker evaluation, you may remove this argument from the bash scripts.

Training SPiN

To train our method on ATLAS, you can directly run:

bash bash/spin/train/train_spin_traintest.sh

The bash script assumes that you are only using one GPU (CUDA_VISIBLE_DEVICES=0). To use multiple GPUs (if you want to train on a different dataset with larger images or using a larger batch size), you can modify the bash script with

export CUDA_VISIBLE_DEVICES=0,1

to use two GPUs.

Additionally, if you would like to train SPiN on your own dataset using our dataloader (src/datasets.py), you will need to convert your data into numpy (.npy) files and create files containing paths to each data sample as done in setup/setup_dataset_atlas.py.

Finally to monitor your training progress (training loss, visualizations, etc.), you can use tensorboard:

tensorboard --logdir trained_spin_models/<model_name>

Note: the frequency of logging tensorboard summaries and storing checkpoints are controlled by the n_summary and n_checkpoint arguments.

Running and training related works

To run the works we compare to on the provided data splits (training, testing, small lesion subset), you may use our evaluation bash scripts located in the bash directory. All of the checkpoints were the best (amongst several runs) produced by the methods below and should reproduce the numbers in the tables above (Table 2, 3 in main text).

U-Net

To evaluate the U-Net model the testing set:

bash bash/u-net/evaluate/evaluate_unet_traintest.sh

To evaluate the U-Net model on our small lesion subset:

bash bash/u-net/evaluate/evaluate_unet_small_lesion.sh

To train U-Net model on ATLAS on the proposed training and testing split:

bash bash/u-net/train/train_unet_traintest.sh

License and disclaimer

This software is property of the UC Regents, and is provided free of charge for research purposes only. It comes with no warranties, expressed or implied, according to these terms and conditions. For commercial use, please contact UCLA TDG.

You might also like...
Official pytorch implementation of paper "Inception Convolution with Efficient Dilation Search" (CVPR 2021 Oral).

IC-Conv This repository is an official implementation of the paper Inception Convolution with Efficient Dilation Search. Getting Started Download Imag

Official PyTorch implementation of RobustNet (CVPR 2021 Oral)
Official PyTorch implementation of RobustNet (CVPR 2021 Oral)

RobustNet (CVPR 2021 Oral): Official Project Webpage Codes and pretrained models will be released soon. This repository provides the official PyTorch

Pytorch implementation for
Pytorch implementation for "Adversarial Robustness under Long-Tailed Distribution" (CVPR 2021 Oral)

Adversarial Long-Tail This repository contains the PyTorch implementation of the paper: Adversarial Robustness under Long-Tailed Distribution, CVPR 20

Official PyTorch Implementation of Convolutional Hough Matching Networks, CVPR 2021 (oral)
Official PyTorch Implementation of Convolutional Hough Matching Networks, CVPR 2021 (oral)

Convolutional Hough Matching Networks This is the implementation of the paper "Convolutional Hough Matching Network" by J. Min and M. Cho. Implemented

Official pytorch implementation of "Feature Stylization and Domain-aware Contrastive Loss for Domain Generalization" ACMMM 2021 (Oral)

Feature Stylization and Domain-aware Contrastive Loss for Domain Generalization This is an official implementation of "Feature Stylization and Domain-

A PyTorch implementation of the baseline method in Panoptic Narrative Grounding (ICCV 2021 Oral)
A PyTorch implementation of the baseline method in Panoptic Narrative Grounding (ICCV 2021 Oral)

A PyTorch implementation of the baseline method in Panoptic Narrative Grounding (ICCV 2021 Oral)

PyTorch implementation for Score-Based Generative Modeling through Stochastic Differential Equations (ICLR 2021, Oral)
PyTorch implementation for Score-Based Generative Modeling through Stochastic Differential Equations (ICLR 2021, Oral)

Score-Based Generative Modeling through Stochastic Differential Equations This repo contains a PyTorch implementation for the paper Score-Based Genera

Paddle implementation for "Highly Efficient Knowledge Graph Embedding Learning with Closed-Form Orthogonal Procrustes Analysis" (NAACL 2021)

ProcrustEs-KGE Paddle implementation for Highly Efficient Knowledge Graph Embedding Learning with Orthogonal Procrustes Analysis πŸ™ˆ A more detailed re

Paddle implementation for "Cross-Lingual Word Embedding Refinement by β„“1 Norm Optimisation" (NAACL 2021)

L1-Refinement Paddle implementation for "Cross-Lingual Word Embedding Refinement by β„“1 Norm Optimisation" (NAACL 2021) πŸ™ˆ A more detailed readme is co

Owner
I am a post-doctoral researcher at the UCLA Vision Lab under the supervision of Professor Stefano Soatto.
null
An official implementation of "Exploiting a Joint Embedding Space for Generalized Zero-Shot Semantic Segmentation" (ICCV 2021) in PyTorch.

Exploiting a Joint Embedding Space for Generalized Zero-Shot Semantic Segmentation This is an official implementation of the paper "Exploiting a Joint

CV Lab @ Yonsei University 35 Oct 26, 2022
U-Net implementation in PyTorch for FLAIR abnormality segmentation in brain MRI

U-Net for brain segmentation U-Net implementation in PyTorch for FLAIR abnormality segmentation in brain MRI based on a deep learning segmentation alg

null 552 Nov 17, 2022
PrimitiveNet: Primitive Instance Segmentation with Local Primitive Embedding under Adversarial Metric (ICCV 2021)

PrimitiveNet Source code for the paper: Jingwei Huang, Yanfeng Zhang, Mingwei Sun. [PrimitiveNet: Primitive Instance Segmentation with Local Primitive

Jingwei Huang 45 Oct 13, 2022
This repo provides the official code for TransBTS: Multimodal Brain Tumor Segmentation Using Transformer (https://arxiv.org/pdf/2103.04430.pdf).

TransBTS: Multimodal Brain Tumor Segmentation Using Transformer This repo is the official implementation for TransBTS: Multimodal Brain Tumor Segmenta

Raymond 241 Nov 9, 2022
Self-supervised Multi-modal Hybrid Fusion Network for Brain Tumor Segmentation

JBHI-Pytorch This repository contains a reference implementation of the algorithms described in our paper "Self-supervised Multi-modal Hybrid Fusion N

FeiyiFANG 5 Dec 13, 2021
BraTs-VNet - BraTS(Brain Tumour Segmentation) using V-Net

BraTS(Brain Tumour Segmentation) using V-Net This project is an approach to dete

Rituraj Dutta 7 Jan 4, 2022
Official PyTorch Implementation of Embedding Transfer with Label Relaxation for Improved Metric Learning, CVPR 2021

Embedding Transfer with Label Relaxation for Improved Metric Learning Official PyTorch implementation of CVPR 2021 paper Embedding Transfer with Label

Sungyeon Kim 36 Nov 22, 2022
PyTorch Implementation of Google Brain's WaveGrad 2: Iterative Refinement for Text-to-Speech Synthesis

WaveGrad2 - PyTorch Implementation PyTorch Implementation of Google Brain's WaveGrad 2: Iterative Refinement for Text-to-Speech Synthesis. Status (202

Keon Lee 57 Nov 20, 2022
Unofficial PyTorch implementation of SimCLR by Google Brain

Unofficial PyTorch implementation of SimCLR by Google Brain

Rishabh Anand 2 Oct 13, 2021
Pytorch Implementation for NeurIPS (oral) paper: Pixel Level Cycle Association: A New Perspective for Domain Adaptive Semantic Segmentation

Pixel-Level Cycle Association This is the Pytorch implementation of our NeurIPS 2020 Oral paper Pixel-Level Cycle Association: A New Perspective for D

null 87 Oct 19, 2022