Git Product home page Git Product logo

cv-ssl-mis's Introduction

CV-SSL-MIS

Exploring CNN and ViT for Semi-Supervised Medical Image Segmentation

Requirements

  • [Pytorch]
  • [MONAI]
  • Some basic python packages: Numpy, Scikit-image, SimpleITK, Scipy, Medpy, nibabel, tqdm ......

News and To-do List

  • Contrastive Learning
  • Various Segmentation Backbone Networks
  • Various Semi-Supervised Frameworks
  • 3D Segmentation ViT Backbone Networks
  • 2D SwinUNet
  • 2D SegFormer
  • Totalsegmentor dataset

Usage

  1. Clone the repo:
git clone https://github.com/ziyangwang007/CV-SSL-MIS.git 
cd CV-SSL-MIS
  1. Download the pre-processed data and put the data in ../data/BraTS2019 or ../data/ACDC or ../data/TotalSegmentator. In this project, we use ACDC, TotalSegmentator for 2D purpose, and BraTS for 3D purpose. You can download the dataset with the list of labeled training, unlabeled training, validation, and testing slices as following:

ACDC from Google Drive Link, or Baidu Netdisk Link with passcode: 'kafc'.

BraTS from Google Drive Link, or Baidu Netdisk Link with passcode: 'kbj3'.

TotalSegmentator from zenodo

  1. Train the model
cd code

You can choose model(unet/vnet/pnet/unetr...) by '--model', dataset(acdc/brats) by '--root_path', ratio of labeled/unlabel training set(10%, 20%, 30%, 50%) by '--labeled_num', experiment name(the path of saving your model weights and inference) by '--exp', iteration number, batch size, multi-class classification and etc in your command line, or leave it with default option.

Fully Supervised - CNN (UNet) -> Paper Link

python train_fully_supervised_2D.py --root_path ../data/ACDC --exp ACDC/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 4 --labeled_num XXX

python train_fully_supervised_3D.py --root_path ../data/BraTS2019 --exp BraTS/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 2 --labeled_num XXX

Fully Supervised - ViT (SwinUNet) -> Paper Link

python train_fully_supervised_2D_ViT.py --root_path ../data/ACDC --exp ACDC/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 4 --labeled_num XXX

Mean Teacher - CNN -> Paper Link

python train_mean_teacher_2D.py --root_path ../data/ACDC --exp ACDC/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 4 --labeled_num XXX

python train_mean_teacher_3D.py --root_path ../data/BraTS2019 --exp BraTS/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 2 --labeled_num XXX

Mean Teacher - ViT -> Paper Link

python train_mean_teacher_ViT.py --root_path ../data/ACDC --exp ACDC/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 4 --labeled_num XXX

Uncertainty-Aware Mean Teacher - CNN -> Paper Link

python train_uncertainty_aware_mean_teacher_2D.py --root_path ../data/ACDC --exp ACDC/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 4 --labeled_num XXX

python train_uncertainty_aware_mean_teacher_3D.py --root_path ../data/BraTS2019 --exp BraTS/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 2 --labeled_num XXX

Uncertainty-Aware Mean Teacher - ViT -> Paper Link

python train_uncertainty_aware_mean_teacher_ViT_2D.py --root_path ../data/ACDC --exp ACDC/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 4 --labeled_num XXX

Adversarial - CNN -> Paper Link

python train_adversarial_network_2D.py --root_path ../data/ACDC --exp ACDC/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 4 --labeled_num XXX

python train_adversarial_network_3D.py --root_path ../data/BraTS2019 --exp BraTS/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 2 --labeled_num XXX

Adversarial - ViT

python train_adversarial_network_2D_ViT.py --root_path ../data/ACDC --exp ACDC/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 4 --labeled_num XXX

Cross Pseudo Supervision CNN -> Paper Link

python train_cross_pseudo_supervision_2D.py --root_path ../data/ACDC --exp ACDC/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 4 --labeled_num XXX

python train_cross_pseudo_supervision_3D.py --root_path ../data/BraTS2019 --exp BraTS/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 2 --labeled_num XXX

Cross Pseudo Supervision - ViT CNN -> Paper Link

python train_cross_teaching_between_cnn_transformer_2D.py --root_path ../data/ACDC --exp ACDC/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 4 --labeled_num XXX

Cross Pseudo Supervision - ViT -> Paper Link

python train_cross_pseudo_supervision_2D_ViT.py --root_path ../data/ACDC --exp ACDC/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 4 --labeled_num XXX

