Spine v3 (#123)

* update spine to ts v2 * conda name * Autoformat code --------- Co-authored-by: louisblankemeier <[email protected]>
StanfordMIMI · Oct 5, 2023 · e8b6e9f · e8b6e9f
1 parent b03bf83
commit e8b6e9f
Show file tree

Hide file tree

Showing 6 changed files with 74 additions and 19 deletions.
diff --git a/bin/C2C-slurm b/bin/C2C-slurm
@@ -13,7 +13,7 @@ def submit_command(command):
     subprocess.run(command.split(" "), check=True, capture_output=False)
 
 
-def python_submit(command, node="amalfi"):
+def python_submit(command, node="siena"):
     bash_file = open("./slurm.sh", "w")
     bash_file.write(f"#!/bin/bash\n{command}")
     bash_file.close()

diff --git a/bin/install.sh b/bin/install.sh
@@ -18,7 +18,7 @@ ABCT_ENV_NAME="c2c_env"
 hasAnaconda=0
 updateEnv=0
 updatePath=1
-pythonVersion="3.8"
+pythonVersion="3.9"
 cudaVersion=""
 
 while [[ $# -gt 0 ]]; do
@@ -94,10 +94,10 @@ if [[ `conda env list | grep $ABCT_ENV_NAME` ]]; then
         exit 0
     else
         conda env remove -n $ABCT_ENV_NAME
-        conda create -y -n $ABCT_ENV_NAME python=3.8
+        conda create -y -n $ABCT_ENV_NAME python=3.9
     fi
 else
-    conda create -y -n $ABCT_ENV_NAME python=3.8
+    conda create -y -n $ABCT_ENV_NAME python=3.9
 fi
 
 conda activate $ABCT_ENV_NAME
@@ -139,7 +139,7 @@ conda activate $ABCT_ENV_NAME
 # echo $currDir
 # exit 1
 
-pip install -e .
+pip install -e . --no-cache-dir
 
 echo ""
 echo ""

diff --git a/comp2comp/models/models.py b/comp2comp/models/models.py
@@ -70,7 +70,8 @@ class Models(enum.Enum):
         5,
         "ts_spine",
         # Category name mapped to channel index
-        {"L5": 18, "L4": 19, "L3": 20, "L2": 21, "L1": 22, "T12": 23},
+        # {"L5": 18, "L4": 19, "L3": 20, "L2": 21, "L1": 22, "T12": 23},
+        {"L5": 27, "L4": 28, "L3": 29, "L2": 30, "L1": 31, "T12": 32},
         False,
         (),
     )

diff --git a/comp2comp/spine/spine.py b/comp2comp/spine/spine.py
@@ -15,18 +15,22 @@
 import pandas as pd
 import wget
 from PIL import Image
-from totalsegmentator.libs import (
-    download_pretrained_weights,
-    nostdout,
-    setup_nnunet,
-)
+from totalsegmentatorv2.python_api import totalsegmentator
 
 from comp2comp.inference_class_base import InferenceClass
 from comp2comp.io import io_utils
 from comp2comp.models.models import Models
 from comp2comp.spine import spine_utils
 from comp2comp.visualization.dicom import to_dicom
 
+# from totalsegmentator.libs import (
+#     download_pretrained_weights,
+#     nostdout,
+#     setup_nnunet,
+# )
+
+
+
 
 class SpineSegmentation(InferenceClass):
     """Spine segmentation."""
@@ -45,11 +49,56 @@ def __call__(self, inference_pipeline):
 
         self.model_dir = inference_pipeline.model_dir
 
-        seg, mv = self.spine_seg(
-            os.path.join(self.output_dir_segmentations, "converted_dcm.nii.gz"),
-            self.output_dir_segmentations + "spine.nii.gz",
-            inference_pipeline.model_dir,
+        # seg, mv = self.spine_seg(
+        #     os.path.join(self.output_dir_segmentations, "converted_dcm.nii.gz"),
+        #     self.output_dir_segmentations + "spine.nii.gz",
+        #     inference_pipeline.model_dir,
+        # )
+        os.environ["TOTALSEG_WEIGHTS_PATH"] = self.model_dir
+
+        seg = totalsegmentator(
+            input=os.path.join(self.output_dir_segmentations, "converted_dcm.nii.gz"),
+            output=os.path.join(self.output_dir_segmentations, "segmentation.nii"),
+            task_ids=[292],
+            ml=True,
+            nr_thr_resamp=1,
+            nr_thr_saving=6,
+            fast=False,
+            nora_tag="None",
+            preview=False,
+            task="total",
+            # roi_subset=[
+            #     "vertebrae_T12",
+            #     "vertebrae_L1",
+            #     "vertebrae_L2",
+            #     "vertebrae_L3",
+            #     "vertebrae_L4",
+            #     "vertebrae_L5",
+            # ],
+            roi_subset=None,
+            statistics=False,
+            radiomics=False,
+            crop_path=None,
+            body_seg=False,
+            force_split=False,
+            output_type="nifti",
+            quiet=False,
+            verbose=False,
+            test=0,
+            skip_saving=True,
+            device="gpu",
+            license_number=None,
+            statistics_exclude_masks_at_border=True,
+            no_derived_masks=False,
+            v1_order=False,
         )
+        mv = nib.load(
+            os.path.join(self.output_dir_segmentations, "converted_dcm.nii.gz")
+        )
+
+        # inference_pipeline.segmentation = nib.load(
+        #     os.path.join(self.output_dir_segmentations, "segmentation.nii")
+        # )
         inference_pipeline.segmentation = seg
         inference_pipeline.medical_volume = mv
         inference_pipeline.save_segmentations = self.save_segmentations
@@ -115,6 +164,7 @@ def spine_seg(
         print("Segmenting spine...")
         st = time()
         os.environ["SCRATCH"] = self.model_dir
+        os.environ["TOTALSEG_WEIGHTS_PATH"] = self.model_dir
 
         # Setup nnunet
         model = "3d_fullres"

diff --git a/comp2comp/utils/process.py b/comp2comp/utils/process.py
@@ -5,10 +5,10 @@
 import os
 import shutil
 import sys
+import time
 import traceback
 from datetime import datetime
 from pathlib import Path
-from time import time
 
 from comp2comp.io import io_utils
 
@@ -51,7 +51,7 @@ def process_3d(args, pipeline_builder):
 
     for path, num in io_utils.get_dicom_or_nifti_paths_and_num(args.input_path):
         try:
-            st = time()
+            st = time.time()
 
             if path.endswith(".nii") or path.endswith(".nii.gz"):
                 print("Processing: ", path)
@@ -107,11 +107,14 @@ def process_3d(args, pipeline_builder):
                 if os.path.exists(segmentations_dir):
                     shutil.rmtree(segmentations_dir)
 
-            print(f"Finished processing {path} in {time() - st:.1f} seconds\n")
+            print(f"Finished processing {path} in {time.time() - st:.1f} seconds\n")
 
         except Exception:
             print(f"ERROR PROCESSING {path}\n")
             traceback.print_exc()
             if os.path.exists(output_dir):
                 shutil.rmtree(output_dir)
+            # remove parent folder if empty
+            if len(os.listdir(os.path.dirname(output_dir))) == 0:
+                shutil.rmtree(os.path.dirname(output_dir))
             continue
diff --git a/setup.py b/setup.py
@@ -41,7 +41,7 @@ def get_version():
     url="https://github.com/StanfordMIMI/Comp2Comp",
     description="Computed tomography to body composition.",
     packages=find_packages(exclude=("configs", "tests")),
-    python_requires=">=3.6",
+    python_requires=">=3.9",
     install_requires=[
         "pydicom",
         "moviepy",
@@ -59,6 +59,7 @@ def get_version():
         "wget",
         "tensorflow==2.12.0",
         "totalsegmentator @ git+https://github.com/StanfordMIMI/TotalSegmentator.git",
+        "totalsegmentatorv2 @ git+https://github.com/StanfordMIMI/TotalSegmentatorV2.git",
     ],
     extras_require={
         "all": ["shapely", "psutil"],