Skip to content

PyTorch implementation of Foveation for Segmentation of Ultra-High Resolution Images

License

Notifications You must be signed in to change notification settings

lxasqjc/Foveation-Segmentation

Repository files navigation

Foveation for Segmentation of Ultra-High Resolution Images in PyTorch

MIT License

This is a PyTorch implementation of Foveation for Segmentation of Ultra-High Resolution Images which extends our prior publication at MICCAI 2020.

Table of Contents

  1. Updates
  2. Introduction
  3. Foveation Module
  4. Citation
  5. Install dependencies
  6. Code Summary
  7. Direct-Executable
  8. Deploy-Your-Model
  9. Deploy-Foveation-Module-Locally
  10. Performance
  11. Acknowledgements

Updates

  • MICCAI 2020 talk on youtube available HERE
  • Pre-processed Gleason2019 datasets is shared HERE
  • We use configuration files to store most options which were in argument parser. The definitions of options are detailed in config/defaults.py (TO-BE-UPDATED).

Introduction

We demonstrates that the trade-off between FoV and resolution affects the segmentation performance on ultra high resolution images. Its influence also varies spatially according to local patterns in different areas. We demonstrate, on three publicly available high-resolution image datasets that our proposed Foveation Module combines the strength of the cases trained with patches of different fixed FoV/resolution trade-off and improves segmentation performance, below is a validation segmentation example on CityScapes compared ours against two fixed trade-off configurations:

Foveation Module

Foveation Module is a learnable dataloader which, for a given ultra high resolution image, adaptively chooses the appropriate configuration (FoV/resolution trade-off) of the input patch to feed to the downstream segmentation model at each spatial location of the image, illustrated as below:

We demonstrate that our Foveation Module can effectively learn spatial distribution ofthe FoV/Resolution trade-off, illustrated as below: teaser

Conceptually, the Foveation Module has an effect of adjusting patch-wise resolution at each location at optimal states for downstream task (segmentation in this case), illustrated as below:

Citation

If you use this code for your research, please cite our paper:

@misc{jin2020foveation,
    title={Foveation for Segmentation of Ultra-High Resolution Images},
    author={Chen Jin and Ryutaro Tanno and Moucheng Xu and Thomy Mertzanidou and Daniel C. Alexander},
    year={2020},
    eprint={2007.15124},
    archivePrefix={arXiv},
    primaryClass={cs.CV}
}

@inproceedings{jin2020foveation,
  title={Foveation for Segmentation of Mega-Pixel Histology Images},
  author={Jin, Chen and Tanno, Ryutaro and Xu, Moucheng and Mertzanidou, Thomy and Alexander, Daniel C},
  booktitle={International Conference on Medical Image Computing and Computer-Assisted Intervention},
  pages={561--571},
  year={2020},
  organization={Springer}
}

Install dependencies

Install dependencies with one of the following options:

  • Method 1: Pip installation:
python -m pip install -r requirements.txt
  • Method 2: Conda installation:
conda env create -f environment.yml
conda activate Fov_Seg

This code was tested with python 3.7, pytorch 1.2 and CUDA 11.0

Code Summary

We provide below three ways for users to easily run/deploy Foveation Segmentation: Direct-Executable, Deploy-Your-Model and Deploy-Foveation-Module-Locally.

  • Direct-Executable: full code direct executable to reproduce paper results, also suitable to run large scale experiments on customerized data in multiple GPU Environment, Syncronized Batch Normalization are deployed.
  • Deploy-Your-Model: Our Foveation Module can attach to any segmentation models. Our framework includes a set of State-of-the-Art segmentation models (PSPNet/UPerNet/Unet/HRNet) for user to choose. It's also straightforward to deploy personlized segmentation model in this framework, instruction are given.
  • Deploy-Foveation-Module-Locally: a Simple-Demo segmentation example to illustrated how to implement Foveation Module to your own code.

Direct-Executable

Data preparation

  1. Download example ultra-high resolution datasets in following table as used in the paper, or prepare of your choice ultra-high resolution images (e.g. medical/aerial/camera-array/material images) for multiclass segmentation purpose.

    Dataset (with Link) Content Resolution (pixels) Number of Classes
    Cityscapes urban scenes 2048x1024 19
    DeepGlobe aerial scenes 2448x2448 6
    Gleason2019 histopathological 5000x5000 4
  2. Your directory tree should be look like this:

$SEG_ROOT/data
├── cityscapes
│   ├── annotations
│   │   ├── testing
│   │   ├── training
│   │   └── validation
│   └── images
│       ├── testing
│       ├── training
│       └── validation
  1. Prepare data list file according to .odgt files in ./data/Cityscape/, a script ./data/write_MIT_format.m has been given to prepare .odgt file, PATH and suffix need to be customerized.

Training

Train a model by selecting the GPUs ($GPUS) and configuration file ($CFG) to use. During training, checkpoints by default are saved in folder ckpt.

python3 train_fove.py --gpus $GPUS --cfg $CFG
  • To choose which gpus to use, you can either do --gpus 0-7, or --gpus 0,2,4,6.
  • You can start with our provided configurations (cityscape example):
