Add support for spinal rootlets segmentation

.github/workflows/ci.yml

@@ -35,6 +35,18 @@ jobs:
           python -m pip install --upgrade pip
           pip install -r requirements.txt
 
+      # https://github.com/spinalcordtoolbox/spinalcordtoolbox/blob/master/.ci.sh
+      - name: Install SCT
+        run: |
+          git clone https://github.com/spinalcordtoolbox/spinalcordtoolbox.git
+          cd spinalcordtoolbox
+          ./.ci.sh -i
+          # NB: install_sct edits ~/.bashrc, but those environment changes don't get passed to subsequent steps in GH Actions.
+          # So, we filter through the .bashrc and pass the values to $GITHUB_ENV and $GITHUB_PATH.
+          # Relevant documentation: https://docs.github.com/en/actions/reference/workflow-commands-for-github-actions#environment-files
+          cat ~/.bashrc | grep "export SCT_DIR" | cut -d " " -f 2 >> $GITHUB_ENV
+          cat ~/.bashrc | grep "export PATH" | grep -o "/.*" | cut -d ':' -f 1 >> $GITHUB_PATH
+
       - name: Run tests with pytest
         run: |
           python -m pytest -v tests/test_utils.py

README.md

@@ -9,6 +9,7 @@ Currently supported labels are:
 - disc labels
 - compression labels
 - ponto-medullary junction (PMJ) label
+- rootlets segmentation
 - centerline
 
 > **Note**

manual_correction.py

@@ -60,6 +60,7 @@ def get_parser():
         "'FILES_LABEL' lists images associated with vertebral labeling, "
         "'FILES_COMPRESSION' lists images associated with compression labeling, "
         "'FILES_PMJ' lists images associated with pontomedullary junction labeling, "
+        "'FILES_ROOTLETS' lists images associated with rootlets segmentation, "
         "and 'FILES_CENTERLINE' lists images associated with centerline. "
         "You can validate your YAML file at this website: http://www.yamllint.com/."
         "\nNote: if you want to iterate over all subjects, you can use the wildcard '*' (Examples: sub-*_T1w.nii.gz, "
@@ -85,6 +86,9 @@ def get_parser():
             FILES_PMJ:
             - sub-001_T1w.nii.gz
             - sub-002_T1w.nii.gz
+            FILES_ROOTLETS:
+            - sub-001_T1w.nii.gz
+            - sub-002_T1w.nii.gz
             FILES_CENTERLINE:
             - sub-001_T1w.nii.gz
             - sub-002_T1w.nii.gz\n
@@ -163,6 +167,11 @@ def get_parser():
         help="FILES-CENTERLINE suffix. Examples: '_centerline' (default), '_label-centerline'.",
         default='_centerline'
     )
