for phase prediction also return pi_time min and max

wasserth · Mar 18, 2024 · bd9e293 · bd9e293
1 parent e205e98
commit bd9e293
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/totalsegmentator/bin/totalseg_get_phase.py b/totalsegmentator/bin/totalseg_get_phase.py
@@ -13,6 +13,11 @@
 from totalsegmentator.statistics import get_basic_statistics
 
 
+"""
+Additional requirements for this script:
+xgboost
+"""
+
 def pi_time_to_phase(pi_time: float) -> str:
     """
     Convert the pi time to a phase and get a probability for the value.
@@ -78,7 +83,8 @@ def get_ct_contrast_phase(ct_img: nib.Nifti1Image):
     print(f"mean: {pi_time} [{preds.min():.1f}-{preds.max():.1f}]")
     phase, probability = pi_time_to_phase(pi_time)
 
-    return {"pi_time": pi_time, "phase": phase, "probability": probability}
+    return {"pi_time": pi_time, "phase": phase, "probability": probability, 
+            "pi_time_min": preds.min(), "pi_time_max": preds.max()}
 
 
 def main():

diff --git a/totalsegmentator/python_api.py b/totalsegmentator/python_api.py
@@ -248,7 +248,7 @@ def totalsegmentator(input: Union[str, Path, Nifti1Image], output: Union[str, Pa
 
     crop_path = output if crop_path is None else crop_path
 
-    if isinstance(input, Nifti1Image) or input.suffix == '.nii' or input.suffixes == ['.nii', '.gz']:
+    if isinstance(input, Nifti1Image) or input.suffix == ".nii" or input.suffixes == [".nii", ".gz"]:
         img_type = "nifti"
     else:
         img_type = "dicom"