This repository includes an official implementation of paper: Advancing Radiograph Representation Learning with Masked Record Modeling (ICLR'23).
Some code is borrowed from MAE, huggingface, and REFERS.
Environmental requirements
-
Ubuntu 18.04 LTS.
-
Python 3.8.11
If you are using anaconda/miniconda, we provide an easy way to prepare the environment for pre-training and finetuning of classification:
conda env create -f environment.yaml
pip install -r requirements.txt
Download the pre-trained weight first!
import torch
import torch.nn as nn
from functools import partial
import timm
assert timm.__version__ == "0.6.12" # version check
from timm.models.vision_transformer import VisionTransformer
def vit_base_patch16(**kwargs):
model = VisionTransformer(norm_layer=partial(nn.LayerNorm, eps=1e-6),**kwargs)
return model
# model definition
model = vit_base_patch16(num_classes=14,drop_path_rate=0.1,global_pool="avg")
checkpoint_model = torch.load("./MRM.pth", map_location="cpu")["model"]
# load the pre-trained model
model.load_state_dict(checkpoint_model, strict=False)- We use MIMIC-CXR-JPG for pre-training. You can acquire more information about this dataset at Johnson et al. MIMIC-CXR-JPG.
- The dataset directory specified in run.sh includes the MIMIC-CXR-JPG dataset and you need to prepare a file
training.csvand put it into the dataset directory. - The file
training.csvincludes two columnsimage_pathandreport_contentfor each line, corresponding to (a) the path to an image and (b) the text of the corresponding report, respectively, which should be organized as follows:
image_path, report_content
/path/to/img1.jpg, FINAL REPORT EXAMINATION: ...
/path/to/img2.jpg, FINAL REPORT CHEST: ...
...,...
- take one line as an example:
-
Download the pre-trained weight of MAE and set
resumeto the path of the pre-trained weight in run.sh. -
Set the data path, GPU IDs, batch size, output directory, and other parameters in run.sh.
-
Start training by running
chmod a+x run.sh
./run.sh
- Download NIH ChestX-ray 14 dataset and split train/valid/test set. The directory should be organized as follows:
NIH_ChestX-ray/
all_classes/
xxxx1.png
xxxx2.png
...
xxxxn.png
train_1.txt
trian_10.txt
train_list.txt
val_list.txt
test_list.txt
- Specify the
dataset_pathin finetuning_1percent.sh, finetuning_10percent.sh, finetuning_100percent.sh, test.py.
-
Download the pre-trained weight from Google Drive and specify
pretrained_pathin finetuning_1percent.sh. -
Start training by running
chmod a+x finetuning_1percent.sh
./finetuning_1percent.sh
| RSNA | warm-up setps | total steps | learning rate |
|---|---|---|---|
| 1% | 50 | 2000 | 3e-3 |
| 10% | 200 | 10000 | 5e-4 |
| 100% | 2000 | 50000 | 5e-4 |
| CheXpert | warm-up setps | total steps | learning rate |
|---|---|---|---|
| 1% | 150 | 2000 | 3e-3 |
| 10% | 1500 | 60000 | 5e-4 |
| 100% | 15000 | 200000 | 5e-4 |
| Covid | warm-up setps | total steps | learning rate |
|---|---|---|---|
| 100% | 50 | 1000 | 3e-2 |
- Download SIIM-ACR Pneumothorax and preprocess the images and annotations. Then organize the directory as follows:
siim/
images/
training/
xxxx1.png
xxxx2.png
...
xxxxn.png
validation/
...
test/
...
annotations/
training/
xxxx1.png
xxxx2.png
...
xxxxn.png
validation/
...
test/
...
We conduct all experiments of segmentation by MMSegmentaiton (version 0.25.0) and it is necessary to set the environment and comprehend the code structures of MMSegmentaiton in advance.
Here we provide the necessary configuration files for reproducing the experiments in the directory Siim_Segmentation. After modifying MMSegmentaiton framework with provided files, start fine-tuning and evaluation with ft.sh and test.sh, respectively.
In the directory DatasetsSplits, we provide dataset splits that may be helpful for organizing the datasets.
We give the train/valid/test splits of CheXpert, NIH ChestX-ray, and RSNA Pneumonia.
For COVID-19 Image Data Collection, we randomly split the train/valid/test set 5 times and we provide the images in the images directory.
For SIIM-ACR_Pneumothorax, please organize the directories of images and annotations as section 5.1 mentioned according to the given splits.
React
Typescript
Django
Laravel
Facebook
Microsoft
Google
Alibaba
D3
Tencent