Repository: spinalcordtoolbox/spinalcordtoolbox
Issue category: Feature Request

Issue:
Unify FSLeyes commands for multiclass DeepSeg models
FSLeyes commands outputted by `sct_deepseg -i t2.nii.gz -task seg_sc_lesion_t2w_sci` differ from the FSLeyes command in [our tutorial](https://github.com/spinalcordtoolbox/sct_tutorial_data/blob/8b91e72c6575cbf40c5a66898fe3a13d9c26790c/single_subject/batch_single_subject.sh#L128):

<img width="1406" alt="Screenshot 2024-11-12 at 11 16 44 AM" src="https://github.com/user-attachments/assets/5204f40c-6d93-4147-adc1-252e99ff292b">

We could consider unifying the commands. I personally like the command from our tutorial, which opens both spinal cord and lesion segmentation in a single FSLeyes window:

```console
fsleyes t2.nii.gz -cm greyscale t2_sc_seg.nii.gz -cm red -a 70.0 t2_lesion_seg.nii.gz -cm blue-lightblue -a 70.0 &
```



Hints:

Patch:
diff --git a/documentation/source/user_section/command-line/deepseg/totalspineseg.rst b/documentation/source/user_section/command-line/deepseg/totalspineseg.rst
new file mode 100644
index 0000000000..2c5423a6d7
--- /dev/null
+++ b/documentation/source/user_section/command-line/deepseg/totalspineseg.rst
@@ -0,0 +1,11 @@
+
+                
+totalspineseg
+=============
+                
+.. argparse::
+   :ref: spinalcordtoolbox.scripts.sct_deepseg.get_parser
+   :prog: sct_deepseg -h
+   :markdownhelp:
+   :noepilog:
+                
\ No newline at end of file
diff --git a/documentation/source/user_section/command-line/sct_deepseg.rst b/documentation/source/user_section/command-line/sct_deepseg.rst
index ef3a7dde11..f73532b201 100644
--- a/documentation/source/user_section/command-line/sct_deepseg.rst
+++ b/documentation/source/user_section/command-line/sct_deepseg.rst
@@ -118,19 +118,24 @@ Other structures
 .. |seg_spinal_rootlets_t2w| image:: https://raw.githubusercontent.com/spinalcordtoolbox/doc-figures/master/sct_deepseg/spinal_rootlets_t2.png
    :target: deepseg/seg_spinal_rootlets_t2w.html
 
-Rootlets segmentation can be performed by running the following sample command:
+.. |totalspineseg| image:: https://raw.githubusercontent.com/spinalcordtoolbox/doc-figures/master/sct_deepseg/totalspineseg.png
+   :target: deepseg/totalspineseg.html
+
+Multiple structures may be segmented by running the following sample command:
 
 .. code::
 
    sct_deepseg -task seg_spinal_rootlets_t2w -i input.nii.gz
 
+You can replace "``seg_spinal_rootlets_t2w``" with any of the task names in the table below to perform different tasks. Click on a task below for more information.
+
 
 .. list-table::
    :align: center
    :widths: 25 25 25
 
    * - |seg_spinal_rootlets_t2w| ``seg_spinal_rootlets_t2w``
-     - 
+     - |totalspineseg| ``totalspineseg``
      -
 
 .. |canal_t2w| image:: https://raw.githubusercontent.com/spinalcordtoolbox/doc-figures/master/sct_deepseg/spinal_canal_t2.png
@@ -200,3 +205,4 @@ If you absolutely require these models, you can downgrade to version of SCT list
    deepseg/seg_tumor_t2
    deepseg/seg_spinal_rootlets_t2w
    deepseg/canal_t2w
+   deepseg/totalspineseg
diff --git a/requirements.txt b/requirements.txt
index 2d22937966..4b36236e24 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -59,6 +59,7 @@ requirements-parser
 scipy
 scikit-image
 scikit-learn
+totalspineseg
 xlwt
 tqdm
 transforms3d
diff --git a/spinalcordtoolbox/deepseg/inference.py b/spinalcordtoolbox/deepseg/inference.py
index e299883e70..c8478aa095 100644
--- a/spinalcordtoolbox/deepseg/inference.py
+++ b/spinalcordtoolbox/deepseg/inference.py
@@ -10,8 +10,11 @@
 import shutil
 import time
 from pathlib import Path
+import glob
+import multiprocessing as mp
 
 from ivadomed import inference as imed_inference
+from totalspineseg.inference import inference as tss_inference
 import nibabel as nib
 import numpy as np
 import torch
@@ -107,7 +110,11 @@ def segment_and_average_volumes(model_paths, input_filenames, options, use_gpu=F
 def segment_non_ivadomed(path_model, model_type, input_filenames, threshold, keep_largest, fill_holes_in_pred,
                          remove_small, use_gpu=False, remove_temp_files=True):
     # MONAI and NNUnet have similar structure, and so we use nnunet+inference functions with the same signature
-    if model_type == "monai":
+    # NB: For TotalSpineSeg, we don't need to create the network ourselves
+    if "totalspineseg" in path_model:
+        def create_net(pm, _): return pm  # just echo `path_model` back as 'net'
+        inference = segment_totalspineseg
+    elif model_type == "monai":
         create_net = ds_monai.create_nnunet_from_plans
         inference = segment_monai
     else:
@@ -323,3 +330,49 @@ def segment_nnunet(path_img, tmpdir, predictor, device: torch.device):
         fnames_out.append(fname_out)
 
     return fnames_out, targets
+
+
+def segment_totalspineseg(path_img, tmpdir, predictor, device):
+    # for totalspineseg, the 'predictor' is just the model path
+    path_model = predictor
+    # fetch the release subdirectory from the model path
+    installed_releases = sorted(
+        os.path.basename(release_path) for release_path in
+        glob.glob(os.path.join(path_model, 'nnUNet', 'results', 'r*'), recursive=True)
+    )
+    # There should always be a release subdirectory, hence the 'assert'
+    assert installed_releases, f"No 'nnUNet/results/rYYYYMMDD' subdirectory found in {path_model}"
+
+    # Copy the file to the temporary directory using shutil.copyfile
+    path_img_tmp = os.path.join(tmpdir, os.path.basename(path_img))
+    shutil.copyfile(path_img, path_img_tmp)
+    logger.info(f'Copied {path_img} to {path_img_tmp}')
+
+    # Create directory for nnUNet prediction
+    tmpdir_nnunet = os.path.join(tmpdir, 'nnUNet_prediction')
+    os.mkdir(tmpdir_nnunet)
+
+    # totalspineseg uses multiprocessing internally, so here we try to
+    # set multiprocessing method to 'spawn' to avoid deadlock issues
+    # - 'spawn' is already default on Windows/macOS, but 'fork' is default on Linux
+    # - 'spawn' is reliable though slower, but 'fork' causes intermittent stalling
+    mp.set_start_method("spawn", force=True)
+
+    tss_inference(
+        input_path=path_img_tmp,
+        output_path=tmpdir_nnunet,
+        data_path=path_model,
+        # totalspineseg requires explicitly specifying the release subdirectory
+        default_release=installed_releases[-1],  # use the most recent release
+        # totalspineseg expects the device type, not torch.device
+        device=device.type,
+        # Try to address stalling due to the use of concurrent.futures in totalspineseg
+        max_workers=1,
+        max_workers_nnunet=1,
+    )
+    fnames_out, targets = [], []
+    for output_dirname in ["step1_canal", "step1_cord", "step1_levels", "step1_output", "step2_output"]:
+        fnames_out.append(os.path.join(tmpdir_nnunet, output_dirname, os.path.basename(path_img)))
+        targets.append(f"_{output_dirname}")
+
+    return fnames_out, targets
diff --git a/spinalcordtoolbox/deepseg/models.py b/spinalcordtoolbox/deepseg/models.py
index baf7ea78d6..08dbee3656 100644
--- a/spinalcordtoolbox/deepseg/models.py
+++ b/spinalcordtoolbox/deepseg/models.py
@@ -14,10 +14,12 @@
 import shutil
 import glob
 from pathlib import Path
+from importlib.metadata import metadata
 
 from spinalcordtoolbox import download
-from spinalcordtoolbox.utils.sys import stylize, __deepseg_dir__
+from spinalcordtoolbox.utils.sys import stylize, __deepseg_dir__, LazyLoader
 
+tss_init = LazyLoader("tss_init", globals(), 'totalspineseg.init_inference')
 
 logger = logging.getLogger(__name__)
 
@@ -184,6 +186,17 @@
         "thr": None,  # Images are already binarized
         "default": False,
     },
+    "totalspineseg": {
+         # NB: Rather than hardcoding the URLs ourselves, use the URLs from the totalspineseg package.
+         # This means that when the totalspineseg package is updated, the URLs will be too, thus triggering
+         # a re-installation of the model URLs
+         "url": dict([meta.split(', ') for meta in metadata('totalspineseg').get_all('Project-URL')
+                      if meta.startswith('Dataset')]),
+         "description": "Instance segmentation of vertebrae, intervertebral discs (IVDs), spinal cord, and spinal canal on multi-contrasts MRI scans.",
+         "contrasts": ["any"],
+         "thr": None,  # Images are already binarized
+         "default": False,
+     },
 }
 
 
