[*] fix remap_labels.py

Author: Kirscher

scripts/remap_labels.py

@@ -13,14 +13,128 @@
 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"
-]
+# Full TotalSegmentator mapping
+total_mapping = {
+    1: "spleen",
+    2: "kidney_right",
+    3: "kidney_left",
+    4: "gallbladder",
+    5: "liver",
+    6: "stomach",
+    7: "pancreas",
+    8: "adrenal_gland_right",
+    9: "adrenal_gland_left",
+    10: "lung_upper_lobe_left",
+    11: "lung_lower_lobe_left",
+    12: "lung_upper_lobe_right",
+    13: "lung_middle_lobe_right",
+    14: "lung_lower_lobe_right",
+    15: "esophagus",
+    16: "trachea",
+    17: "thyroid_gland",
+    18: "small_bowel",
+    19: "duodenum",
+    20: "colon",
+    21: "urinary_bladder",
+    22: "prostate",
+    23: "kidney_cyst_left",
+    24: "kidney_cyst_right",
+    25: "sacrum",
+    26: "vertebrae_S1",
+    27: "vertebrae_L5",
+    28: "vertebrae_L4",
+    29: "vertebrae_L3",
+    30: "vertebrae_L2",
+    31: "vertebrae_L1",
+    32: "vertebrae_T12",
+    33: "vertebrae_T11",
+    34: "vertebrae_T10",
+    35: "vertebrae_T9",
+    36: "vertebrae_T8",
+    37: "vertebrae_T7",
+    38: "vertebrae_T6",
+    39: "vertebrae_T5",
+    40: "vertebrae_T4",
+    41: "vertebrae_T3",
+    42: "vertebrae_T2",
+    43: "vertebrae_T1",
+    44: "vertebrae_C7",
+    45: "vertebrae_C6",
+    46: "vertebrae_C5",
+    47: "vertebrae_C4",
+    48: "vertebrae_C3",
+    49: "vertebrae_C2",
+    50: "vertebrae_C1",
+    51: "heart",
+    52: "aorta",
+    53: "pulmonary_vein",
+    54: "brachiocephalic_trunk",
+    55: "subclavian_artery_right",
+    56: "subclavian_artery_left",
+    57: "common_carotid_artery_right",
+    58: "common_carotid_artery_left",
+    59: "brachiocephalic_vein_left",
+    60: "brachiocephalic_vein_right",
+    61: "atrial_appendage_left",
+    62: "superior_vena_cava",
+    63: "inferior_vena_cava",
+    64: "portal_vein_and_splenic_vein",
+    65: "iliac_artery_left",
+    66: "iliac_artery_right",
+    67: "iliac_vena_left",
+    68: "iliac_vena_right",
+    69: "humerus_left",
+    70: "humerus_right",
+    71: "scapula_left",
+    72: "scapula_right",
+    73: "clavicula_left",
+    74: "clavicula_right",
+    75: "femur_left",
+    76: "femur_right",
+    77: "hip_left",
+    78: "hip_right",
+    79: "spinal_cord",
+    80: "gluteus_maximus_left",
+    81: "gluteus_maximus_right",
+    82: "gluteus_medius_left",
+    83: "gluteus_medius_right",
+    84: "gluteus_minimus_left",
+    85: "gluteus_minimus_right",
+    86: "autochthon_left",
+    87: "autochthon_right",
+    88: "iliopsoas_left",
+    89: "iliopsoas_right",
+    90: "brain",
+    91: "skull",
+    92: "rib_left_1",
+    93: "rib_left_2",
+    94: "rib_left_3",
+    95: "rib_left_4",
+    96: "rib_left_5",
+    97: "rib_left_6",
+    98: "rib_left_7",
+    99: "rib_left_8",
+    100: "rib_left_9",
+    101: "rib_left_10",
+    102: "rib_left_11",
+    103: "rib_left_12",
+    104: "rib_right_1",
+    105: "rib_right_2",
+    106: "rib_right_3",
+    107: "rib_right_4",
+    108: "rib_right_5",
+    109: "rib_right_6",
+    110: "rib_right_7",
+    111: "rib_right_8",
+    112: "rib_right_9",
+    113: "rib_right_10",
+    114: "rib_right_11",
+    115: "rib_right_12",
+    116: "sternum",
+    117: "costal_cartilages"
+}
 
-# Define desired mapping: keys are ROI names, values are target label integers
+# Define a desired remapping.
 desired_mapping = {
     "spleen": 1,
     "kidney_right": 2,
@@ -30,35 +144,23 @@
     "stomach": 6,
     "pancreas": 7,
     "esophagus": 8,
-    "small_bowel": 9,      # corresponds to "Small-Intestine"
+    "small_bowel": 9,  # corresponds to "small_bowel"
     "duodenum": 10,
-    "urinary_bladder": 11,  # corresponds to "Bladder"
+    "urinary_bladder": 11,
     "prostate": 12,
-    "spinal_cord": 13       # corresponds to "Spinal-Canal"
+    #"spinal_cord": 13
+    # Add additional mappings if needed.
 }
 
-# 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 for the remapped segmentation.
+new_data = np.copy(data)
 
-# 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
+# Iterate through each label in the TotalSegmentator mapping.
+for orig_label, structure in total_mapping.items():
+    # Determine the new label
+    new_label = desired_mapping.get(structure, 0)
+    new_data[data == orig_label] = new_label
 
-# Save the new segmentation with the same header and affine
+# 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)
+nib.save(new_img, output_file)