Change labels

README.md

@@ -1,6 +1,6 @@
 # TotalSpineSeg
 
-TotalSpineSeg is a tool for automatic instance segmentation and labeling of all vertebrae, intervertebral discs (IVDs), spinal cord, and spinal canal in MRI images. It follows the [TotalSegmentator classes](https://github.com/wasserth/TotalSegmentator/tree/v1.5.7#class-details) with additional classes for IVDs, spinal cord, and spinal canal (see [list of classes](#list-of-classes)). The model is based on [nnUNet](https://github.com/MIC-DKFZ/nnUNet) as the backbone for training and inference.
+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 [nnUNet](https://github.com/MIC-DKFZ/nnUNet) as the backbone for training and inference.
 
 If you use this model, please cite our work:
 > Warszawer Y, Molinier N, Valošek J, Shirbint E, Benveniste PL, Achiron A, Eshaghi A and Cohen-Adad J. _Fully Automatic Vertebrae and Spinal Cord Segmentation Using a Hybrid Approach Combining nnU-Net and Iterative Algorithm_.	Proceedings of the 32th Annual Meeting of ISMRM. 2024
@@ -222,53 +222,53 @@ For a more detailed view of the output examples, you can check the [PDF version]
 
 | Label | Name |
 |:------|:-----|
-| 18 | vertebrae_L5 |
-| 19 | vertebrae_L4 |
-| 20 | vertebrae_L3 |
-| 21 | vertebrae_L2 |
-| 22 | vertebrae_L1 |
-| 23 | vertebrae_T12 |
-| 24 | vertebrae_T11 |
-| 25 | vertebrae_T10 |
-| 26 | vertebrae_T9 |
+| 1 | spinal_cord |
+| 2 | spinal_canal |
+| 10 | vertebrae_C1 |
+| 11 | vertebrae_C2 |
+| 12 | vertebrae_C3 |
+| 13 | vertebrae_C4 |
+| 14 | vertebrae_C5 |
+| 15 | vertebrae_C6 |
+| 16 | vertebrae_C7 |
+| 20 | vertebrae_T1 |
+| 21 | vertebrae_T2 |
+| 22 | vertebrae_T3 |
+| 23 | vertebrae_T4 |
+| 24 | vertebrae_T5 |
+| 25 | vertebrae_T6 |
+| 26 | vertebrae_T7 |
 | 27 | vertebrae_T8 |
-| 28 | vertebrae_T7 |
-| 29 | vertebrae_T6 |
-| 30 | vertebrae_T5 |
-| 31 | vertebrae_T4 |
-| 32 | vertebrae_T3 |
-| 33 | vertebrae_T2 |
-| 34 | vertebrae_T1 |
-| 35 | vertebrae_C7 |
-| 36 | vertebrae_C6 |
-| 37 | vertebrae_C5 |
-| 38 | vertebrae_C4 |
-| 39 | vertebrae_C3 |
-| 40 | vertebrae_C2 |
-| 41 | vertebrae_C1 |
-| 92 | sacrum |
-| 200 | spinal_cord |
-| 201 | spinal_canal |
-| 202 | disc_L5_S |
-| 203 | disc_L4_L5 |
-| 204 | disc_L3_L4 |
-| 205 | disc_L2_L3 |
-| 206 | disc_L1_L2 |
-| 207 | disc_T12_L1 |
-| 208 | disc_T11_T12 |
-| 209 | disc_T10_T11 |
-| 210 | disc_T9_T10 |
-| 211 | disc_T8_T9 |
-| 212 | disc_T7_T8 |
-| 213 | disc_T6_T7 |
-| 214 | disc_T5_T6 |
-| 215 | disc_T4_T5 |
-| 216 | disc_T3_T4 |
-| 217 | disc_T2_T3 |
-| 218 | disc_T1_T2 |
-| 219 | disc_C7_T1 |
-| 220 | disc_C6_C7 |
-| 221 | disc_C5_C6 |
-| 222 | disc_C4_C5 |
-| 223 | disc_C3_C4 |
-| 224 | disc_C2_C3 |
+| 28 | vertebrae_T9 |
+| 29 | vertebrae_T10 |
+| 30 | vertebrae_T11 |
+| 31 | vertebrae_T12 |
+| 40 | vertebrae_L1 |
+| 41 | vertebrae_L2 |
+| 42 | vertebrae_L3 |
+| 43 | vertebrae_L4 |
+| 44 | vertebrae_L5 |
+| 50 | sacrum |
+| 60 | disc_C2_C3 |
+| 61 | disc_C3_C4 |
+| 62 | disc_C4_C5 |
+| 63 | disc_C5_C6 |
+| 64 | disc_C6_C7 |
+| 70 | disc_C7_T1 |
+| 71 | disc_T1_T2 |
+| 72 | disc_T2_T3 |
+| 73 | disc_T3_T4 |
+| 74 | disc_T4_T5 |
+| 75 | disc_T5_T6 |
+| 76 | disc_T6_T7 |
+| 77 | disc_T7_T8 |
+| 78 | disc_T8_T9 |
+| 79 | disc_T9_T10 |
+| 80 | disc_T10_T11 |
+| 81 | disc_T11_T12 |
+| 90 | disc_T12_L1 |
+| 91 | disc_L1_L2 |
+| 92 | disc_L2_L3 |
+| 93 | disc_L3_L4 |
+| 94 | disc_L4_L5 |
+| 100 | disc_L5_S |

pyproject.toml

@@ -75,6 +75,7 @@ totalspineseg_average4d = "totalspineseg.utils.average4d:main"
 totalspineseg_crop_image2seg = "totalspineseg.utils.crop_image2seg:main"
 totalspineseg_extract_soft = "totalspineseg.utils.extract_soft:main"
 totalspineseg_extract_levels = "totalspineseg.utils.extract_levels:main"
+totalspineseg_extract_alternate = "totalspineseg.utils.extract_alternate:main"
 totalspineseg_add_nnunet_trainer = "totalspineseg.utils.add_nnunet_trainer:main"
 
 [build-system]

totalspineseg/__init__.py

@@ -3,6 +3,7 @@
 from .utils.cpdir import cpdir_mp
 from .utils.crop_image2seg import crop_image2seg, crop_image2seg_mp
 from .utils.extract_levels import extract_levels, extract_levels_mp
+from .utils.extract_alternate import extract_alternate, extract_alternate_mp
 from .utils.extract_soft import extract_soft, extract_soft_mp
 from .utils.fill_canal import fill_canal, fill_canal_mp
 from .utils.iterative_label import iterative_label, iterative_label_mp

totalspineseg/inference.py

@@ -233,6 +233,27 @@ def main():
             max_workers=max_workers,
             quiet=quiet,
         )
+        preview_jpg_mp(
+            output_path / 'input',
+            output_path / 'preview',
+            segs_path=output_path / 'localizers',
+            output_suffix='_loc_tags',
+            override=True,
+            max_workers=max_workers,
+            quiet=quiet,
+            label_texts_right={
+                10: 'C1', 11: 'C2', 12: 'C3', 13: 'C4', 14: 'C5', 15: 'C6', 16: 'C7',
+                20: 'T1', 21: 'T2', 22: 'T3', 23: 'T4', 24: 'T5', 25: 'T6', 26: 'T7',
+                27: 'T8', 28: 'T9', 29: 'T10', 30: 'T11', 31: 'T12',
+                40: 'L1', 41: 'L2', 42: 'L3', 43: 'L4', 44: 'L5'
+            },
+            label_texts_left={
+                50: 'Sacrum', 60: 'C2C3', 61: 'C3C4', 62: 'C4C5', 63: 'C5C6', 64: 'C6C7', 70: 'C7T1',
+                71: 'T1T2', 72: 'T2T3', 73: 'T3T4', 74: 'T4T5', 75: 'T5T6', 76: 'T6T7', 77: 'T7T8',
+                78: 'T8T9', 79: 'T9T10', 80: 'T10T11', 81: 'T11T12', 90: 'T12L1',
+                91: 'L1L2', 92: 'L2L3', 93: 'L3L4', 94: 'L4L5', 100: 'L5S'
+            },
+        )
 
     if not quiet: print('\n' 'Converting 4D images to 3D:')
     average4d_mp(
@@ -323,10 +344,11 @@ def main():
         iterative_label_mp(
             output_path / 'step1_output',
             output_path / 'step1_output',
+            selected_disc_landmarks=[2, 5, 3, 4],
             disc_labels=[1, 2, 3, 4, 5],
-            init_disc={2:224, 5:202, 3:219, 4:207},
-            output_disc_step=-1,
-            map_input_dict={6:92, 7:201, 8:201, 9:200},
+            disc_landmark_labels=[2, 3, 4, 5],
+            disc_landmark_output_labels=[60, 70, 90, 100],
+            map_input_dict={6:50, 7:2, 8:2, 9:1},
             override=True,
             max_workers=max_workers,
             quiet=quiet,
@@ -336,11 +358,12 @@ def main():
             output_path / 'step1_output',
             output_path / 'step1_output',
             locs_path=output_path / 'localizers',
+            selected_disc_landmarks=[2, 5],
             disc_labels=[1, 2, 3, 4, 5],
-            init_disc={2:224, 5:202},
-            output_disc_step=-1,
-            loc_disc_labels=list(range(202, 225)),
-            map_input_dict={6:92, 7:201, 8:201, 9:200},
+            disc_landmark_labels=[2, 3, 4, 5],
+            disc_landmark_output_labels=[60, 70, 90, 100],
+            loc_disc_labels=list(range(60, 101)),
+            map_input_dict={6:50, 7:2, 8:2, 9:1},
             override=True,
             max_workers=max_workers,
             quiet=quiet,
@@ -351,8 +374,8 @@ def main():
     fill_canal_mp(
         output_path / 'step1_output',
         output_path / 'step1_output',
-        canal_label=201,
-        cord_label=200,
+        canal_label=2,
+        cord_label=1,
         largest_canal=True,
         largest_cord=True,
         override=True,
@@ -380,14 +403,29 @@ def main():
         max_workers=max_workers,
         quiet=quiet,
     )
+    preview_jpg_mp(
+        output_path / 'input',
+        output_path / 'preview',
+        segs_path=output_path / 'step1_output',
+        output_suffix='_step1_output_tags',
+        override=True,
+        max_workers=max_workers,
+        quiet=quiet,
+        label_texts_left={
+            60: 'C2C3', 61: 'C3C4', 62: 'C4C5', 63: 'C5C6', 64: 'C6C7', 70: 'C7T1',
+            71: 'T1T2', 72: 'T2T3', 73: 'T3T4', 74: 'T4T5', 75: 'T5T6', 76: 'T6T7', 77: 'T7T8',
+            78: 'T8T9', 79: 'T9T10', 80: 'T10T11', 81: 'T11T12', 90: 'T12L1',
+            91: 'L1L2', 92: 'L2L3', 93: 'L3L4', 94: 'L4L5', 100: 'L5S'
+        },
+    )
 
     if not quiet: print('\n' 'Extracting spinal cord soft segmentation from step 1 model output:')
     extract_soft_mp(
         output_path / 'step1_raw',
         output_path / 'step1_output',
         output_path / 'step1_cord',
         label=9,
-        seg_labels=[200],
+        seg_labels=[1],
         dilate=1,
         override=True,
         max_workers=max_workers,
@@ -400,7 +438,7 @@ def main():
         output_path / 'step1_output',
         output_path / 'step1_canal',
         label=7,
-        seg_labels=[200, 201],
+        seg_labels=[1, 2],
         dilate=1,
         override=True,
         max_workers=max_workers,
@@ -417,9 +455,8 @@ def main():
     extract_levels_mp(
         output_path / 'step1_output',
         output_path / 'step1_levels',
-        canal_labels=[200, 201],
-        c2c3_label=224,
-        step=-1,
+        canal_labels=[1, 2],
+        disc_labels=list(range(60, 65)) + list(range(70, 82)) + list(range(90, 95)) + [100],
         override=True,
         max_workers=max_workers,
         quiet=quiet,
@@ -459,18 +496,14 @@ def main():
             quiet=quiet,
         )
 
-        # Load label mappings from JSON file
-        with open(resources_path / 'labels_maps' / 'nnunet_step2_input.json', 'r', encoding='utf-8') as map_file:
-            map_dict = json.load(map_file)
-
         if not quiet: print('\n' 'Mapping the IVDs labels from the step1 model output to the odd IVDs:')
         # This will also delete labels without odd IVDs
-        map_labels_mp(
+        extract_alternate_mp(
             output_path / 'step2_input',
             output_path / 'step2_input',
-            map_dict=map_dict,
             seg_suffix='_0001',
             output_seg_suffix='_0001',
+            labels=list(range(60, 101)),
             override=True,
             max_workers=max_workers,
             quiet=quiet,
@@ -543,17 +576,15 @@ def main():
             iterative_label_mp(
                 output_path / 'step2_output',
                 output_path / 'step2_output',
+                selected_disc_landmarks=[4, 7, 5, 6],
                 disc_labels=[1, 2, 3, 4, 5, 6, 7],
+                disc_landmark_labels=[4, 5, 6, 7],
+                disc_landmark_output_labels=[60, 70, 90, 100],
                 vertebrae_labels=[9, 10, 11, 12, 13, 14],
+                vertebrae_landmark_output_labels=[12, 20, 40, 50],
                 vertebrae_extra_labels=[8],
-                init_disc={4:224, 7:202, 5:219, 6:207},
-                init_vertebrae={11:40, 14:17, 12:34, 13:23},
-                step_diff_label=True,
-                step_diff_disc=True,
-                output_disc_step=-1,
-                output_vertebrae_step=-1,
-                map_output_dict={17:92},
-                map_input_dict={14:92, 15:201, 16:201, 17:200},
+                map_output_dict={17:50},
+                map_input_dict={14:50, 15:2, 16:2, 17:1},
                 override=True,
                 max_workers=max_workers,
                 quiet=quiet,
@@ -563,19 +594,16 @@ def main():
                 output_path / 'step2_output',
                 output_path / 'step2_output',
                 locs_path=output_path / 'localizers',
+                selected_disc_landmarks=[4, 7],
                 disc_labels=[1, 2, 3, 4, 5, 6, 7],
+                disc_landmark_labels=[4, 5, 6, 7],
+                disc_landmark_output_labels=[60, 70, 90, 100],
                 vertebrae_labels=[9, 10, 11, 12, 13, 14],
+                vertebrae_landmark_output_labels=[12, 20, 40, 50],
                 vertebrae_extra_labels=[8],
-                init_disc={4:224, 7:202},
-                init_vertebrae={11:40, 14:17},
-                loc_disc_labels=list(range(202, 225)),
-                loc_vertebrae_labels=list(range(18, 42)) + [92],
-                step_diff_label=True,
-                step_diff_disc=True,
-                output_disc_step=-1,
-                output_vertebrae_step=-1,
-                map_output_dict={17:92},
-                map_input_dict={14:92, 15:201, 16:201, 17:200},
+                loc_disc_labels=list(range(60, 101)),
+                map_output_dict={17:50},
+                map_input_dict={14:50, 15:2, 16:2, 17:1},
                 override=True,
                 max_workers=max_workers,
                 quiet=quiet,
@@ -586,8 +614,8 @@ def main():
         fill_canal_mp(
             output_path / 'step2_output',
             output_path / 'step2_output',
-            canal_label=201,
-            cord_label=200,
+            canal_label=2,
+            cord_label=1,
             largest_canal=True,
             largest_cord=True,
             override=True,
@@ -615,6 +643,27 @@ def main():
             max_workers=max_workers,
             quiet=quiet,
         )
+        preview_jpg_mp(
+            output_path / 'input',
+            output_path / 'preview',
+            segs_path=output_path / 'step2_output',
+            output_suffix='_step2_output_tags',
+            override=True,
+            max_workers=max_workers,
+            quiet=quiet,
+            label_texts_right={
+                10: 'C1', 11: 'C2', 12: 'C3', 13: 'C4', 14: 'C5', 15: 'C6', 16: 'C7',
+                20: 'T1', 21: 'T2', 22: 'T3', 23: 'T4', 24: 'T5', 25: 'T6', 26: 'T7',
+                27: 'T8', 28: 'T9', 29: 'T10', 30: 'T11', 31: 'T12',
+                40: 'L1', 41: 'L2', 42: 'L3', 43: 'L4', 44: 'L5'
+            },
+            label_texts_left={
+                50: 'Sacrum', 60: 'C2C3', 61: 'C3C4', 62: 'C4C5', 63: 'C5C6', 64: 'C6C7', 70: 'C7T1',
+                71: 'T1T2', 72: 'T2T3', 73: 'T3T4', 74: 'T4T5', 75: 'T5T6', 76: 'T6T7', 77: 'T7T8',
+                78: 'T8T9', 79: 'T9T10', 80: 'T10T11', 81: 'T11T12', 90: 'T12L1',
+                91: 'L1L2', 92: 'L2L3', 93: 'L3L4', 94: 'L4L5', 100: 'L5S'
+            },
+        )
 
 if __name__ == '__main__':
     main()