Competition Page: https://www.drivendata.org/competitions/99/biomass-estimation/page/534/
Team: Just4Fun
Contact: [email protected]
Source Code: https://github.com/quqixun/BioMassters
- S1 and S2 features and AGBM labels were carefully preprocessed according to statistics of data. See code in process.py and ./libs/process for details.
- Training data was splited into 5 folds for cross validation in split.py.
- Processed S1 and S2 features were concatenated to 3D tensor in shape [B, 15, 12, 256, 256] as input, targets were AGBM labels in shape [B, 1, 256, 256].
- We applied Swin UNETR with the attention from Swin Transformer V2 as the regression model. In ./libs/models, Swin UNETR was adapted from the implementation by MONAI project.
- In training steps, Swin UNETR was optimized by weighted MAE and SSIM. RMSE of validation data was used to select the best model.
- We trained Swin UNETR using each fold, and got 5 models.
- For each testing sample, the average of 5 predictions was the final result.
- Ubuntu 20.04 LTS
- CUDA 11.3 or later
- Any GPU with at least 40Gb VRAM for training
- Any GPU with at least 8Gb VRAM for predicting
conda create --name biomassters python=3.9
conda activate biomassters
pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
pip install -r requirements.txt
- Download metadata from DATA DOWNLOAD page, and put all files in ./data/information as following structure:
./data/information
├── biomassters-download-instructions.txt # Instructions to download satellite images and AGBM data
├── features_metadata.csv # Metadata for satellite images
└── train_agbm_metadata.csv # Metadata for training set AGBM tifs
- Download image data by running
./scripts/download.sh
, data is saved in ./data/source:
./data/source
├── test
│ ├── aa5e092e
│ │ ├── S1
│ │ │ ├── aa5e092e_S1_00.tif
│ │ │ ├── ...
│ │ │ └── aa5e092e_S1_11.tif
│ │ └── S2
│ │ ├── aa5e092e_S2_00.tif
│ │ ├── ...
│ │ └── aa5e092e_S2_11.tif
| ├── ...
│ └── fff812c0
└── train
├── aa018d7b
├── ...
└── fff05995
- Calculate statistics for normalization and split dataset into 5 folds by running
./scripts/process.sh
:
./data/source
├── plot # plot of data distribution
├── splits.pkl # 5 folds for cross validation
├── stats_log2.pkl # statistics of log2 transformed dataset
├── stats_plain.pkl # statistics of original dataset
├── test
└── train
You don't have to run the above script again since all outputs can be found in ./data/source.
Train model with arguments (see ./script/train.sh):
data_root
: root directory of training datasetexp_root
: root directory to save checkpoints, logs and modelsconfig_file
: file path of configurationsprocess_method
: processing method to calculate statistics,log2
orplain
, default isplain
folds
: list of folds, separated by,
device=0
process=plain
folds=0,1,2,3,4
data_root=./data/source
config_file=./configs/swin_unetr/exp1.yaml
CUDA_VISIBLE_DEVICES=$device \
python train.py \
--data_root $data_root \
--exp_root ./experiments/$process \
--config_file $config_file \
--process_method $process \
--folds $folds
Models and logs will be saved in ./experiments/plain/swin_unetr/exp1.
Training on 5 folds will take about 1 week if only one GPU is available.
If you have 5 GPUs, you can run each fold training on each GPU, and it will take less than 2 days.
You can download the trained model from Google Drive (tbd) or Baidu Disc (code:4dbj).
Make predictions with almost the same arguments as training (see ./script/predict.sh):
data_root
: root directory of training datasetexp_root
: root directory to save checkpoints, logs and modelsoutput_root
: root directory to save predictionsconfig_file
: file path of configurationsprocess_method
: processing method to calculate statistics,log2
orplain
, default isplain
folds
: list of folds, separated by,
apply_tta
: if apply test-time augmentation, default isFalse
device=0
process=plain
folds=0,1,2,3,4
apply_tta=false
data_root=./data/source
config_file=./configs/swin_unetr/exp1.yaml
CUDA_VISIBLE_DEVICES=$device \
python predict.py \
--data_root $data_root \
--exp_root ./experiments/$process \
--output_root ./predictions/$process \
--config_file $config_file \
--process_method $process \
--folds $folds \
--apply_tta $apply_tta
Predictions will be saved in ./predictions/plain/swin_unetr/exp1/folds_0-1-2-3-4.
Predicting testing samples on 5 folds and calculating the average will take about 30 minutes.
You can download the submission from Google Drive (tbd) or Baidu Disc (code:0me5).
Metrics of submitted models and predictions on validation dataset and testing dataset.
Metrics | Val Fold 0 | Val Fold 1 | Val Fold 2 | Val Fold 3 | Val Fold 4 | Val Average | Test |
---|---|---|---|---|---|---|---|
LMAE | 0.03562 | 0.03516 | 0.03527 | 0.03522 | 0.03626 | - | - |
LSSIM | 0.04758 | 0.04684 | 0.04713 | 0.04691 | 0.04834 | - | - |
RMSE | 27.9676 | 27.4368 | 27.5011 | 27.8954 | 28.0946 | 27.7781 | 27.3891 |
- Swin UNETR: Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images. [Paper]
- Swin Transformer V2: Scaling Up Capacity and Resolution. [Paper , Code]
- Implementation of Swin UNETR by MONAI project. [Code]
- Differentiable structure similarity metric. [Code]
- MIT License