python3 train_fove.py --gpus GPUS --cfg config/foveation-cityscape-hrnetv2.yaml
  • Alternatively, you can quick start with provided bash script
bash quick_start_bash/cityscape_gumbel_softmax_1_10th.sh
  • You can also override options in commandline, for example python3 train_fove.py TRAIN.num_epoch 10 .

Evaluation

  1. Evaluate a trained model on the validation set. run eval_multipro.py with VAL.* arguments configured.
  • You can start with our provided configurations:
python3 eval_multipro.py --gpus GPUS --cfg config/foveation-cityscape-hrnetv2.yaml
  • Alternatively, you can quick start with provided bash script
bash quick_start_bash/deepglob_argmax_st_1_10th_eval.sh

Deploy-Your-Model

  1. The segmentation models in our framework are split into encoder and decoder, where encoders are usually modified directly from classification networks, and decoders consist of final convolutions and upsampling. Following architectures are supported:

    Encoder Decoder
    MobileNetV2dilated C1 (one convolution module)
    ResNet18/ResNet18dilated C1_deepsup (C1 + deep supervision trick)
    ResNet50/ResNet50dilated Pyramid Pooling Module
    ResNet101/ResNet101dilated PPM_deepsup (PPM + deep supervision trick)
    Unet UPerNet (Pyramid Pooling + FPN head)
    HRNetV2 (W48) -
  2. To deploy your personlized model, first prepare your model follow models/attention_u_net.py and put it under models/;

  3. Then simply add your model in def build_encoder and def build_decoder functions in models/models.py.

Deploy-Foveation-Module-Locally

A Simple-Demo segmentation example is given as follows:

  1. dataset for simple-demo are prepared at data/CITYSCAPE_single_example
  2. quick test simple demo with:
python3 train_sim_demo.py
  1. implement foveation module to your own code following train_sim_demo.py, includes:
  • import our libs:
from config import cfg
from dataset import TrainDataset, b_imresize
from models import ModelBuilder, SegmentationModule, FovSegmentationModule
from train_fove import checkpoint_last, train, create_optimizers
from eval import eval_during_train
  • add foveation module in network builder, train and evaluation (i.e. lines following if cfg.MODEL.foveation: in train_sim_demo.py)

Performance

  1. Segmentation performance measured in IoU/mIoU (%) on Cityscapes. For class names: All referred to all classes average, and the rest single class codes are as follows: DeepGlobe: Ro.-road, Sw.-sidewalk, Bu.-building, W.-wall, F.-fence, P.-pole, T.L.-traffic light, R.S.-raffic sign, V.-vegetation, Te.-terrain, Sk.-sky, P.-person, RI.-rider, C.-car, Tru.-truck, Tra.-train, M.-motorcycle, Bi.-bicycle.
Class All Ro. Sw. Bu. W. F. P. T.L. R.S. V. Link
Baseline-Patch 1 67.7 97.6 80.8 89.8 49.7 51.4 53.9 39.0 63.3 89.6 OneDrive/BaiduYun(Access Code:123)
Baseline-Patch 2 68.0 97.3 81.0 89.5 41.9 52.2 51.5 32.9 66.0 89.2 OneDrive/BaiduYun(Access Code:123)
Baseline-Patch 3 69.8 97.5 81.2 90.3 47.0 55.8 54.2 46.8 68.6 90.0 OneDrive/BaiduYun(Access Code:123)
Baseline-Patch 4 70.6 97.0 82.1 89.8 47.3 53.4 57.6 56.9 72.1 90.3 OneDrive/BaiduYun(Access Code:123)
Baseline-Patch 5 70.2 95.0 84.4 90.3 45.6 56.2 63.4 61.4 76.5 91.4 OneDrive/BaiduYun(Access Code:123)
Baseline-Random 50.3 94.1 68.1 80.7 36.5 40.2 18.8 13.3 31.8 79.6 OneDrive/BaiduYun(Access Code:123)
Baseline-Average 56.6 96.3 73.6 85.7 31.6 40.4 39.0 19.8 49.8 86.5 OneDrive/BaiduYun(Access Code:123)
Ensemble 72.5 97.7 83.8 91.1 51.7 58.2 58.9 50.9 72.6 90.9 OneDrive/BaiduYun(Access Code:123)
Ours-Mean 76.1 98.2 84.9 92.3 55.2 61.4 64.5 71.4 80.1 92.5 OneDrive/BaiduYun(Access Code:123)
Ours-Mode 71.4 97.5 80.8 90.8 45.4 55.3 55.6 56.5 67.7 90.7 OneDrive/BaiduYun(Access Code:123)
Ours-GSM 70.5 97.7 82.2 89.1 44.8 55.4 64.5 66.1 76.2 91.8 OneDrive/BaiduYun(Access Code:123)
Class Te. Sk. P. Ri. C. Tru. Bus Tra. M. Bi. Link
Baseline-Patch 1 59.4 90.4 73.9 51.6 91.9 54.9 74.0 61.3 45.3 69.1 OneDrive/BaiduYun(Access Code:123)
Baseline-Patch 2 55.9 90.1 75.3 55.4 92.0 54.1 80.2 74.8 43.7 69.8 OneDrive/BaiduYun(Access Code:123)
Baseline-Patch 3 57.4 88.8 77.1 56.9 92.8 57.4 76.8 63.6 51.5 72.5 OneDrive/BaiduYun(Access Code:123)
Baseline-Patch 4 56.7 85.6 79.0 60.0 92.8 54.0 78.2 58.1 56.8 74.2 OneDrive/BaiduYun(Access Code:123)
Baseline-Patch 5 59.9 68.9 81.2 58.0 93.2 53.0 69.6 50.6 59.3 76.8 OneDrive/BaiduYun(Access Code:123)
Baseline-Random 50.7 58.4 53.8 34.4 79.1 49.3 46.3 37.8 31.2 51.6 OneDrive/BaiduYun(Access Code:123)
Baseline-Average 47.5 89.3 65.7 43.6 89.0 47.6 45.7 38.6 23.2 61.6 OneDrive/BaiduYun(Access Code:123)
Ensemble 60.9 90.9 79.4 60.1 93.2 57.1 79.6 69.0 56.4 74.8 OneDrive/BaiduYun(Access Code:123)
Ours-Mean 64.2 94.1 82.8 64.0 94.5 63.6 81.0 63.3 59.2 78.2 OneDrive/BaiduYun(Access Code:123)
Ours-Mode 60.8 92.1 75.5 55.8 93.1 66.0 76.4 69.9 54.5 71.6 OneDrive/BaiduYun(Access Code:123)
Ours-GSM 62.1 75.1 80.8 60.8 92.8 53.4 74.8 42.6 52.9 76.6 OneDrive/BaiduYun(Access Code:123)
  1. Segmentation performance measured in IoU/mIoU (%) on DeepGlobe. For class names: All referred to all classes average, and the rest single class codes are as follows: U.-Urban, A.-Agriculture, R.-Rangeland, F.-Forest, W.-Water, B.-Barren.