+    parser.add_argument(
+        '-suffix-files-rootlets',
+        help="FILES-ROOTLETS suffix. Examples: '_label-rootlets_dseg' (default), '_rootlets'.",
+        default='_label-rootlets_dseg'
+    )
     parser.add_argument(
         '-label-disc-list',
         help="Comma-separated list containing individual values and/or intervals for disc labeling. Example: '1:4,6,8' "
@@ -249,6 +258,21 @@ def get_parser():
              }\n
              """),
     )
+    parser.add_argument(
+        '-change-orient',
+        type=str,
+        help=
+        "R|Orientation to show the image in the viewer. If provided, the image and label will be reoriented before "
+        "opening the viewer. After manual correction, the image and label will be reoriented back to the original "
+        "orientation.\n"
+        "Warning: be aware of this issue when using this flag: "
+        "https://github.com/spinalcordtoolbox/manual-correction/issues/101",
+        choices=['LAS', 'LAI', 'LPS', 'LPI', 'LSA', 'LSP', 'LIA', 'LIP', 'RAS', 'RAI', 'RPS', 'RPI', 'RSA', 'RSP',
+                 'RIA', 'RIP', 'ALS', 'ALI', 'ARS', 'ARI', 'ASL', 'ASR', 'AIL', 'AIR', 'PLS', 'PLI', 'PRS', 'PRI',
+                 'PSL', 'PSR', 'PIL', 'PIR', 'SLA', 'SLP', 'SRA', 'SRP', 'SAL', 'SAR', 'SPL', 'SPR', 'ILA', 'ILP',
+                 'IRA', 'IRP', 'IAL', 'IAR', 'IPL', 'IPR'],
+        default=''
+    )
     parser.add_argument(
         '-v', '--verbose',
         help="Full verbose (for debugging)",
@@ -649,6 +673,12 @@ def generate_qc(fname, fname_label, task, fname_qc, subject, config_file, qc_les
             archive_qc(fname_qc, config_file)
         else:
             print("WARNING: SC segmentation file not found: {}. QC report will not be generated.".format(fname_seg))
+
+    # Skip QC for the spinal rootlets segmentation as `sct_qc` does not support it
+    # Context: https://github.com/spinalcordtoolbox/spinalcordtoolbox/issues/4166#issuecomment-1654175610
+    elif task == 'FILES_ROOTLETS':
+        pass
+
     else:
         subprocess.check_call(['sct_qc',
                                '-i', fname,
@@ -721,6 +751,7 @@ def main():
         'FILES_LABEL': args.suffix_files_label,             # e.g., _labels or _label-disc
         'FILES_COMPRESSION': args.suffix_files_compression,  # e.g., _label-compression
         'FILES_PMJ': args.suffix_files_pmj,                 # e.g., _pmj or _label-pmj
+        'FILES_ROOTLETS': args.suffix_files_rootlets,       # e.g., _rootlets or _label-rootlets
         'FILES_CENTERLINE': args.suffix_files_centerline    # e.g., _centerline or _label-centerline
     }
     path_img = utils.get_full_path(args.path_img)
@@ -847,6 +878,16 @@ def main():
                     # For example: '/Users/user/dataset/derivatives/labels/sub-001/anat/sub-001_T2w_seg.nii.gz'
                     # The information regarding the modified data will be stored within the sidecar .json file
                     fname_out = utils.add_suffix(os.path.join(path_out, subject, ses, contrast, filename), suffix_dict[task])
+
+                    # Change orientation of the input image (if different from the original orientation)
+                    if args.change_orient:
+                        # Get image and label orientation
+                        image_orig_orient = utils.get_orientation(fname)
+                        label_orig_orient = utils.get_orientation(fname_label)
+                        # Change orientation of the input image for better visualization
+                        if image_orig_orient != args.change_orient or label_orig_orient != args.change_orient:
+                            utils.change_orientation(fname, args.change_orient)
+                            utils.change_orientation(fname_label, args.change_orient)
 
                     # Create subject folder in output if they do not exist
                     os.makedirs(os.path.join(path_out, subject, ses, contrast), exist_ok=True)
@@ -869,7 +910,7 @@ def main():
                             elif create_empty_mask:
                                 utils.create_empty_mask(fname, fname_out)
 
-                            if task in ['FILES_SEG', 'FILES_GMSEG']:
+                            if task in ['FILES_SEG', 'FILES_GMSEG', 'FILES_ROOTLETS']:
                                 if not args.add_seg_only:
                                     correct_segmentation(fname, fname_out, fname_other_contrast, args.viewer, param_fsleyes)
                             elif task == 'FILES_LESION':
@@ -909,6 +950,14 @@ def main():
                     # Keep track of corrected files in YAML.
                     dict_yml = utils.track_corrections(files_dict=dict_yml.copy(), config_path=args.config, file_path=fname, task=task)
 
+                    # Change orientation of the input image back to the original orientation
+                    if args.change_orient:
+                        image_current_orientation = utils.get_orientation(fname)
+                        label_current_orientation = utils.get_orientation(fname_label)
+                        if image_current_orientation != image_orig_orient or label_current_orientation != label_orig_orient:
+                            utils.change_orientation(fname, image_orig_orient)
+                            utils.change_orientation(fname_out, label_orig_orient)
+
             else:
                 sys.exit("ERROR: The list of files to correct is empty. \nMaybe, you have already corrected all the "
                          "files? Please, check the YAML file: {}".format(args.config))

tests/test_utils.py

@@ -7,8 +7,12 @@
 #######################################################################
 
 import os
+
+import numpy as np
+import nibabel as nib
+
 from utils import fetch_subject_and_session, add_suffix, remove_suffix, splitext, curate_dict_yml, get_full_path, \
-    check_files_exist, fetch_yaml_config, track_corrections
+    check_files_exist, fetch_yaml_config, track_corrections, get_orientation, change_orientation
 
 
 def test_fetch_subject_and_session():
@@ -198,3 +202,56 @@ def test_track_corrections(tmp_path):
 
     # Assert that the config file was updated correctly
     assert dict_files_updated == dict_files_test
+
+
+def create_dummy_nii_file(tmp_path, filename):
+    """
+    Create a dummy nifti file for testing purposes
+    :param tmp_path: Path to the temporary directory
+    :param filename: Name of the nifti file
+    :return: Path to the created nifti file
+    """
+    path_data = tmp_path / "BIDS" / "sub-001" / "ses-01" / "anat"
+    os.makedirs(path_data, exist_ok=True)
+
+    # Note: we have to create a 3D array to save it as a nifti file to simulate a real nifti file
+    data = np.random.rand(10, 10, 10)
+    affine = np.eye(4)
+    img = nib.Nifti1Image(data, affine)
+    fname_path = path_data / filename
+    nib.save(img, fname_path)
+
+    return fname_path
+
+
+def test_get_orientation(tmp_path):
+    """
+    Test that the get_orientation function returns the expected orientation
+    In this case, we create and check a nifti file with LPI orientation
+    """
+    # Create a test dummy nifti file using nibabel
+    fname_path = create_dummy_nii_file(tmp_path, "sub-001_ses-01_T1w.nii.gz")
+
+    # Get the orientation
+    orientation = get_orientation(fname_path)
+
+    # Assert that the orientation is correct
+    assert orientation == "LPI"
+
+
+def test_change_orientation(tmp_path):
+    """
+    Test that the change_orientation function changes the orientation of the nifti file
+    In this case, we reorient the file from LPI to AIL
+    """
+    # Create a test dummy nifti file using nibabel
+    fname_path = create_dummy_nii_file(tmp_path, "sub-001_ses-01_T1w.nii.gz")
+
+    # Change orientation to AIL
+    change_orientation(fname_path, "AIL")
+
+    # Get the orientation
+    orientation = get_orientation(fname_path)
+
+    # Assert that the orientation is correct
+    assert orientation == "AIL"

utils.py

@@ -332,3 +332,43 @@ def track_corrections(files_dict, config_path, file_path, task):
     return files_dict
 
 
+def get_orientation(file_path):
+    """
+    Get the orientation of the input nifti file
+    :param file_path: path to the nifti file
+    :return: actual orientation of the nifti file, e.g., 'RPI'
+    """
+
+    def _parse_orientation(output_bytes: bytes) -> str:
+        """
+        Parse the image orientation from the provided output bytes.
+        Args: output_bytes (bytes): The input bytes containing the output from sct_image command.
+        Returns:
+        str: The parsed image orientation.
+        """
+
+        output_string = output_bytes.decode('utf-8')
+        lines = output_string.strip().split('\n')
+        # Get only the string containing the orientation, e.g., 'RPI'
+        orientation = lines[-1].strip()
+        return orientation
+
+    # Note: we use bash command 'sct_image' to get the orientation instead of SCT's Image class because this would
+    # introduce dependency on SCT conda environment
+    output = subprocess.run(['sct_image', '-i', file_path, '-getorient'], capture_output=True)
+    orientation = _parse_orientation(output.stdout)  # e.g., 'RPI'
+
+    return orientation
+
+
+def change_orientation(file_path, orientation):
+    """
+    Change the orientation of the input nifti file
+    :param file_path: path to the nifti file
+    :param orientation: desired orientation of the nifti file, e.g., 'RPI'
+    """
+
+    # Change the orientation of the input nifti file
+    # Note: The image is currently being overwritten
+    subprocess.run(['sct_image', '-i', file_path, '-setorient', orientation, '-o', file_path], capture_output=True)
+    print(f"Orientation of {file_path} has been changed to {orientation}")