merged

nodes/totalsegmentator/Dockerfile

@@ -8,8 +8,8 @@ RUN pip install git+https://github.com/wasserth/TotalSegmentator.git
 RUN python -c "from totalsegmentator.libs import download_pretrained_weights; [download_pretrained_weights(task_id) for task_id in [291, 292, 293, 294, 295, 297, 298, 258, 150, 260, 503, 315, 299, 300]]"
 
 # #RH-Node
-RUN pip install git+https://github.com/CAAI/[email protected]
 
+RUN pip install git+https://github.com/CAAI/[email protected]
 COPY totalsegmentator.py /app/totalsegmentator.py
 WORKDIR /app
 ENV PYTHONUNBUFFERED 1

nodes/totalsegmentator/test_totalsegmentator.py

@@ -5,6 +5,7 @@
     #"in_file": "/homes/claes/projects/LowdosePET/Cu64DOTATATE/Data/mnc/DOTA_000/CT.nii",
     "in_file": "/homes/hinge/Projects/rh-library/nodes/totalsegmentator/testlarge.nii.gz",#"/homes/hinge/lymphoma/data/raw/RIS_PACS/LM12533_02/CTres_crop.nii.gz",#"/homes/claes/projects/LowdosePET/PE2I/data_Vision/nii/0rCyjRXXbm_0/ACCT.nii.gz",
     'fast': True,
+    "task":"body",
    # 'roi_subset': 'brain skull',
     'out_segmentation': 'test.nii.gz'
 }

nodes/totalsegmentator/totalsegmentator.py

@@ -10,6 +10,7 @@ class TotalSegmentatorInput(BaseModel):
     in_file:FilePath
     fast:Optional[bool]=False
     roi_subset:Optional[str]=""
+    task:str="total"
 
 class TotalSegmentatorOutput(BaseModel):
     out_segmentation:FilePath
@@ -24,20 +25,20 @@ class TotalSegmentatorNode(RHNode):
     required_num_threads = 2
     required_gb_memory = 12    
 
-    def process(inputs, job):
+    def process(inputs:TotalSegmentatorInput, job) -> TotalSegmentatorOutput:
         out_file = job.directory / 'segmentation.nii.gz'
-
-        cmd = ["TotalSegmentator", "-i", str(inputs.in_file), "-o", str(out_file)]
+        
+        cmd = ["TotalSegmentator", "-i", str(inputs.in_file), "-o", str(out_file), "-ta", inputs.task]
 
         cmd_args = ["--ml"]
 
         if inputs.fast:
             cmd_args += ['--fast']
 
-        shape = nib.load(str(inputs.in_file)).shape
-        if shape[-1] > 590:
-            print("Large image, running with --body_seg --force_split --nr_thr_saving 1")
-            cmd_args+=["--body_seg","--force_split","--nr_thr_saving","1"]
+        #shape = nib.load(str(inputs.in_file)).shape
+        # if shape[-1] > 590:
+        #     print("Large image, running with --body_seg --force_split --nr_thr_saving 1")
+        #     cmd_args+=["--body_seg","--force_split","--nr_thr_saving","1"]
 
         if not inputs.roi_subset == "":
             cmd_args += ['--roi_subset'] + inputs.roi_subset.split(" ")