Merge pull request #125 from sct-pipeline/nk/contrast-agnostic-v2

Update model with new datasets, pathologies and contrasts
sct-pipeline · Nov 13, 2024 · 12ab4e0 · 12ab4e0
2 parents ad7ee3d + f9e5833
commit 12ab4e0
Show file tree

Hide file tree

Showing 24 changed files with 2,542 additions and 709 deletions.
diff --git a/configs/train_all.yaml b/configs/train_all.yaml
@@ -1,5 +1,5 @@
 seed: 50
-save_test_preds: True
+save_test_preds: False
 
 directories:
   # Path to the saved models directory
@@ -12,14 +12,9 @@ directories:
 
 dataset:
   # Dataset name (will be used as "group_name" for wandb logging)
-  name: lifelong-contrast-agnostic
+  name: contrast-agnostic-v3
   # Path to the dataset directory containing all datalists (.json files)
-  # root_dir: /home/GRAMES.POLYMTL.CA/u114716/contrast-agnostic/datalists/lifelong-contrast-agnostic
-  root_dir: /home/GRAMES.POLYMTL.CA/u114716/contrast-agnostic/datalists/aggregation-20240517
-  # 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"]
+  root_dir: /home/GRAMES.POLYMTL.CA/u114716/contrast-agnostic/datalists/v2-final-aggregation-20241017
 
 preprocessing:
   # Online resampling of images to the specified spacing.
@@ -34,7 +29,7 @@ opt:
   max_epochs: 250
   batch_size: 2
   # Interval between validation checks in epochs
-  check_val_every_n_epochs: 1
+  check_val_every_n_epochs: 50
   # Early stopping patience (this is until patience * check_val_every_n_epochs)
   early_stopping_patience: 20
 
@@ -72,19 +67,3 @@ model:
       [1, 2, 2],
     ]
     enable_deep_supervision: True
-
-  mednext:
-    num_input_channels: 1
-    base_num_features: 32
-    num_classes: 1
-    kernel_size: 3            # 3x3x3 and 5x5x5 were tested in publication
-    block_counts: [2,2,2,2,1,1,1,1,1] # number of blocks in each layer
-    norm_type: 'layer'
-    enable_deep_supervision: True
-
-  swinunetr:
-    spatial_dims: 3
-    depths: [2, 2, 2, 2]
-    num_heads: [3, 6, 12, 24]   # number of heads in multi-head Attention
-    feature_size: 36
-    use_pretrained: False
diff --git a/csa_generate_figures/analyse_csa_across.py b/csa_generate_figures/analyse_csa_across.py
@@ -13,16 +13,16 @@
 import matplotlib.pyplot as plt
 
 # Setting the hue order as specified
-# HUE_ORDER = ["softseg_bin", "deepseg_2d", "nnunet", "monai", "mednext", "swinunetr", "swinpretrained", "ensemble"]
+HUE_ORDER = ["softseg_bin", "deepseg", "plain_320", "plain_384", "resencM"]
 # HUE_ORDER = ["softseg_bin", "deepseg_2d", "monai_single", "monai_7datasets", "swinunetr_7datasets"]
-HUE_ORDER = ["softseg_bin", "monai_v21", "monai_v23", "monai_v2x"]
+# HUE_ORDER = ["softseg_bin", "monai_v21", "monai_v23", "monai_v2x"]
 HUE_ORDER_THR = ["GT", "15", "1", "05", "01", "005"]
 HUE_ORDER_RES = ["1mm", "05mm", "15mm", "3mm", "2mm"]
 CONTRAST_ORDER = ["DWI", "MTon", "MToff", "T1w", "T2star", "T2w"]
 
 FONTSIZE = 12
-# XTICKS = ["GT", "DeepSeg2D", "C-A\nv2.1", "C-A\nv2.3", "C-A\nv2.x"] 
-XTICKS = ["GT", "contrast-agnostic\nv2.1", "contrast-agnostic\nv2.3", "contrast-agnostic\nv2.x"] 
+XTICKS = ["GT", "DeepSeg2D", "C-A\nplain_320", "C-A\nplain_384", "C-A\nresencM"]
+# XTICKS = ["GT", "contrast-agnostic\nv2.1", "contrast-agnostic\nv2.3", "contrast-agnostic\nv2.x"] 
 
 
 def save_figure(file_path, save_fname):
