From e8b6e9f8fc1d83cb50d00a9f6e6d3f70823e4efc Mon Sep 17 00:00:00 2001 From: Louis Blankemeier Date: Thu, 5 Oct 2023 14:02:01 -0700 Subject: [PATCH] Spine v3 (#123) * update spine to ts v2 * conda name * Autoformat code --------- Co-authored-by: louisblankemeier --- bin/C2C-slurm | 2 +- bin/install.sh | 8 ++--- comp2comp/models/models.py | 3 +- comp2comp/spine/spine.py | 68 +++++++++++++++++++++++++++++++++----- comp2comp/utils/process.py | 9 +++-- setup.py | 3 +- 6 files changed, 74 insertions(+), 19 deletions(-) diff --git a/bin/C2C-slurm b/bin/C2C-slurm index d304f85..9b0daaa 100755 --- 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 index 4e4c8bb..9d00a70 100755 --- 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 index 3f6bcf4..b972f64 100644 --- 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 index 0ae96fc..42c8fd9 100644 --- a/comp2comp/spine/spine.py +++ b/comp2comp/spine/spine.py @@ -15,11 +15,7 @@ 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 @@ -27,6 +23,14 @@ 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 index b10198a..491f957 100644 --- 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 index 363ee97..578e336 100644 --- 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"],