Skip to content

The official implementation of "Delving into Masked Autoencoders for Multi-Label Thorax Disease Classification"

License

Notifications You must be signed in to change notification settings

lambert-x/medical_mae

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

35 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Medical Masked Autoencoders

Paper

This repository provides the official implementation of training Vision Transformers (ViT) for (2D) medical imaging tasks as well as the usage of the pre-trained ViTs in the following paper:

Delving into Masked Autoencoders for Multi-Label Thorax Disease Classification
Junfei Xiao, Yutong Bai, Alan Yuille, Zongwei Zhou
Johns Hopkins University
IEEE/CVF Winter Conference on Applications of Computer Vision (WACV), 2023
paper | code

TO DO

  • Instructions for preparing datasets.
  • Instructions for pretraining and fine-tuning.

Image reconstruction demo

Installing Requirements

Our codebase follows the MAE Official and uses some additional packages. You may use one of the following commands to build environments with Conda and Pip.

Conda:

conda create -n medical_mae -f medical_mae.yml 

Pip:

conda create -n medical_mae python=3.8
conda activate medical_mae
pip install -r requirements.txt 

Preparing Datasets:

The MIMIC-CXR, CheXpert, and ChestX-ray14 datasets are public available on their official sites. You can download or request the access to them under the agreements.

You may also download them through the following links for research only and follow the official agreements.

MIMIC-CXR (JPG): https://physionet.org/content/mimic-cxr-jpg/2.0.0/

CheXpert (v1.0-small): https://www.kaggle.com/datasets/ashery/chexpert

ChestX-ray14 : https://www.kaggle.com/datasets/nih-chest-xrays/data

Pre-training on ImageNet or Chest X-rays

The pre-training instruction is in PRETRAIN.md.

Fine-tuning with pre-trained checkpoints

The fine-tuning instruction is in FINETUNE.md.

The following table provides the pre-trained checkpoints used in Table 1:

You can download all the weights in the following table with this link (google drive).

Model Pretrained Dataset Method Pretrained Finetuned (NIH Chest X-ray) mAUC
DenseNet-121 ImageNet Categorization torchvision official google drive 82.2
ResNet-50 ImageNet MoCo v2 google drive google drive 80.9
ResNet-50 ImageNet BYOL google drive google drive 81.0
ResNet-50 ImageNet SwAV google drive google drive 81.5
DenseNet-121 X-rays (0.3M) MoCo v2 google drive google drive 80.6
DenseNet 121 X-rays (0.3M) MAE google drive google drive 81.2
ViT-Small/16 ImageNet Categorization DeiT Official google drive 79.6
ViT-Small/16 ImageNet MAE google drive google drive 78.6
ViT-Small/16 X-rays (0.3M) MAE google drive google drive 82.3
ViT-Base/16 X-rays (0.5M) MAE google drive google drive 83.0
Model Pretrained Dataset Finetuned (Chest X-ray) mAUC Finetuned (CheXpert) mAUC Finetuned (COVIDx) Accuracy
ViT-Small/16 X-rays (0.3M) google drive 82.3 google drive 89.2 google drive 95.2
ViT-Base/16 X-rays (0.5M) google drive 83.0 google drive 89.3 google drive 95.3

Citation

If you use this code or use our pre-trained weights for your research, please cite our papers:

@inproceedings{xiao2023delving,
  title={Delving into masked autoencoders for multi-label thorax disease classification},
  author={Xiao, Junfei and Bai, Yutong and Yuille, Alan and Zhou, Zongwei},
  booktitle={Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision},
  pages={3588--3600},
  year={2023}
}

License

This repo is under Apache 2.0 license.

Acknowledgement

This work was supported by the Lustgarten Foundation for Pancreatic Cancer Research.

Our code is built upon facebookresearch/mae.

About

The official implementation of "Delving into Masked Autoencoders for Multi-Label Thorax Disease Classification"

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published