MCR_RRR FLARE21 submission: Cpu-Only aBdominal oRgan segmentAtion (COBRA)

Training data preprocessing

Before training our model, the training data needs to be pre-processing. For this use train_preprocessing.py in the parent directory.
This script will:

convert images and gold standard segmentations from .nii.gz to .npy format

correct orientation errors

segment the "body" using thresholding

resample all images to the same dimensions (96 x 192 x 192 voxels)

save the resulting data to a directory ready for training

NOTE - Three directory paths need to be changed here:

imdir: directory containing the raw .nii.gz format CT images

maskdir: directory containing the raw .nii.gz format gold standard segmentations

out_dir: where the preprocessed images and masks will be saved

Model training

Once pre-processed, the data is ready for model training!
To train a model, use python train.py --fold_num i where i is an integer in [1,2,3,4,5]
or sh train_multi_folds.sh

NOTE - Two directory paths need to be changed here:

source_dir: where the preprocessed images and masks are saved by the preprocessing script

checkpoint_dir: where model weights and training statitics will be saved

Current hyper-parameter settings in train.py will reproduce our submitted model.

ONNX compilation + Inference

Compile

Once the model has been trained, it can be compiled by running cd ONNX && bash ./export.sh.
ARGS:

MODEL_PATH: path to trained model.

OUTPUT_NAME: compiled model will be named: OUTPUT_NAME.onnx.

check_model: checks model was compiled correctly.

check_output: validates compiled model outputs against original pytorch model.
NOTE: If you get the following error AssertionError: Not equal to tolerance rtol=0.001, atol=1e-05. Re-running the command should fix the issue.

Inference

Inference is performed by the inference.py script in the docker directory.
This script operates end-to-end, reading .nii.gz format CTs and writing .nii.gz format segmentations (no preprocessing necessary).

Extended datasets

We additionally trained our model with three datasets availbale here: https://github.com/JunMa11/AbdomenCT-1K

For each of these new datasets we retain identical model setup and training hyperparameters (this may not be ideal!)

AbdomenCT-1K

docker link: https://hub.docker.com/repository/docker/afgreen/rrr_mcr_abdomenct_1k

results (tested on 200 scans reserved from the full dataset):

AbdomenCT-1K with pseudo tumor labels

docker link: https://hub.docker.com/repository/docker/afgreen/rrr_mcr_abdomenct_1k_tumor

results (tested on the 285 scans without pseudo-tumor labels):

AbdomenCT-1K - 50 cases with 12 annotated organs

docker link: https://hub.docker.com/repository/docker/afgreen/rrr_mcr_abdomenct_1k_12organ

results (tested on the 959 scans without all 12 organs labelled):