• Stars
    star
    356
  • Rank 119,446 (Top 3 %)
  • Language
    Python
  • License
    MIT License
  • Created over 4 years ago
  • Updated over 1 year ago

Reviews

There are no reviews yet. Be the first to send feedback to the community and the maintainers!

Repository Details

Official Pytorch Code of KiU-Net for Image/3D Segmentation - MICCAI 2020 (Oral), IEEE TMI

KiU-Net-pytorch PWC

Paper (Conference) | Paper (Journal Extension) | Project Page | Slides

Official Pytorch Code for the paper "KiU-Net: Towards Accurate Segmentation of Biomedical Images using Over-complete Representations" , presented at MICCAI 2020 and its

Journal Extension: "KiU-Net: Overcomplete Convolutional Architectures for Biomedical Image and Volumetric Segmentation" , IEEE Transactions on Medical Imaging (IEEE TMI)

About this repo:

This repo hosts the code for the following networks:

  1. KiU-Net 2D
  2. KiU-Net 3D
  3. Res-KiU-Net
  4. Dense-KiU-Net

It also has dataloaders organized for generic 2D image segmentation and 3D volumetric segmentation for BraTS, LiTS dataset; for easy benchmarking of medical image and volumetric segmentation algorithms.

Introduction

In a generic "encoder-decoder" architecture , the initial few blocks of the encoder learn low-level features of the data while the later blocks learn the high-level features. Eventually, the encoder learns to map the data to lower dimensionality (in the spatial sense). The increasing receptive field size over the depth of the network, constrains the network to focus more on the higher-level features. In our proposed architecture , we introduce Ki-Net where we use overcomplete representations which constraints the receptive field from increasing. This is done by a simple change in the architecture of encoder where max-pooling is replaced by up-sampling. This helps the filters in deep layers focus more on the low-level details helping in fine segmentation. This network, when augmented with U-Net is termed as KiU-Net which results in significant improvements in the case of segmenting small anatomical landmarks and blurred noisy boundaries while obtaining better overall performance. More details can be found in the paper.

Prerequisites:

  • Python 3.6
  • Pytorch 1.4

Pytorch Installation

Links for downloading the public Datasets:

  1. GLAS Dataset - Link (Original) | Link (Resized)
  2. RITE Dataset - Link (Original) | Link (Resized)
  3. BraTS Dataset - Link
  4. LiTS Dataset - Link
  5. Unfortunately, Brain Anatomy US dataset introduced in the paper cannot be made public because of pending IRB approval.

Using the Code for your dataset

  • Clone this repository:
git clone https://github.com/jeya-maria-jose/KiU-Net-pytorch
cd KiU-Net-pytorch

Dataset Preparation

Prepare the dataset in the following format for easy use of the code. The train and test folders should contain two subfolders each: img and label. Make sure the images their corresponding segmentation masks are placed under these folders and have the same name for easy correspondance. Please change the data loaders to your need if you prefer not preparing the dataset in this format.

Train Folder-----
      img----
          0001.png
          0002.png
          .......
      label---
          0001.png
          0002.png
          .......
Validation Folder-----
      img----
          0001.png
          0002.png
          .......
      label---
          0001.png
          0002.png
          .......
Test Folder-----
      img----
          0001.png
          0002.png
          .......
      label---
          0001.png
          0002.png
          .......
  • This code is written for binary segmentation of an grayscale image. For using RGB images, just change the number of channels in the first and last layers in kiunet class found in arch/ae.py . For using KiU-Net for segmentation of more than one class, change the number of channels in the softmax layer in kiunet class found in arch/ae.py . The code is written for images of dimenstions 128x128. Either resize your images to the exact dimension for using the code directly or make changes in the ratios in interpolation functions accordingly in case of using dimensions which are not powers of 2.

  • The ground truth images should have pixels corresponding to the labels. Example: In case of binary segmentation, the pixels in the GT should be 0 or 1.

Training Command:

python train.py --train_dataset "enter train directory" --val_dataset "enter validation directory" --direc 'path for results to be saved' --batch_size 1 --epoch 400 --save_freq 10 --modelname "kiunet" --learning_rate 0.0001
Change modelname to kiunet3d, reskiunet or densekiunet to train those architectures

Testing Command:

