🏆 Champion Solution for FLARE24-Task3

This repository is the official implementation of A 3D Unsupervised Domain Adaption Framework Combining Style Translation and Self-Training for Abdominal Organs Segmentation of Team tju_vil_pioneers on FLARE24 challenge.

🔍 Overview

This work addresses the challenge of adapting CT-trained segmentation models to MR images by generating synthetic MR data, employing self-training strategies, and a two-stage segmentation framework in the FLARE24 dataset. For more details, see the pipeline diagram below:

⚙️ Environments and Requirements

• Ubuntu 20.04.6 LTS
• Intel(R) Xeon(R) Platinum 8153 CPU @ 2.00GHz, RAM 192GB , NVIDIA Tesla V100 (32G)
• CUDA >= 11.3
• python >= 3.7.13

To set up the environment, follow these steps:

conda create -n FLARE24 python=3.7.13
conda activate FLARE24
pip install -r requirements.txt

PS: You should install torch==1.12.0+cu113 before pip install requirements.

pip install torch==1.12.0+cu113 --extra-index-url https://download.pytorch.org/whl/cu113

⚠️ Important Notice

Before starting any training or inference, make sure to modify the base paths in all configuration files:

1. Modify the base paths in configs/xxx/xxx_base.yaml: xxx includes preprocess, train, inference, etc.

DATASET:
  BASE_DIR: "your/path/to/datasets"  # Dataset root directory

2. All training and inference configuration files inherit from base.yaml, so this step is mandatory.

3. Please ensure to use absolute paths instead of relative paths to avoid potential path resolution issues.

💾 Dataset

The training Data and validation data are provided by the FLARE24. In short, there are 2300 partial labeled CT and 4000+ unlabeled MR data for training, 110 public cases for validation and 300 hidden cases for the final test.

The CT scans are from the FLARE 2022 dataset where 50 cases have ground-truth labels and the remaining cases have pseudo labels (generated by the FLARE 2022 winning solution, ~90% DSC score).

1.CT scans need to be manually adjusted, placing the 2000 pseudo-labeled and 300 manually labeled data into the corresponding folders under the CT directory.

2.The two MR datasets 'AMOS_MR_good_spacing-833' and 'LLD-MMRI-3984' should both be placed under the MR/Training directory.

|-- datasets
|   |-- CT
|   |   |-- CT2MR_image
|   |   |-- CT_image
|   |   `-- CT_label
|   |-- MR
|   |   |-- PublicValidation
|   |   |   |-- imagesVal
|   |   |   `-- labelsVal
|   |   `-- Training
|   |       |-- AMOS_MR_good_spacing-833
|   |       |-- LLD-MMRI-3984
|   |       |-- PLabel
|   |       |-- PLabel_image
|   |       |-- preprocess_data_pair
|   |       `-- regis_data
|   |-- processed_data
|   |   |-- coarse
|   |   |   |-- combined_data
|   |   |   `-- small_segnet
|   |   `-- fine
|   |       |-- big_segnet
|   |       |-- combined_data
|   |       `-- small_segnet
|   `-- static_info

🪄 Preprocessing

Our preprocessing pipeline consists of three main steps:

Step 1: CT and Synthetic MR Data Preprocessing

This step preprocesses the original CT data and fake MR data generated through style translation:

Note1: Before preprocess, please refer to "Style_Translation/README.md" for style translation detailed information.

Note2: Before starting the preprocessing pipeline, please modify the data and file paths in preprocess/preprocess_base.yaml. You need to update:

• MR_DATA_PREPROCESS.ROOT_PATH: The absolute path to your datasets directory
• DATASET.BASE_DIR: The absolute path to your datasets directory
• FINE_MODEL_PATH and COARSE_MODEL_PATH: The paths to your downloaded checkpoint directories

# Process original CT data
python ./preprocess/data_preprocess.py --cfg ./configs/preprocess/preprocess_step1_CT.yaml

# Process style translation fake MR data
python ./preprocess/data_preprocess.py --cfg ./configs/preprocess/preprocess_step1_FakeMR.yaml

Step 2: MR Data Preprocessing and Pseudo-label Generation

This step handles two MR datasets (AMOS and LLD), including data filtering, registration, and pseudo-label generation:

# Process AMOS MR dataset
python ./preprocess/data_preprocess.py --cfg ./configs/preprocess/preprocess_step2_amos.yaml

# Process LLD MR dataset
python ./preprocess/data_preprocess.py --cfg ./configs/preprocess/preprocess_step2_lld.yaml

Step 3: Fine Segmentation Data Preprocessing

This step prepares data for the fine segmentation stage by adjusting patch size to [96,192,192]:

# Process AMOS MR dataset
python ./preprocess/data_preprocess.py --cfg ./configs/preprocess/preprocess_step3_amos_fine.yaml

# Process LLD MR dataset
python ./preprocess/data_preprocess.py --cfg ./configs/preprocess/preprocess_step3_lld_fine.yaml

🖥️ Train

Training our model consists of several main steps, as depicted in the pipeline figure above. Follow these instructions to conduct the training process effectively:

Step 1: Domain Translation

Before training the segmentation models, we first translate the source domain images to target domain using CycleGAN.

• Train stage one image-to-image translation model for domain translation

python stage_1_i2i_train.py --name sourceAtotargetB

• Generate target-like source domain images

python stage_1.5_i2i_inference.py --ckpt_path YOUR_PATH --source_npy_dirpath SOURCE_PATH --target_npy_dirpath TARGET_PATH --save_npy_dirpath SAVE_PATH --k_means_clusters 6

Step 2: Big SegNet Training

python train.py --cfg ./configs/train/train_big_segnet.yaml

Step 3: Small SegNet Training

# coarse stage
python train.py --cfg train_small_segnet_coarse_stage.yaml
# fine stage
python train.py --cfg train_small_segnet_fine_stage.yaml

You can download trained models here:

BaiduNetDisk

To fine-tune the model on a customized dataset, run this command: An example of small_segnet_fine_stage

python train.py --cfg train_small_segnet_fine_stage_fintune.yaml

🗳️ Inference

To infer the testing cases, run this command: An experiment of small segnet

python inference.py --cfg ./configs/inference/inference_small_segnet.yaml 

📊 Evaluation

To compute the evaluation metrics, run:

python eval.py --cfg ./configs/eval/eval_small_segnet.yaml

📋 Docker

To run the inference using Docker, use the following command:

Note: This is the official inference script. When running predictions, please replace /data/js/inputs/ and /data/js/outputs/ with your own input and output directories. The input MR images must be in .nii.gz format.

docker container run --gpus "device=0" \
    -m 28G \
    --name tju_vil_pioneers \
    --rm \
    -v /data/js/inputs/:/workspace/inputs/ \
    -v /data/js/outputs/:/workspace/outputs/ \
    tju_vil_pioneers:latest /bin/bash -c "sh predict.sh"

You can download the Docker image from:

Docker Container download link BaiduNetDisk

Extraction code: o9za

📋 Results

Our method achieves the following performance on FLARE24

Dataset Name	DSC(%)	NSD(%)
Validation Dataset	79.42%	86.46%
Test Dataset	(?)	(?)

🎆 Contributing

This project is licensed under the Apache License 2.0, and please adhere to the licenses of the models and datasets used.

🙏 Acknowledgement

We thank the contributors of FLARE24 datasets.