[+] add TS inference and remapping scripts

Author: Kirscher

scripts/remap_labels.py

@@ -0,0 +1,64 @@
+import sys
+import nibabel as nib
+import numpy as np
+
+if len(sys.argv) != 3:
+    print("Usage: remap_labels.py <input_segmentation> <output_segmentation>")
+    sys.exit(1)
+
+input_file = sys.argv[1]
+output_file = sys.argv[2]
+
+# Load the segmentation image
+img = nib.load(input_file)
+data = img.get_fdata().astype(np.int16)
+
+# Define the ROI list (in the order used for segmentation)
+roi_list = [
+    "spleen", "kidney_right", "kidney_left", "gallbladder", "liver", "stomach",
+    "pancreas", "esophagus", "small_bowel", "duodenum", "urinary_bladder", "prostate",
+    "spinal_cord"
+]
+
+# Define desired mapping: keys are ROI names, values are target label integers
+desired_mapping = {
+    "spleen": 1,
+    "kidney_right": 2,
+    "kidney_left": 3,
+    "gallbladder": 4,
+    "liver": 5,
+    "stomach": 6,
+    "pancreas": 7,
+    "esophagus": 8,
+    "small_bowel": 9,      # corresponds to "Small-Intestine"
+    "duodenum": 10,
+    "urinary_bladder": 11,  # corresponds to "Bladder"
+    "prostate": 12,
+    "spinal_cord": 13       # corresponds to "Spinal-Canal"
+}
+
+# Get the unique labels (ignoring background 0)
+unique_vals = np.unique(data)
+unique_vals = unique_vals[unique_vals != 0]
+unique_vals = np.sort(unique_vals)
+
+if len(unique_vals) != len(roi_list):
+    print("Warning: Expected {} ROI labels, but found {} in the segmentation.".format(len(roi_list), len(unique_vals)))
+
+# Build a remapping dictionary based on the assumed order.
+remap_dict = {}
+for i, orig_label in enumerate(unique_vals):
+    if i < len(roi_list):
+        roi = roi_list[i]
+        remap_dict[orig_label] = desired_mapping[roi]
+    else:
+        remap_dict[orig_label] = 0  # fallback to background if extra
+
+# Create a new data array with remapped labels
+new_data = np.zeros_like(data, dtype=np.int16)
+for orig, new in remap_dict.items():
+    new_data[data == orig] = new
+
+# Save the new segmentation with the same header and affine
+new_img = nib.Nifti1Image(new_data, img.affine, img.header)
+nib.save(new_img, output_file)

scripts/run_TotalSegmentator.sh

@@ -0,0 +1,96 @@
+#!/bin/bash
+# segment_folder.sh
+#
+# This script loops over all nii.gz files in an input folder,
+# runs TotalSegmentator with a restricted ROI subset, and then
+# remaps the resulting multi-label segmentation to a custom label mapping.
+#
+# The desired final mapping is:
+#   background:      0
+#   Spleen:          1
+#   Kidney-Right:    2
+#   Kidney-Left:     3
+#   Gall-Bladder:    4
+#   Liver:           5
+#   Stomach:         6
+#   Pancreas:        7
+#   Esophagus:       8
+#   Small-Intestine: 9
+#   Duodenum:        10
+#   Bladder:         11
+#   Prostate:        12
+#   Spinal-Canal:    13
+#
+# Note:
+#   The names used with --roi_subset must match the TotalSegmentator classes.
+#   Here we use:
+#       spleen kidney_right kidney_left gallbladder liver stomach pancreas esophagus small_bowel duodenum urinary_bladder prostate spinal_cord
+#   which correspond to the desired labels (with small naming differences).
+#
+# Usage:
+#   ./segment_folder.sh /path/to/input_folder /path/to/output_folder
+
+# Check for two arguments
+if [ "$#" -ne 2 ]; then
+    echo "Usage: $0 <input_folder> <output_folder>"
+    exit 1
+fi
+
+INPUT_FOLDER="$1"
+OUTPUT_FOLDER="$2"
+
+# Create output folder if it doesn't exist
+mkdir -p "$OUTPUT_FOLDER"
+
+# Create a temporary directory for TotalSegmentator outputs
+TMP_DIR=$(mktemp -d)
+
+# Function to process a single file
+process_file() {
+    local file="$1"
+    local filename=$(basename "$file")
+    local base="${filename%.nii.gz}"
+    echo "Processing $filename..."
+
+    # Create a temporary folder for this file's segmentation
+    local FILE_TMP_DIR="$TMP_DIR/$base"
+    mkdir -p "$FILE_TMP_DIR"
+
+    # Run TotalSegmentator using ROI subset and multi-label output
+    if ! TotalSegmentator -i "$file" -o "$FILE_TMP_DIR" \
+        --roi_subset spleen kidney_right kidney_left gallbladder liver stomach pancreas esophagus small_bowel duodenum urinary_bladder prostate spinal_cord \
+        --ml; then
+        echo "TotalSegmentator failed for $file; skipping."
+        return
+    fi
+
+    # Log the contents of the temporary directory
+    echo "Contents of $FILE_TMP_DIR:"
+    ls -l "$FILE_TMP_DIR"
+
+    # Find the segmentation file produced by TotalSegmentator.
+    local SEG_FILE=$(find "$TMP_DIR" -maxdepth 1 -name "$base.nii" | head -n 1)
+    if [ -z "$SEG_FILE" ]; then
+        echo "No segmentation file found for $file; skipping."
+        return
+    fi
+
+    # Run the Python script to remap labels
+    if ! python3 scripts/remap_labels.py "$SEG_FILE" "$OUTPUT_FOLDER/${base}.nii.gz"; then
+        echo "Label remapping failed for $file; skipping."
+        return
+    fi
+
+    echo "Saved remapped segmentation as ${base}.nii.gz in $OUTPUT_FOLDER"
+}
+
+export -f process_file
+export TMP_DIR
+export OUTPUT_FOLDER
+
+# Loop over all nii.gz files in the input folder and process them in parallel
+find "$INPUT_FOLDER" -name "*.nii.gz" -print0 | xargs -0 -n 1 -P 4 bash -c 'process_file "$@"' _
+
+# Clean up temporary files
+rm -rf "$TMP_DIR"
+echo "Processing complete."