Contrastive Learning - Cross Pseudo Supervision - CNN ViT

python train_Contrastive_Cross_CNN_ViT_2D.py --root_path ../data/ACDC --exp ACDC/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 4 --labeled_num XXX

Contrastive Learning - Cross Pseudo Supervision - CNN -> Paper Link

python train_Contrastive_Cross_CNN_2D.py --root_path ../data/ACDC --exp ACDC/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 4 --labeled_num XXX

Fixmatch - CNN -> Paper Link

python train_Fixmatch_CNN_2D.py --root_path ../data/ACDC --exp ACDC/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 4 --labeled_num XXX

Contrastive Learning - Fixmatch - Mean Teacher - ViT -> Paper Link

python train_Contrastive_Consistency_ViT_2D.py --root_path ../data/ACDC --exp ACDC/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 4 --labeled_num XXX

Adversarial Consistency - ViT -> Paper Link

python train_adversarial_consistency_ViT_2D.py --root_path ../data/ACDC --exp ACDC/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 4 --labeled_num XXX

Co-Training - CNN -> Paper Link

python train_deep_co_training_2D.py --root_path ../data/ACDC --exp ACDC/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 4 --labeled_num XXX

Co-Training - ViT

python train_deep_co_training_2D_ViT.py --root_path ../data/ACDC --exp ACDC/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 4 --labeled_num XXX

MixUp - CNN -> Paper Link

python train_interpolation_consistency_training_2D.py --root_path ../data/ACDC --exp ACDC/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 4 --labeled_num XXX

python train_interpolation_consistency_training_3D.py --root_path ../data/BraTS2019 --exp BraTS/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 2 --labeled_num XXX

MixUp - ViT

python train_interpolation_consistency_training_2D_ViT.py --root_path ../data/ACDC --exp ACDC/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 4 --labeled_num XXX

Semi CNN-ViT -> Paper Link

python train_cnn_meet_vit_2D.py --root_path ../data/ACDC --exp ACDC/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 4 --labeled_num XXX

Triple-View Segmentation CNN -> Paper Link

python train_tripleview_2D(demo).py --root_path ../data/ACDC --exp ACDC/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 4 --labeled_num XXX

Examiner-Student-Teacher CNN -> Paper TBC (after October)

python train_exam_student_teacher_3D.py --root_path ../data/ACDC --exp ACDC/XXX --model XXX -max_iterations XXX -batch_size XXX --base_lr XXX --num_classes 2 --labeled_num XXX
  1. Test the model
python test_2D_fully.py -root_path ../data/XXX --exp ACDC/XXX -model XXX --num_classes 4 --labeled_num XXX

python test_3D.py -root_path ../data/XXX --exp ACDC/XXX -model XXX --num_classes 4 --labeled_num XXX

python test_CNNVIT.py -root_path ../data/XXX --exp ACDC/XXX -model XXX --num_classes 4 --labeled_num XXX

Check trained model and inference

cd model

Acknowledgement

This code is mainly based on SSL4MIS, MONAI.

Some of the other code is from SegFormer, SwinUNet, Segmentation Models, UAMT, nnUNet.

cv-ssl-mis's People

Contributors

ziyangwang007 avatar

Recommend Projects

  • React photo React

    A declarative, efficient, and flexible JavaScript library for building user interfaces.

  • Vue.js photo Vue.js

    🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.

  • Typescript photo Typescript

    TypeScript is a superset of JavaScript that compiles to clean JavaScript output.

  • TensorFlow photo TensorFlow

    An Open Source Machine Learning Framework for Everyone

  • Django photo Django

    The Web framework for perfectionists with deadlines.

  • D3 photo D3

    Bring data to life with SVG, Canvas and HTML. 📊📈🎉

Recommend Topics

  • javascript

    JavaScript (JS) is a lightweight interpreted programming language with first-class functions.

  • web

    Some thing interesting about web. New door for the world.

  • server

    A server is a program made to process requests and deliver data to clients.

  • Machine learning

    Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.

  • Game

    Some thing interesting about game, make everyone happy.

Recommend Org

  • Facebook photo Facebook

    We are working to build community through open source technology. NB: members must have two-factor auth.

  • Microsoft photo Microsoft

    Open source projects and samples from Microsoft.

  • Google photo Google

    Google ❤️ Open Source for everyone.

  • D3 photo D3

    Data-Driven Documents codes.