bugfixes for roi_subset mr

Author: wasserth

tests/reference_files/example_seg_roi_subset_mr.nii.gz

@@ -58,6 +58,13 @@ def test_prediction_liver_roi_subset(self):
         dice = dice_score_multilabel(img_ref, img_new)
         images_equal = dice > 0.99
         self.assertTrue(images_equal, f"roi subset prediction not correct (dice: {dice:.6f})")
+        
+    def test_prediction_liver_roi_subset_mr(self):
+        img_ref = nib.load("tests/reference_files/example_seg_roi_subset_mr.nii.gz").get_fdata()
+        img_new = nib.load("tests/unittest_prediction_roi_subset_mr.nii.gz").get_fdata()
+        dice = dice_score_multilabel(img_ref, img_new)
+        images_equal = dice > 0.99
+        self.assertTrue(images_equal, f"roi subset MR prediction not correct (dice: {dice:.6f})")
 
     def test_preview(self):
         preview_exists = os.path.exists("tests/unittest_prediction_fast/preview_total.png")

tests/test_end_to_end.py

@@ -4,18 +4,22 @@ set -e
 # ./tests/tests.sh
 
 
-# Test multilabel prediction
+# Test - multilabel prediction
 TotalSegmentator -i tests/reference_files/example_ct_sm.nii.gz -o tests/unittest_prediction.nii.gz -bs --ml -d cpu
 pytest -v tests/test_end_to_end.py::test_end_to_end::test_prediction_multilabel
 
-# Test organ prediction - roi subset
+# Test - roi subset
 # 2 cpus:
 #   example_ct_sm.nii.gz: 34s, 3.0GB
 #   example_ct.nii.gz: 36s, 3.0GB
 TotalSegmentator -i tests/reference_files/example_ct_sm.nii.gz -o tests/unittest_prediction_roi_subset.nii.gz --ml -rs liver brain -d cpu
 pytest -v tests/test_end_to_end.py::test_end_to_end::test_prediction_liver_roi_subset
 
-# Test organ predictions - fast - statistics
+# Test - roi subset - MR
+TotalSegmentator -i tests/reference_files/example_mr_sm.nii.gz -o tests/unittest_prediction_roi_subset_mr.nii.gz -ta total_mr --ml -rs liver brain -d cpu
+pytest -v tests/test_end_to_end.py::test_end_to_end::test_prediction_liver_roi_subset_mr
+
+# Test - fast - statistics
 # 2 cpus: (statistics <1s)
 #   example_ct_sm.nii.gz: 13s, 4.1GB
 #   example_ct.nii.gz: 16s, 4.1GB
@@ -24,15 +28,15 @@ pytest -v tests/test_end_to_end.py::test_end_to_end::test_prediction_fast
 pytest -v tests/test_end_to_end.py::test_end_to_end::test_statistics
 pytest -v tests/test_end_to_end.py::test_end_to_end::test_preview
 
-# Test organ predictions - fast - multilabel
+# Test - fast - multilabel
 TotalSegmentator -i tests/reference_files/example_ct_sm.nii.gz -o tests/unittest_prediction_fast.nii.gz --fast --ml -d cpu
 pytest -v tests/test_end_to_end.py::test_end_to_end::test_prediction_multilabel_fast
 
-# Test organ predictions - fast - multilabel - force split
+# Test - fast - multilabel - force split
 TotalSegmentator -i tests/reference_files/example_ct.nii.gz -o tests/unittest_prediction_fast_force_split.nii.gz --fast --ml -fs -d cpu
 pytest -v tests/test_end_to_end.py::test_end_to_end::test_prediction_multilabel_fast_force_split
 
-# Test organ predictions - fast - multilabel - body_seg
+# Test - fast - multilabel - body_seg
 TotalSegmentator -i tests/reference_files/example_ct_sm.nii.gz -o tests/unittest_prediction_fast_body_seg.nii.gz --fast --ml -bs -d cpu
 pytest -v tests/test_end_to_end.py::test_end_to_end::test_prediction_multilabel_fast_body_seg
 

tests/tests.sh

@@ -8,6 +8,7 @@ set -e
 TotalSegmentator -i tests/reference_files/example_ct_sm.nii.gz -o tests/reference_files/example_seg.nii.gz -bs --ml -d cpu
 TotalSegmentator -i tests/reference_files/example_mr_sm.nii.gz -o tests/reference_files/example_seg_mr.nii.gz -ta total_mr --ml -d cpu
 TotalSegmentator -i tests/reference_files/example_ct_sm.nii.gz -o tests/reference_files/example_seg_roi_subset.nii.gz --ml -rs liver brain -d cpu