@@ -363,6 +376,13 @@
                              'false positives segmentations of other anatomical structures.',
          'url': 'https://github.com/ivadomed/model-canal-seg',
          'models': ['model_seg_canal_t2w']},
+    'totalspineseg':
+        {'description': 'Intervertebral discs labeling and vertebrae segmentation',
+         'long_description': 'TotalSpineSeg is a tool for automatic instance segmentation of all vertebrae, intervertebral discs (IVDs), '
+                             'spinal cord, and spinal canal in MRI images. It is robust to various MRI contrasts, acquisition orientations, '
+                             'and resolutions. The model used in TotalSpineSeg is based on nnU-Net as the backbone for training and inference.',
+         'url': 'https://github.com/neuropoly/totalspineseg',
+         'models': ['totalspineseg']},
 }
 
 
@@ -408,11 +428,17 @@ def install_model(name_model, custom_url=None):
     else:
         if not isinstance(url_field, dict):
             raise ValueError("Invalid url field in MODELS")
-        urls_used = {}
-        for seed_name, model_urls in url_field.items():
-            logger.info(f"\nInstalling '{seed_name}'...")
-            urls_used[seed_name] = download.install_data(model_urls,
-                                                         folder(os.path.join(name_model, seed_name)), keep=True)
+        # totalspineseg handles data downloading itself, so just pass the urls along
+        if name_model == 'totalspineseg':
+            tss_init.init_inference(data_path=Path(folder(name_model)), quiet=False, dict_urls=url_field,
+                                    store_export=False)  # Avoid having duplicate .zip files stored on disk
+            urls_used = url_field
+        else:
+            urls_used = {}
+            for seed_name, model_urls in url_field.items():
+                logger.info(f"\nInstalling '{seed_name}'...")
+                urls_used[seed_name] = download.install_data(model_urls,
+                                                             folder(os.path.join(name_model, seed_name)), keep=True)
     # Write `source.json` (for model provenance / updating)
     source_dict = {
         'model_name': name_model,
diff --git a/spinalcordtoolbox/reports/qc2.py b/spinalcordtoolbox/reports/qc2.py
index fd9b8cac2e..1c5ed23b5e 100644
--- a/spinalcordtoolbox/reports/qc2.py
+++ b/spinalcordtoolbox/reports/qc2.py
@@ -272,10 +272,17 @@ def sct_deepseg(
         dataset=dataset,
         subject=subject,
     ) as imgs_to_generate:
+        # Custom QC to handle multiclass segmentation outside the spinal cord
         if "seg_spinal_rootlets_t2w" in argv:
             sct_deepseg_spinal_rootlets_t2w(
-                imgs_to_generate, fname_input, fname_seg, fname_seg2, species)
-        if plane == 'Axial':
+                imgs_to_generate, fname_input, fname_seg, fname_seg2, species,
+                radius=(23, 23))
+        elif "totalspineseg" in argv:
+            sct_deepseg_spinal_rootlets_t2w(
+                imgs_to_generate, fname_input, fname_seg, fname_seg2, species,
+                radius=(40, 40), outline=False)
+        # Non-rootlets, axial/sagittal DeepSeg QC report
+        elif plane == 'Axial':
             sct_deepseg_axial(
                 imgs_to_generate, fname_input, fname_seg, fname_seg2, species)
         else:
