Spine cropper

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]

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
+
+