+TotalSegmentator -i tests/reference_files/example_mr_sm.nii.gz -o tests/reference_files/example_seg_roi_subset_mr.nii.gz -ta total_mr --ml -rs liver brain -d cpu
 # TotalSegmentator -i tests/reference_files/example_ct_sm.nii.gz -o tests/reference_files/example_seg_fast --fast --statistics -sii -p -d cpu
 TotalSegmentator -i tests/reference_files/example_ct_sm.nii.gz -o tests/reference_files/example_seg_fast.nii.gz --fast --ml -d cpu
 TotalSegmentator -i tests/reference_files/example_ct.nii.gz -o tests/reference_files/example_seg_fast_force_split.nii.gz --fast --ml -fs -d cpu

tests/update_test_files.sh

@@ -363,6 +363,14 @@
     "face": {
         1: "face"
     },
+    "face_mr": {
+        1: "face"
+    },
+    # those classes need to be removed
+    "face_mr_auxiliary": {
+        2: "brain",
+        3: "liver"
+    },
     "test": {
         1: "carpal"
     }

totalsegmentator/map_to_binary.py

@@ -130,6 +130,9 @@
     "face": [
         ["face"]
     ],
+    "face_mr": [
+        ["face"]
+    ],
     # "aortic_branches_test": [
     #     ["brachiocephalic_trunk", "subclavian_artery_right", "subclavian_artery_left", "aorta",
     #     "common_carotid_artery_right", "common_carotid_artery_left"],

totalsegmentator/preview.py

@@ -305,6 +305,11 @@ def totalsegmentator(input: Union[str, Path, Nifti1Image], output: Union[str, Pa
     else:
         download_pretrained_weights(task_id)
 
+    # For MR always run 3mm model for roi_subset, because 6mm too bad results
+    #  (runtime for 3mm still very good for MR)
+    if task.endswith("_mr") and roi_subset is not None:
+        roi_subset_robust = roi_subset
+
     if roi_subset_robust is not None:
         roi_subset = roi_subset_robust
         robust_rs = True
@@ -313,10 +318,10 @@ def totalsegmentator(input: Union[str, Path, Nifti1Image], output: Union[str, Pa
 
     if roi_subset is not None and type(roi_subset) is not list:
         raise ValueError("roi_subset must be a list of strings")
-    if roi_subset is not None and task != "total":
-        raise ValueError("roi_subset only works with task 'total'")
+    if roi_subset is not None and not task.startswith("total"):
+        raise ValueError("roi_subset only works with task 'total' or 'total_mr'")
 
-    if task == "total_mr" or task == "tissue_types_mr":
+    if task.endswith("_mr"):
         if body_seg:
             body_seg = False
             print("INFO: For MR models the argument '--body_seg' is not supported and will be ignored.")
@@ -335,13 +340,15 @@ def totalsegmentator(input: Union[str, Path, Nifti1Image], output: Union[str, Pa
         else:
             crop_model_task = 733 if task == "total_mr" else 298
             crop_spacing = 6.0
+        crop_task = "total_mr" if task == "total_mr" else "total"
+        crop_trainer = "nnUNetTrainer_DASegOrd0_NoMirroring" if task == "total_mr" else "nnUNetTrainer_4000epochs_NoMirroring"
         organ_seg, _, _ = nnUNet_predict_image(input, None, crop_model_task, model="3d_fullres", folds=[0],
-                            trainer="nnUNetTrainer_4000epochs_NoMirroring", tta=False, multilabel_image=True, resample=crop_spacing,
-                            crop=None, crop_path=None, task="total", nora_tag="None", preview=False,
+                            trainer=crop_trainer, tta=False, multilabel_image=True, resample=crop_spacing,
+                            crop=None, crop_path=None, task_name=crop_task, nora_tag="None", preview=False,
                             save_binary=False, nr_threads_resampling=nr_thr_resamp, nr_threads_saving=1,
                             crop_addon=crop_addon, output_type=output_type, statistics=False,
                             quiet=quiet, verbose=verbose, test=0, skip_saving=False, device=device)
-        class_map_inv = {v: k for k, v in class_map["total"].items()}
+        class_map_inv = {v: k for k, v in class_map[crop_task].items()}
         crop_mask = np.zeros(organ_seg.shape, dtype=np.uint8)
         organ_seg_data = organ_seg.get_fdata()
         # roi_subset_crop = [map_to_total[roi] if roi in map_to_total else roi for roi in roi_subset]