Merge pull request #36 from CAAI/dev/ts2

Totalseg update

.gitignore

@@ -2,7 +2,7 @@
 __pycache__/
 *.py[cod]
 *$py.class
-
+*.nii.gz
 # C extensions
 *.so
 .vscode

nodes/totalsegmentator/test_totalsegmentator.py

@@ -1,24 +1,38 @@
+#from rhnode import RHJob
 from rhnode import RHJob
 
 data = {
     #"in_file": "/homes/claes/projects/LowdosePET/Cu64DOTATATE/Data/mnc/DOTA_000/CT.nii",
-    "in_file": "/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",
+    "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,
-    'roi_subset': 'brain skull',
-    #'out_segmentation': 'test.nii.gz'
+   # 'roi_subset': 'brain skull',
+    'out_segmentation': 'test.nii.gz'
 }
 
 node = RHJob(
     node_name="totalsegmentator",
     inputs=data,
-    manager_address="titan5:9030",
+    #manager_address="http://127.0.0.1:8009",
+    node_address="localhost:8009",
     check_cache=False,
+    resources_included=True,
+    included_cuda_device=0
 )
 # Wait for the node to finish
 node.start()
 
 output = node.wait_for_finish()
 print(output)
 
+# totalsegmentator(
+#         str(data["in_file"]), "/homes/hinge/Projects/rh-library/nodes/totalsegmentator/out.nii.gz",ml=True, nr_thr_resamp=1, nr_thr_saving=1,
+#         fast=False, nora_tag="None", preview=False, task="total", 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=False, device="gpu", license_number=None,
+#         statistics_exclude_masks_at_border=True, no_derived_masks=False,
+#         v1_order=False, fastest=False, roi_subset_robust=None
+# )
+
 # Alternatively to interrupt the job:
-# node.stop()
+# node.stop()

nodes/totalsegmentator/totalsegmentator.py

@@ -3,6 +3,8 @@
 from typing import Optional
 import subprocess
 import os
+import nibabel as nib
+
 
 class TotalSegmentatorInput(BaseModel):
     in_file:FilePath
@@ -11,6 +13,8 @@ class TotalSegmentatorInput(BaseModel):
 
 class TotalSegmentatorOutput(BaseModel):
     out_segmentation:FilePath
+    out_version:str
+    out_args:str
 
 class TotalSegmentatorNode(RHNode):
     input_spec = TotalSegmentatorInput
@@ -23,18 +27,30 @@ class TotalSegmentatorNode(RHNode):
     def process(inputs, job):
         out_file = job.directory / 'segmentation.nii.gz'
 
-        cmd = ["TotalSegmentator", "-i", str(inputs.in_file), "-o", str(out_file), "--ml"]
+        cmd = ["TotalSegmentator", "-i", str(inputs.in_file), "-o", str(out_file)]
+
+        cmd_args = ["--ml"]
 
         if inputs.fast:
-            cmd += ['--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"]
+
         if not inputs.roi_subset == "":
-            cmd += ['--roi_subset'] + inputs.roi_subset.split(" ")
+            cmd_args += ['--roi_subset'] + inputs.roi_subset.split(" ")
 
         all_env_vars = os.environ.copy()
         all_env_vars.update({"CUDA_VISIBLE_DEVICES": str(job.device)})
-        out = subprocess.check_output(cmd, text=True,env=all_env_vars)
+
+        print("TOTALSEGMENTATOR STARTING VERSION")
+        version = subprocess.check_output("TotalSegmentator --version".split(" "),env=all_env_vars)        
+        print("TOTALSEGMENTATOR STARTING SEGMENTATION")
+        out = subprocess.check_output(cmd+cmd_args, text=True,env=all_env_vars)
+        print("TOTALSEGMENTATOR ENDING")
 
-        return TotalSegmentatorOutput(out_segmentation=out_file)
+        return TotalSegmentatorOutput(out_segmentation=out_file,out_version=version,out_args=" ".join(cmd_args))
 
 app = TotalSegmentatorNode()