Port towards config file-based training

.gitignore

@@ -1 +1,2 @@
-*.idea
+*.idea
+*__pycache__/

README.md

@@ -1,8 +1,23 @@
 # Towards Contrast-agnostic Soft Segmentation of the Spinal Cord
+
+[![arXiv](https://img.shields.io/badge/arXiv-2310.15402-b31b1b.svg)](https://arxiv.org/abs/2310.15402)
+
 Official repository for contrast-agnostic spinal cord segmentation project using SoftSeg. 
 
 This repo contains all the code for data preprocessing, training and running inference on other datasets. The code is mainly based on [Spinal Cord Toolbox](https://spinalcordtoolbox.com) and [MONAI](https://github.com/Project-MONAI/MONAI) (PyTorch).
 
+**CITATION INFO**: If you find this work and/or code useful for your research, please cite our paper:
+
+```
+@article{bedard2023towards,
+  title={Towards contrast-agnostic soft segmentation of the spinal cord},
+  author={B{\'e}dard, Sandrine and Enamundram, Naga Karthik and Tsagkas, Charidimos and Pravat{\`a}, Emanuele and Granziera, Cristina and Smith, Andrew and Weber II, Kenneth Arnold and Cohen-Adad, Julien},
+  journal={arXiv preprint arXiv:2310.15402},
+  year={2023}
+  url={https://arxiv.org/abs/2310.15402}
+}
+```
+
 ## Table of contents
 * [1. Main Dependencies](#1-main-dependencies)
 * [2. Dataset](#2-dataset)
@@ -149,6 +164,8 @@ The training script expects a datalist file in the Medical Decathlon format cont
 python monai/create_msd_data.py -pd ~/duke/projects/ivadomed/contrast-agnostic-seg/data_processed_sg_2023-08-08_NO_CROP\data_processed_clean> -po ~/datasets/contrast-agnostic/ --contrast all --label-type soft --seed 42
 ```
 
+The dataset split containing the training, validation, and test subjects can be found in the `monai/data_split_all_soft_seed15.yaml` file. 
+
 > **Note** 
 > The output of the above command is just `.json` file pointing to the image-label pairs in the original BIDS dataset. It _does not_ copy the existing data to the output folder. 
 

configs/train_all.yaml

@@ -0,0 +1,63 @@
+seed: 15
+save_test_preds: True
+
+directories:
+  # Path to the saved models directory
+  models_dir: /home/GRAMES.POLYMTL.CA/u114716/contrast-agnostic/saved_models/followup
+  # Path to the saved results directory
+  results_dir: /home/GRAMES.POLYMTL.CA/u114716/contrast-agnostic/results/models_followup
+  # Path to the saved wandb logs directory
+  # if None, starts training from scratch. Otherwise, resumes training from the specified wandb run folder
+  wandb_run_folder: None  
+
+dataset:
+  # Dataset name (will be used as "group_name" for wandb logging)
+  name: spine-generic
+  # Path to the dataset directory containing all datalists (.json files)
+  root_dir: /home/GRAMES.POLYMTL.CA/u114716/contrast-agnostic/datalists/spine-generic/seed15
+  # Type of contrast to be used for training. "all" corresponds to training on all contrasts
+  contrast: all   # choices: ["t1w", "t2w", "t2star", "mton", "mtoff", "dwi", "all"]
+  # Type of label to be used for training. 
+  label_type: soft_bin  # choices: ["hard", "soft", "soft_bin"]
+
+preprocessing:
+  # Online resampling of images to the specified spacing.
+  spacing: [1.0, 1.0, 1.0]
+  # Center crop/pad images to the specified size. (NOTE: done after resampling)
+  # values correspond to R-L, A-P, I-S axes of the image after 1mm isotropic resampling.
+  crop_pad_size: [64, 192, 320]   
+
+opt:
+  name: adam
+  lr: 0.001
+  max_epochs: 200
+  batch_size: 2
+  # Interval between validation checks in epochs
+  check_val_every_n_epochs: 10
+  # Early stopping patience (this is until patience * check_val_every_n_epochs)
+  early_stopping_patience: 20     
+
+
+model:
+  # Model architecture to be used for training (also to be specified as args in the command line)
+  nnunet:
+    # NOTE: these info are typically taken from nnUNetPlans.json (if an nnUNet model is trained)
+    base_num_features: 32
+    max_num_features: 320
+    n_conv_per_stage_encoder: [2, 2, 2, 2, 2, 2]
+    n_conv_per_stage_decoder: [2, 2, 2, 2, 2]
+    pool_op_kernel_sizes: [
+      [1, 1, 1],
+      [2, 2, 2],
+      [2, 2, 2],
+      [2, 2, 2],
+      [2, 2, 2],
+      [1, 2, 2] 
+    ]
+    enable_deep_supervision: True
+
+  unetr:
+    feature_size: 16
+    hidden_size: 768    # dimensionality of hidden embeddings
+    mlp_dim: 2048       # dimensionality of the MLPs
+    num_heads: 12       # number of heads in multi-head Attention