@@ -385,6 +392,8 @@ def sct_deepseg_spinal_rootlets_t2w(
     fname_seg_sc: str,
     fname_seg_lesion: Optional[str],
     species: str,
+    radius: Sequence[int],
+    outline: bool = True
 ):
     """
     Generate a QC report for `sct_deepseg -task seg_spinal_rootlets_t2w`.
@@ -400,8 +409,6 @@ def sct_deepseg_spinal_rootlets_t2w(
     img_seg_sc = Image(fname_seg_sc).change_orientation('SAL')
     img_seg_lesion = Image(fname_seg_lesion).change_orientation('SAL') if fname_seg_lesion else None
 
-    # Rootlets need a larger "base" radius as they exist outside the SC
-    radius = (23, 23)
     # The radius size is suited to the species-specific resolutions. But, since we plan to skip
     # resampling, we need to instead adjust the crop radius to suit the *actual* resolution.
     p_original = img_seg_sc.dim[5]  # dim[0:3] => shape, dim[4:7] => pixdim, so dim[5] == pixdim[1]
@@ -459,8 +466,9 @@ def sct_deepseg_spinal_rootlets_t2w(
                   alpha=1.0,
                   interpolation='none',
                   aspect=1.0)
-        # linewidth 0.5 is too thick, 0.25 is too thin
-        plot_outlines(img, ax=ax, facecolor='none', edgecolor='black', linewidth=0.3)
+        if outline:
+            # linewidth 0.5 is too thick, 0.25 is too thin
+            plot_outlines(img, ax=ax, facecolor='none', edgecolor='black', linewidth=0.3)
         add_segmentation_labels(ax, img, colors=colormaps[i].colors, radius=tuple(r*scale for r in radius))
 
     ax.get_xaxis().set_visible(False)
@@ -798,8 +806,9 @@ def add_segmentation_labels(ax: mpl_axes.Axes, seg_mosaic: np.ndarray, colors: l
     # Fetch mosaic shape properties
     bbox = [2*radius[0], 2*radius[1]]
     grid_shape = [s // bb for s, bb in zip(seg_mosaic.shape, bbox)]
-    # Fetch set of labels in the mosaic
-    labels = [v for v in np.unique(seg_mosaic) if v]
+    # Fetch set of labels in the mosaic (including labels in between min/max)
+    labels = [float(val) for val in range(int(np.unique(seg_mosaic).min()),
+                                          int(np.unique(seg_mosaic).max())+1)]
     # Iterate over each sub-array in the mosaic
     for row in range(grid_shape[0]):
         for col in range(grid_shape[1]):
@@ -810,8 +819,9 @@ def add_segmentation_labels(ax: mpl_axes.Axes, seg_mosaic: np.ndarray, colors: l
             # Check for nonzero labels, then draw text for each label found
             labels_in_arr = [v for v in np.unique(arr) if v]
             for idx_pos, l_arr in enumerate(labels_in_arr, start=1):
+                lr_shift = -4 * (len(str(int(l_arr))) - 1)
                 y, x = (extents[0].stop - 6*idx_pos + 3,  # Shift each subsequent label up in case there are >1
-                        extents[1].stop - 6)
+                        extents[1].stop - 6 + lr_shift)   # Shift labels left if double/triple digit
                 color = colors[0] if len(colors) == 1 else colors[labels.index(l_arr)]
                 ax.text(x, y, str(int(l_arr)), color=color, size=4).set_path_effects([
                     mpl_patheffects.Stroke(linewidth=1, foreground='black'),
@@ -916,7 +926,8 @@ def assign_label_colors_by_groups(labels):
     should each be assigned their own distinct colormap, as to group them semantically.
     """
     # Arrange colormaps for max contrast between colormaps, and max contrast between colors in colormaps
