diff --git a/comp2comp/spine/spine.py b/comp2comp/spine/spine.py index 42c8fd9..988871a 100644 --- a/comp2comp/spine/spine.py +++ b/comp2comp/spine/spine.py @@ -15,13 +15,14 @@ import pandas as pd import wget from PIL import Image + 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 comp2comp.io import io_utils # from totalsegmentator.libs import ( # download_pretrained_weights, @@ -241,7 +242,7 @@ def __init__(self, lower_level: str = "L5", upper_level: str = "L1", save=True): super().__init__() self.lower_level = lower_level self.upper_level = upper_level - ts_spine_full_model = Models.model_from_name("ts_spine_full") + ts_spine_full_model = Models.model_from_name("ts_spine") categories = ts_spine_full_model.categories try: self.lower_level_index = categories[self.lower_level] diff --git a/comp2comp/utils/process.py b/comp2comp/utils/process.py index 491f957..c5c2741 100644 --- a/comp2comp/utils/process.py +++ b/comp2comp/utils/process.py @@ -10,6 +10,7 @@ from datetime import datetime from pathlib import Path + from comp2comp.io import io_utils @@ -118,3 +119,5 @@ def process_3d(args, pipeline_builder): if len(os.listdir(os.path.dirname(output_dir))) == 0: shutil.rmtree(os.path.dirname(output_dir)) continue + +