@@ -61,8 +61,8 @@ def extract_contrast_and_details(filename, across="Method"):
     if across == "Method":
         # pattern = r'.*_(DWI|MTon|MToff|T1w|T2star|T2w).*_(softseg_bin|deepseg_2d|soft_input|bin_input).*'
         # pattern = r'.*_(DWI|MTon|MToff|T1w|T2star|T2w).*_(softseg_bin|deepseg_2d|nnunet|monai|mednext|swinunetr|swinpretrained|ensemble).*'
-        # pattern = r'.*_(DWI|MTon|MToff|T1w|T2star|T2w).*_(softseg_bin|deepseg_2d|monai_single|monai_7datasets|swinunetr_7datasets).*'
-        pattern = r'.*_(DWI|MTon|MToff|T1w|T2star|T2w).*_(softseg_bin|monai_v21|monai_v23|monai_v2x).*'
+        pattern = r'.*_(DWI|MTon|MToff|T1w|T2star|T2w).*_(softseg_bin|deepseg|plain_320|plain_384|resencM).*'
+        # pattern = r'.*_(DWI|MTon|MToff|T1w|T2star|T2w).*_(softseg_bin|monai_v21|monai_v23|monai_v2x).*'
         match = re.search(pattern, filename)
         if match:
             return match.group(1), match.group(2)

diff --git a/csa_generate_figures/analyse_perslice_csa_across.py b/csa_generate_figures/analyse_perslice_csa_across.py
@@ -13,15 +13,15 @@
 import matplotlib.pyplot as plt
 
 # Setting the hue order as specified
-# HUE_ORDER = ["softseg_bin", "deepseg_2d", "nnunet", "monai", "mednext", "swinunetr", "swinpretrained", "ensemble"]
+HUE_ORDER = ["softseg_bin", "deepseg", "plain_320", "plain_384", "resencM"]
 # HUE_ORDER = ["softseg_bin", "deepseg_2d", "monai_single", "monai_7datasets", "swinunetr_7datasets"]
-HUE_ORDER = ["softseg_bin", "monai_v21", "monai_v23", "monai_v2x"]
+# HUE_ORDER = ["softseg_bin", "monai_v21", "monai_v23", "monai_v2x"]
 HUE_ORDER_RES = ["1mm", "05mm", "15mm", "3mm", "2mm"]
 CONTRAST_ORDER = ["DWI", "MTon", "MToff", "T1w", "T2star", "T2w"]
 
 FONTSIZE = 12
-# XTICKS = ["GT", "DeepSeg2D", "C-A\nOriginal", "C-A\n7-datasets", "SwinUNETR\n7-datasets"]
-XTICKS = ["GT", "contrast-agnostic\nv2.1", "contrast-agnostic\nv2.3", "contrast-agnostic\nv2.x"] 
+XTICKS = ["GT", "DeepSeg2D", "C-A\nplain_320", "C-A\nplain_384", "C-A\nresencM"]
+# XTICKS = ["GT", "contrast-agnostic\nv2.1", "contrast-agnostic\nv2.3", "contrast-agnostic\nv2.x"] 
 
 
 def save_figure(file_path, save_fname):
@@ -59,9 +59,9 @@ def extract_contrast_and_details(filename, across="Method"):
     # pattern = r'.*iso-(\d+mm).*_(propseg|deepseg_2d|nnunet_3d_fullres|monai).*'
     if across == "Method":
         # pattern = r'.*_(DWI|MTon|MToff|T1w|T2star|T2w).*_(softseg_bin|deepseg_2d|soft_input|bin_input).*'