Class All U. A. R. F. W. B. Link
Baseline-Patch 1 73.8 79.4 88.2 45.1 77.2 87.4 65.4 OneDrive/BaiduYun(Access Code:123)
Baseline-Patch 2 73.4 79.0 88.4 43.6 77.7 85.6 66.2 OneDrive/BaiduYun(Access Code:123)
Baseline-Patch 3 56.4 57.0 81.4 27.8 62.3 69.1 40.7 OneDrive/BaiduYun(Access Code:123)
Baseline-Patch 4 72.3 77.2 87.8 41.8 78.5 84.2 64.2 OneDrive/BaiduYun(Access Code:123)
Baseline-Patch 5 71.3 76.9 87.0 39.6 79.1 83.0 62.2 OneDrive/BaiduYun(Access Code:123)
Baseline-Random 46.4 56.3 77.0 20.1 41.0 45.8 38.5 OneDrive/BaiduYun(Access Code:123)
Baseline-Average 73.5 78.4 88.1 45.0 80.8 82.3 66.7 OneDrive/BaiduYun(Access Code:123)
Ensemble 74.3 79.4 88.7 44.5 80.0 86.3 66.8 OneDrive/BaiduYun(Access Code:123)
Ours-Mean 74.0 79.5 88.7 45.2 77.6 86.4 66.6 OneDrive/BaiduYun(Access Code:123)
Ours-Mode 73.2 77.5 88.2 43.4 79.9 84.4 65.5 OneDrive/BaiduYun(Access Code:123)
Ours-GSM 73.8 78.4 88.5 44.7 80.1 84.4 66.8 OneDrive/BaiduYun(Access Code:123)
GLnet 71.6 - - - - - -
  1. Segmentation performance measured in IoU/mIoU (%) on Gleason2019.
Class All Benign Grade 3 Grade 4 Link
Baseline-Patch 1 69.3 81.0 61.8 65.0 OneDrive/BaiduYun(Access Code:123)
Baseline-Patch 2 70.1 80.5 64.4 65.3 OneDrive/BaiduYun(Access Code:123)
Baseline-Patch 3 70.7 81.0 64.9 66.1 OneDrive/BaiduYun(Access Code:123)
Baseline-Patch 4 67.3 76.4 61.4 64.0 OneDrive/BaiduYun(Access Code:123)
Baseline-Patch 5 61.4 67.5 55.6 61.0 OneDrive/BaiduYun(Access Code:123)
Baseline-Random 64.9 72.2 58.6 63.8 OneDrive/BaiduYun(Access Code:123)
Baseline-Average 65.7 78.8 52.2 66.0 OneDrive/BaiduYun(Access Code:123)
Ensemble 71.0 80.8 65.4 66.9 -
Ours-Mean 71.1 82.4 63.0 67.8 OneDrive/BaiduYun(Access Code:123)
Ours-Mode 69.5 81.6 61.3 65.6 OneDrive/BaiduYun(Access Code:123)
Ours-GSM 67.4 79.4 58.3 64.4 OneDrive/BaiduYun(Access Code:123)
Clinical experts 66.9 83.9 56.4 60.3 -

Acknowledgements

Releases

No releases published

Packages

No packages published

Languages