python test.py --loaddirec "./saved_model_path/model_name.pth" --val_dataset "test dataset directory" --direc 'path for results to be saved' --batch_size 1 --modelname "kiunet"

The results including predicted segmentations maps will be placed in the results folder along with the model weights. Run the performance metrics code in MATLAB for calculating DICE Coefficient and Jaccard Index.

KiU-Net 3D:

The model file of KiU-Net 3D can be found at arch/ae.py .

Using the Code for BraTS Dataset

Check out readme in "BRATS/" folder for details on training KiU-Net 3D for BraTS dataset.

Using the Code for LiTS Dataset

Check out readme in "LiTS/" folder for details on training KiU-Net 3D for LiTS dataset.

Acknowledgement:

The dataloader code is inspired from pytorch-UNet , BraTS 2018 dataloader and LiTS dataloader.

Citation:

@inproceedings{valanarasu2020kiu,
  title={KiU-Net: Towards Accurate Segmentation of Biomedical Images Using Over-Complete Representations},
  author={Valanarasu, Jeya Maria Jose and Sindagi, Vishwanath A and Hacihaliloglu, Ilker and Patel, Vishal M},
  booktitle={Medical Image Computing and Computer Assisted Intervention--MICCAI 2020: 23rd International Conference, Lima, Peru, October 4--8, 2020, Proceedings, Part IV 23},
  pages={363--373},
  year={2020},
  organization={Springer}
}
@article{valanarasu2021kiu,
  title={Kiu-net: Overcomplete convolutional architectures for biomedical image and volumetric segmentation},
  author={Valanarasu, Jeya Maria Jose and Sindagi, Vishwanath A and Hacihaliloglu, Ilker and Patel, Vishal M},
  journal={IEEE Transactions on Medical Imaging},
  year={2021},
  publisher={IEEE}
}

Open an issue in case of any queries or mail me directly.

More Repositories

1

Medical-Transformer

Official Pytorch Code for "Medical Transformer: Gated Axial-Attention for Medical Image Segmentation" - MICCAI 2021
Python
797
star
2

UNeXt-pytorch

Official Pytorch Code base for "UNeXt: MLP-based Rapid Medical Image Segmentation Network", MICCAI 2022
Python
462
star
3

TransWeather

Pytorch Code for the paper TransWeather - CVPR 2022
Python
160
star
4

Cuff_less_BP_Prediction

Prediction of Blood Pressure from ECG and PPG signals using regression methods.
Python
146
star
5

Interactive-Portrait-Harmonization

Code Base for the work "Interactive Portrait Harmonization"
28
star
6

On-The-Fly-Adaptation

Code base for "On-the-Fly Test-time Adaptation for Medical Image Segmentation"
Python
26
star
7

Derain_OUCD_Net

Official Pytorch Code for "Exploring Overcomplete Representations for Single Image Deraining using CNNs" - IEEE Journal of STSP
Python
14
star
8

Overcomplete-Deep-Subspace-Clustering

Official Tensorflow Code for the paper "Overcomplete Deep Subspace Clustering Networks" - WACV 2021
Python
13
star
9

Image-Recovery-Using-Conditional-Adversarial-Networks

Analyzing Conditional Adversarial Networks to solve image recovery problems like shadow recovery, denoising and deblurring - CVIP 2019
Python
10
star
10

Unet_DWT

Unet based on Wavelet coefficients for segmentation
Python
8
star
11

essential_codes

Just a collection of essential set of codes that can be plugged in at needed places
Python
3
star
12

rebotnet-web

Website for Rebotnet
JavaScript
3
star
13

sparseSGD

Tweaking SGD by imposing sparsity to improve opimization for deep learning tasks.
Python
3
star
14

transweather-web

Website for TransWeather paper: https://jeya-maria-jose.github.io/transweather-web/
HTML
2
star
15

Vision-Based-Texting

Approach towards texting using eye gestures taken from a camera.
Python
1
star
16

Wind_LandArea_Aerial_Seg

Supporting Code for the paper "A novel application of deep learning to determine the actual land transformed by wind power"
Python
1
star
17

IPH-web

Website for "Interactive Portrait Harmonization"
JavaScript
1
star
18

quad-pid-control-

Python
1
star
19

Performance_Metrics

Python
1
star
20

rebot-net

Network architecture for ReBotNet, WACV 2025
Python
1
star