-        # pattern = r'.*_(DWI|MTon|MToff|T1w|T2star|T2w).*_(softseg_bin|deepseg_2d|nnunet|monai|mednext|swinunetr|swinpretrained|ensemble).*'
+        pattern = r'.*_(DWI|MTon|MToff|T1w|T2star|T2w).*_(softseg_bin|deepseg|plain_320|plain_384|resencM).*'
         # pattern = r'.*_(DWI|MTon|MToff|T1w|T2star|T2w).*_(softseg_bin|deepseg_2d|monai_single|monai_7datasets|swinunetr_7datasets).*'
-        pattern = r'.*_(DWI|MTon|MToff|T1w|T2star|T2w).*_(softseg_bin|monai_v21|monai_v23|monai_v2x).*'
+        # pattern = r'.*_(DWI|MTon|MToff|T1w|T2star|T2w).*_(softseg_bin|monai_v21|monai_v23|monai_v2x).*'
         match = re.search(pattern, filename)
         if match:
             return match.group(1), match.group(2)

diff --git a/csa_qc_evaluation_spine_generic/comparison_across_models.sh b/csa_qc_evaluation_spine_generic/comparison_across_models.sh
@@ -77,7 +77,8 @@ label_vertebrae(){
   FILELABEL="${file}_discs"
 
   # Label vertebral levels
-  sct_label_utils -i ${file}.nii.gz -disc ${FILELABEL}.nii.gz -o ${FILESEG}_labeled.nii.gz
+  # sct_label_utils -i ${file}.nii.gz -disc ${FILELABEL}.nii.gz -o ${FILESEG}_labeled.nii.gz
+  sct_label_utils -i ${FILESEG}.nii.gz -disc ${FILELABEL}.nii.gz -o ${FILESEG}_labeled.nii.gz
 
   # # Run QC
   # sct_qc -i ${file}.nii.gz -s ${file_seg}_labeled.nii.gz -p sct_label_vertebrae -qc ${PATH_QC} -qc-subject ${SUBJECT}
@@ -245,22 +246,28 @@ segment_sc_MONAI(){
   local contrast="$4"   # used only for saving output file name
   local csv_fname="$5"   # used for saving output file name
 
-	if [[ $model == 'monai_single' ]]; then
+	if [[ $model == 'plain_320' ]]; then
 		# FILESEG="${file%%_*}_${contrast}_seg_monai_orig"
     FILESEG="${file%%_*}_${contrast}_seg_${model}"
 		PATH_MODEL=${PATH_MONAI_MODEL_1}
     model_name='monai'
+    max_feat=320
+    pad='edge'
 
-	elif [[ $model == 'monai_v23' ]]; then
+	elif [[ $model == 'plain_384' ]]; then
 		# FILESEG="${file%%_*}_${contrast}_seg_monai_ll"
     FILESEG="${file%%_*}_${contrast}_seg_${model}"
 		PATH_MODEL=${PATH_MONAI_MODEL_2}
     model_name='monai'
+    max_feat=384
+    pad='edge'
 
-  elif [[ $model == 'monai_v2x' ]]; then
+  elif [[ $model == 'resencM' ]]; then
     FILESEG="${file%%_*}_${contrast}_seg_${model}"
     PATH_MODEL=${PATH_MONAI_MODEL_3}
-    model_name='monai'
+    model_name='monai-resencM'
+    max_feat=384
+    pad='edge'
 
 	elif [[ $model == 'swinunetr' ]]; then
     FILESEG="${file%%_*}_${contrast}_seg_swinunetr"
@@ -275,7 +282,7 @@ segment_sc_MONAI(){
   echo "Running inference from model at ${PATH_MODEL}"
   # NOTE: surprisingly, the `edge` padding is resulting in higher abs. csa error compared to `constant` (zero) padded inputs.
   # Run SC segmentation
-  python ${PATH_MONAI_SCRIPT} --path-img ${file}.nii.gz --path-out . --chkp-path ${PATH_MODEL} --device gpu --model ${model_name} --pred-type soft --pad-mode constant
+  python ${PATH_MONAI_SCRIPT} --path-img ${file}.nii.gz --path-out . --chkp-path ${PATH_MODEL} --device gpu --model ${model_name} --pred-type soft --pad-mode ${pad} --max-feat ${max_feat}
   # python ${PATH_MONAI_SCRIPT} --path-img ${file}.nii.gz --path-out . --chkp-path ${PATH_MODEL} --device gpu --model monai --pred-type soft --pad-mode constant
   # Rename MONAI output
   mv ${file}_pred.nii.gz ${FILESEG}.nii.gz
@@ -293,8 +300,8 @@ segment_sc_MONAI(){
   echo "${FILESEG},${slicewise_dice}"
   echo "${FILESEG},${slicewise_dice}" >> ${PATH_RESULTS}/slicewise_dice.csv
 
-  # # Generate QC report with soft prediction
-  # sct_qc -i ${file}.nii.gz -s ${FILESEG}.nii.gz -p sct_deepseg_sc -qc ${PATH_QC} -qc-subject ${SUBJECT}
+  # Generate QC report 
+  sct_qc -i ${file}.nii.gz -s ${FILESEG}.nii.gz -p sct_deepseg_sc -qc ${PATH_QC} -qc-subject ${SUBJECT}
 
   # Compute CSA averaged across all slices C2-C3 vertebral levels for plotting the STD across contrasts
   # NOTE: this is per-level because not all contrasts have thes same FoV (C2-C3 is what all contrasts have in common)
@@ -364,7 +371,7 @@ contrasts="space-other_T1w space-other_T2w space-other_T2star flip-1_mt-on_space
 # contrasts="space-other_T1w rec-average_dwi"
 
 # output csv filename
-csv_fname="csa_softIn_CL_deploy"   # "csa_label_inputs"
+csv_fname="csa_model_sizes"   # "csa_label_inputs"
 
 # Loop across contrasts
 for contrast in ${contrasts}; do
@@ -435,12 +442,12 @@ for contrast in ${contrasts}; do
   sct_process_segmentation -i ${FILEBIN}.nii.gz -perslice 1 -vertfile ${file}_seg-manual_labeled.nii.gz -o $PATH_RESULTS/${csv_fname}_softseg_bin_perslice.csv -append 1
 
   # 3. Segment SC using different methods, binarize at 0.5 and compute CSA
-	CUDA_VISIBLE_DEVICES=1 segment_sc_MONAI ${file} "${file}_seg-manual" 'monai_v21' ${contrast} ${csv_fname}
-  CUDA_VISIBLE_DEVICES=1 segment_sc_MONAI ${file} "${file}_seg-manual" 'monai_v23' ${contrast} ${csv_fname}
-  CUDA_VISIBLE_DEVICES=2 segment_sc_MONAI ${file} "${file}_seg-manual" 'monai_v2x' ${contrast} ${csv_fname}
+	CUDA_VISIBLE_DEVICES=1 segment_sc_MONAI ${file} "${file}_seg-manual" 'plain_320' ${contrast} ${csv_fname}
+  CUDA_VISIBLE_DEVICES=2 segment_sc_MONAI ${file} "${file}_seg-manual" 'plain_384' ${contrast} ${csv_fname}
+  CUDA_VISIBLE_DEVICES=3 segment_sc_MONAI ${file} "${file}_seg-manual" 'resencM' ${contrast} ${csv_fname}
   # CUDA_VISIBLE_DEVICES=2 segment_sc_MONAI ${file} "${file}_seg-manual" 'swinunetr' ${contrast} ${csv_fname}
   # CUDA_VISIBLE_DEVICES=3 segment_sc_nnUNet ${file} "${file}_seg-manual" '3d_fullres' ${contrast} ${csv_fname}
-  # segment_sc ${file} "${file}_seg-manual" 'deepseg' ${deepseg_input_c} ${contrast} ${csv_fname}
+  segment_sc ${file} "${file}_seg-manual" 'deepseg' ${deepseg_input_c} ${contrast} ${csv_fname}
   # TODO: run on deep/progseg after fixing the contrasts for those
   # segment_sc ${file_res} 't2' 'propseg' '' "${file}_seg-manual" ${native_res}
 

diff --git a/datasplits/datasplit_dcm-brno_seed50.yaml b/datasplits/datasplit_dcm-brno_seed50.yaml
@@ -0,0 +1,156 @@
+test:
+- sub-1860B6472B
+- sub-2295B4676B
+- sub-2336B
+- sub-2347B6454B
+- sub-2353B
+- sub-2450B6177B
+- sub-2694B
+- sub-2779B4786B
+- sub-2804B4632B
+- sub-3066B
+- sub-3121B
+- sub-3148B
+- sub-3206B
+- sub-3261B4658B
+- sub-3758B6378B
+- sub-4403B5238B
+train:
+- sub-1838B6577B
+- sub-2060B4733B
+- sub-2247B6492B
+- sub-2249B4725B
+- sub-2255B6372B
+- sub-2271B
+- sub-2284B4723B
+- sub-2287B
+- sub-2296B4806B
+- sub-2319B4634B
+- sub-2320B6468B
+- sub-2321B6243B
+- sub-2322B
+- sub-2330B6137B
+- sub-2333B6135B
+- sub-2334B4627B
+- sub-2340B6562B
+- sub-2346B
+- sub-2348B4595B
+- sub-2349B6345B
+- sub-2358B4598B
+- sub-2371B4687B
+- sub-2383B6196B
+- sub-2384B
+- sub-2386B
+- sub-2389B
+- sub-2390B4949B
+- sub-2391B6466B
+- sub-2407B5757B
+- sub-2411B4591B
+- sub-2412B
+- sub-2413B4629B
+- sub-2414B
+- sub-2416B6027B
+- sub-2417B4965B
+- sub-2419B
+- sub-2420B
+- sub-2421B
+- sub-2446B6192B
+- sub-2451B
+- sub-2472B
+- sub-2479B6195B
+- sub-2481B6079B
+- sub-2599B4623B
+- sub-2600B4616B
+- sub-2606B6390B
+- sub-2608B
+- sub-2610B6476B
+- sub-2614B
+- sub-2644B5869B
+- sub-2648B
+- sub-2649B
+- sub-2654B4633B
+- sub-2667B6188B
+- sub-2683B
+- sub-2698B
+- sub-2699B
+- sub-2702B6403B
+- sub-2717B
+- sub-2737B
+- sub-2741B4963B
+- sub-2742B
+- sub-2750B4646B
+- sub-2751B
+- sub-2752B
+- sub-2756B
+- sub-2757B
+- sub-2773B4661B
+- sub-2774B4763B
+- sub-2778B6482B
+- sub-2825B4881B
+- sub-2847B6453B
+- sub-2873B
+- sub-2884B4714B
+- sub-2886B5253B
+- sub-2887B4699B
+- sub-2902B4626B
+- sub-2988B6186B
+- sub-2999B
+- sub-3005B
+- sub-3006B
+- sub-3007B
+- sub-3013B6405B
+- sub-3024B
+- sub-3056B6483B
+- sub-3060B6384B
+- sub-3075B6206B
+- sub-3077B
+- sub-3094B
+- sub-3111B
+- sub-3120B6558B
+- sub-3132B4599B
+- sub-3140B
+- sub-3146B
+- sub-3167B
+- sub-3170B
+- sub-3194B
+- sub-3207B4664B
+- sub-3209B
+- sub-3214B
+- sub-3215B
+- sub-3230B
+- sub-3231B
+- sub-3249B
+- sub-3262B
+- sub-3281B4590B
+- sub-3289B4615B
+- sub-3427B4966B
+- sub-3608B6555B
+- sub-3636B
+- sub-3640B
+- sub-3646B
+- sub-3651B
+- sub-3782B6193B
+- sub-3793B4681B
+- sub-3998B6406B
+- sub-4088B6310B
+- sub-4196B5202B
+- sub-4793B6025B
+- sub-4891B6572B
+- sub-4960B6571B
+val:
+- sub-1836B6029B
+- sub-2259B4883B
+- sub-2315B4686B
+- sub-2316B5038B
+- sub-2355B
+- sub-2372B2863B
+- sub-2418B4628B
+- sub-2590B4647B
+- sub-2613B
+- sub-2723B4648B
+- sub-3179B
+- sub-3227B6487B
+- sub-3232B5049B
+- sub-3237B
+- sub-3641B
+- sub-3759B6530B