-    distinct_colormaps = ['Blues', 'Reds', 'Greens', 'Oranges', 'Purples']
+    # Put reds first because SC labels tend to be ~1, so they will usually be first in a list of labels
+    distinct_colormaps = ['Reds', 'Blues', 'Greens', 'Oranges', 'Purples']
     colormap_sampling = [0.25, 0.5, 0.75, 0.5]  # light -> medium -> dark -> medium -> (repeat)
 
     # Split labels into subgroups --> we split the groups wherever the difference between labels is > 1
diff --git a/spinalcordtoolbox/scripts/sct_deepseg.py b/spinalcordtoolbox/scripts/sct_deepseg.py
index a958ba1ee3..5a66a19d3a 100755
--- a/spinalcordtoolbox/scripts/sct_deepseg.py
+++ b/spinalcordtoolbox/scripts/sct_deepseg.py
@@ -356,9 +356,14 @@ def main(argv: Sequence[str]):
 
     if arguments.qc is not None:
         # Models can have multiple input images -- create 1 QC report per input image.
-        # Models can also either have 1 OR 2 outputs per input, so we may need to split output_filenames into 2 lists
         if len(output_filenames) == len(input_filenames):
             iterator = zip(input_filenames, output_filenames, [None] * len(input_filenames))
+        # Special case: totalspineseg which outputs 5 files per 1 input file
+        # Just use the 5th image ([4]) which represents the step2 output
+        elif arguments.task[0] == 'totalspineseg':
+            assert len(output_filenames) == 5 * len(input_filenames)
+            iterator = zip(input_filenames, output_filenames[4::5], [None] * len(input_filenames))
+        # Other models typically have 2 outputs per input (e.g. SC + lesion), so use both segs
         else:
             assert len(output_filenames) == 2 * len(input_filenames)
             iterator = zip(input_filenames, output_filenames[0::2], output_filenames[1::2])
