From 273fc47cb3f8a9cbcc12f18ba9a3f7fc7c02bc3c Mon Sep 17 00:00:00 2001
From: Nathan Molinier <nathan.molinier@gmail.com>
Date: Mon, 11 Mar 2024 10:37:23 -0400
Subject: [PATCH] update init_data_config with contrast method

---
 scripts/init_data_config.py | 30 +++++++++++++++++++++---------
 scripts/utils.py            | 27 +++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 9 deletions(-)

diff --git a/scripts/init_data_config.py b/scripts/init_data_config.py
index 9126c21..66eac79 100644
--- a/scripts/init_data_config.py
+++ b/scripts/init_data_config.py
@@ -1,4 +1,6 @@
-# Based on https://github.com/spinalcordtoolbox/disc-labeling-hourglass
+"""
+Script copied from https://github.com/spinalcordtoolbox/disc-labeling-hourglass
+"""
 
 import os
 import argparse
@@ -7,10 +9,7 @@
 import itertools
 import numpy as np
 
-from utils import CONTRAST, get_img_path_from_mask_path, fetch_contrast
-
-
-CONTRAST_LOOKUP = {tuple(sorted(value)): key for key, value in CONTRAST.items()}
+from utils import get_img_path_from_mask_path, get_cont_path_from_other_cont, fetch_contrast, fetch_subject_and_session
 
 
 # Determine specified contrasts
@@ -26,10 +25,17 @@ def init_data_config(args):
     file_paths = [os.path.abspath(path.replace('\n', '')) for path in open(args.txt)]
     if args.type == 'LABEL':
         label_paths = file_paths
-        img_paths = [get_img_path_from_mask_path(lp) for lp in label_paths]
+        img_paths = [get_img_path_from_label_path(lp) for lp in label_paths]
         file_paths = label_paths + img_paths
     elif args.type == 'IMAGE':
         img_paths = file_paths
+    elif args.type == 'CONTRAST':
+        if not args.cont: # If the target contrast is not specified
+            raise ValueError(f'When using the type CONTRAST, please specify the target contrast using the flag "--cont"')
+        img_paths = file_paths
+        new_contrast = args.cont
+        label_paths = [get_cont_path_from_other_cont(ip) for ip in img_paths]
+        file_paths = label_paths + img_paths
     else:
         raise ValueError(f"invalid args.type: {args.type}")
     missing_paths = [
@@ -50,7 +56,7 @@ def init_data_config(args):
         raise ValueError('Please store all the BIDS datasets inside the same parent folder !')
 
     # Look up the right code for the set of contrasts present
-    contrasts = CONTRAST_LOOKUP[tuple(sorted(set(map(fetch_contrast, img_paths))))]
+    contrasts = "_".join(tuple(sorted(set(map(fetch_contrast, img_paths)))))
 
     config = {
         'TYPE': args.type,
@@ -58,6 +64,10 @@ def init_data_config(args):
         'DATASETS_PATH': dataset_parent_path
     }
 
+    # Add target contrast when the type CONTRAST is used
+    if args.type == 'CONTRAST':
+        config['TARGET_CONTRAST'] = args.cont
+
     # Split into training, validation, and testing sets
     split_ratio = (1 - (args.split_validation + args.split_test), args.split_validation, args.split_test) # TRAIN, VALIDATION, and TEST
     config_paths = label_paths if args.type == 'LABEL' else img_paths
@@ -91,8 +101,10 @@ def pairwise(iterable):
     ## Parameters
     parser.add_argument('--txt', required=True,
                         help='Path to TXT file that contains only image or label paths. (Required)')
-    parser.add_argument('--type', choices=('LABEL', 'IMAGE'),
-                        help='Type of paths specified. Choices "LABEL" or "IMAGE". (Required)')
+    parser.add_argument('--type', choices=('LABEL', 'IMAGE', 'CONTRAST'),
+                        help='Type of paths specified. Choices are "LABEL", "IMAGE" or "CONTRAST". (Required)')
+    parser.add_argument('--cont', type=str, default='',
+                        help='If the type CONTRAST is selected, this variable specifies the wanted contrast for target.')
     parser.add_argument('--split-validation', type=float, default=0.1,
                         help='Split ratio for validation. Default=0.1')
     parser.add_argument('--split-test', type=float, default=0.1,
diff --git a/scripts/utils.py b/scripts/utils.py
index 2367811..48d860a 100644
--- a/scripts/utils.py
+++ b/scripts/utils.py
@@ -88,6 +88,33 @@ def get_mask_path_from_img_path(img_path, deriv_sub_folders, short_suffix='_seg'
             mask_path.append(paths[0])
     return mask_path
 
+
+def get_cont_path_from_other_cont(str_path, cont):
+    """
+    :param str_path: absolute path to the input nifti img. Example: /<path_to_BIDS_data>/sub-amuALT/anat/sub-amuALT_T1w.nii.gz
+    :param cont: contrast of the target output image stored in the same data folder. Example: T2w
+    :return: path to the output target image. Example: /<path_to_BIDS_data>/sub-amuALT/anat/sub-amuALT_T2w.nii.gz
+
+    """
+    # Load path
+    path = Path(str_path)
+
+    # Extract file extension
+    ext = ''.join(path.suffixes)
+
+    # Remove input contrast from name
+    path_list = path.name.split('_')
+    suffixes_pos = [1 if len(part.split('-')) == 1 else 0 for part in path_list]
+    contrast_idx = suffixes_pos.index(1) # Find suffix
+
+    # New image name
+    img_name = '_'.join(path_list[:contrast_idx]+[cont]) + ext
+
+    # Recreate img path
+    img_path = os.path.join(str(path.parent), img_name)
+
+    return img_path
+
 def get_deriv_sub_from_img_path(img_path, derivatives_folder='derivatives'):
     """
     This function returns the derivatives path of the subject from an image path or an empty string if the path does not exists. Images need to be stored in a BIDS compliant dataset.