@@ -378,11 +383,14 @@ def main(argv: Sequence[str]):
                 plane=arguments.qc_plane,
             )
 
+    images = [arguments.i[0]]
+    im_types = ['anat']
+    opacities = ['']
     for output_filename in output_filenames:
-        img_kind = check_image_kind(Image(output_filename))
-        display_viewer_syntax([arguments.i[0], output_filename],
-                              im_types=['anat', img_kind],
-                              opacities=['', '0.7'], verbose=verbose)
+        images.append(output_filename)
+        im_types.append(check_image_kind(Image(output_filename)))
+        opacities.append('0.7')
+    display_viewer_syntax(images, im_types=im_types, opacities=opacities, verbose=verbose)
 
 
 if __name__ == "__main__":


Test patch:
diff --git a/testing/cli/test_sct_deepseg.py b/testing/cli/test_sct_deepseg.py
index b6f4e9e8fb..a012f959f9 100644
--- a/testing/cli/test_sct_deepseg.py
+++ b/testing/cli/test_sct_deepseg.py
@@ -158,6 +158,12 @@ def test_segment_nifti_softseg_error_with_fill_holes(tmp_path):
      ["_GM_seg", "_WM_seg"],
      'seg_mouse_gm_wm_t1w',
      0.5),
+    (sct_test_path('t2', 't2.nii.gz'),
+     [None, None, None, None, None],
+     't2_deepseg.nii.gz',
+     ["_step1_canal", "_step1_cord", "_step1_levels", "_step1_output", "_step2_output"],
+     'totalspineseg',
+     0),
 ])
 def test_segment_nifti_multiclass(fname_image, fnames_seg_manual, fname_out, suffixes, task, thr,
                                   tmp_path):