Fixed #19 Added --body_seg option

README.md

@@ -48,6 +48,7 @@ If you have a powerful GPU is available then a full-quality segmentation can be
   - Input volume: input CT image
   - Segmentation task: instead of the default "total" segmentation, a more specialized segmentation model can be chosen
   - Fast: performs segmentation faster, but at lower resolution
+  - Body crop: crops the images to the body region before processing them, saves GPU memory
 - Outputs
   - Segmentation: it will contain a brain segment, which specifies the brain region
   - Show 3D: show/hide segments in 3D views

TotalSegmentator/Resources/UI/TotalSegmentator.ui

@@ -12,8 +12,8 @@
   </property>
   <layout class="QVBoxLayout" name="verticalLayout">
    <item>
-    <widget class="ctkCollapsibleButton" name="inputsCollapsibleButton">
-     <property name="text">
+    <widget class="ctkCollapsibleButton" name="inputsCollapsibleButton" native="true">
+     <property name="text" stdset="0">
       <string>Inputs</string>
      </property>
      <layout class="QFormLayout" name="formLayout_2">
@@ -25,22 +25,22 @@
        </widget>
       </item>
       <item row="0" column="1">
-       <widget class="qMRMLNodeComboBox" name="inputVolumeSelector">
+       <widget class="qMRMLNodeComboBox" name="inputVolumeSelector" native="true">
         <property name="toolTip">
          <string>Input abdominal, chest, or whole body CT.</string>
         </property>
-        <property name="nodeTypes">
+        <property name="nodeTypes" stdset="0">
          <stringlist>
           <string>vtkMRMLScalarVolumeNode</string>
          </stringlist>
         </property>
-        <property name="showChildNodeTypes">
+        <property name="showChildNodeTypes" stdset="0">
          <bool>false</bool>
         </property>
-        <property name="addEnabled">
+        <property name="addEnabled" stdset="0">
          <bool>false</bool>
         </property>
-        <property name="removeEnabled">
+        <property name="removeEnabled" stdset="0">
          <bool>false</bool>
         </property>
        </widget>
@@ -90,12 +90,30 @@
         </item>
        </layout>
       </item>
+      <item row="3" column="0">
+       <widget class="QLabel" name="label_2">
+        <property name="text">
+         <string>Body crop:</string>
+        </property>
+       </widget>
+      </item>
+      <item row="3" column="1">
+       <layout class="QHBoxLayout" name="horizontalLayout_2">
+        <item>
+         <widget class="QCheckBox" name="bodycropCheckBox">
+          <property name="text">
+           <string/>
+          </property>
+         </widget>
+        </item>
+       </layout>
+      </item>
      </layout>
     </widget>
    </item>
    <item>
-    <widget class="ctkCollapsibleButton" name="outputsCollapsibleButton">
-     <property name="text">
+    <widget class="ctkCollapsibleButton" name="outputsCollapsibleButton" native="true">
+     <property name="text" stdset="0">
       <string>Outputs</string>
      </property>
      <layout class="QFormLayout" name="formLayout_4">
@@ -107,37 +125,37 @@
        </widget>
       </item>
       <item row="0" column="1">
-       <widget class="qMRMLNodeComboBox" name="outputSegmentationSelector">
+       <widget class="qMRMLNodeComboBox" name="outputSegmentationSelector" native="true">
         <property name="toolTip">
          <string>This will store the segmentation result.</string>
         </property>
-        <property name="nodeTypes">
+        <property name="nodeTypes" stdset="0">
          <stringlist>
           <string>vtkMRMLSegmentationNode</string>
          </stringlist>
         </property>
-        <property name="showChildNodeTypes">
+        <property name="showChildNodeTypes" stdset="0">
          <bool>false</bool>
         </property>
-        <property name="baseName">
+        <property name="baseName" stdset="0">
          <string>Brain</string>
         </property>
-        <property name="noneEnabled">
+        <property name="noneEnabled" stdset="0">
          <bool>true</bool>
         </property>
-        <property name="addEnabled">
+        <property name="addEnabled" stdset="0">
          <bool>true</bool>
         </property>
-        <property name="removeEnabled">
+        <property name="removeEnabled" stdset="0">
          <bool>true</bool>
         </property>
-        <property name="editEnabled">
+        <property name="editEnabled" stdset="0">
          <bool>true</bool>
         </property>
-        <property name="renameEnabled">
+        <property name="renameEnabled" stdset="0">
          <bool>true</bool>
         </property>
-        <property name="noneDisplay">
+        <property name="noneDisplay" stdset="0">
          <string>Create new segmentation on Apply</string>
         </property>
        </widget>
@@ -149,23 +167,23 @@
     </widget>
    </item>
    <item>
-    <widget class="ctkCollapsibleButton" name="advancedCollapsibleButton">
-     <property name="text">
+    <widget class="ctkCollapsibleButton" name="advancedCollapsibleButton" native="true">
+     <property name="text" stdset="0">
       <string>Advanced</string>
      </property>
-     <property name="collapsed">
+     <property name="collapsed" stdset="0">
       <bool>true</bool>
      </property>
      <layout class="QFormLayout" name="formLayout_3">
       <item row="3" column="0" colspan="2">
        <widget class="ctkFittedTextBrowser" name="packageInfoTextBrowser">
-        <property name="collapsed">
+        <property name="collapsed" stdset="0">
          <bool>false</bool>
         </property>
-        <property name="showDetailsText">
+        <property name="showDetailsText" stdset="0">
          <string/>
         </property>
-        <property name="hideDetailsText">
+        <property name="hideDetailsText" stdset="0">
          <string/>
         </property>
        </widget>

TotalSegmentator/TotalSegmentator.py

@@ -198,6 +198,7 @@ def updateGUIFromParameterNode(self, caller=None, event=None):
         task = self._parameterNode.GetParameter("Task")
         self.ui.taskComboBox.setCurrentIndex(self.ui.taskComboBox.findData(task))
         self.ui.fastCheckBox.checked = self._parameterNode.GetParameter("Fast") == "true"
+        self.ui.bodycropCheckBox.checked = self._parameterNode.GetParameter("BodyCrop") == "true"
         self.ui.useStandardSegmentNamesCheckBox.checked = self._parameterNode.GetParameter("UseStandardSegmentNames") == "true"
         self.ui.outputSegmentationSelector.setCurrentNode(self._parameterNode.GetNodeReference("OutputSegmentation"))
 
@@ -234,6 +235,7 @@ def updateParameterNodeFromGUI(self, caller=None, event=None):
         self._parameterNode.SetNodeReferenceID("InputVolume", self.ui.inputVolumeSelector.currentNodeID)
         self._parameterNode.SetParameter("Task", self.ui.taskComboBox.currentData)
         self._parameterNode.SetParameter("Fast", "true" if self.ui.fastCheckBox.checked else "false")
+        self._parameterNode.SetParameter("BodyCrop", "true" if self.ui.bodycropCheckBox.checked else "false")
         self._parameterNode.SetParameter("UseStandardSegmentNames", "true" if self.ui.useStandardSegmentNamesCheckBox.checked else "false")
         self._parameterNode.SetNodeReferenceID("OutputSegmentation", self.ui.outputSegmentationSelector.currentNodeID)
 
@@ -262,7 +264,7 @@ def onApplyButton(self):
 
             # Compute output
             self.logic.process(self.ui.inputVolumeSelector.currentNode(), self.ui.outputSegmentationSelector.currentNode(),
-                self.ui.fastCheckBox.checked, self.ui.taskComboBox.currentData)
+                self.ui.fastCheckBox.checked, self.ui.bodycropCheckBox.checked, self.ui.taskComboBox.currentData)
 
         self.ui.statusLabel.appendPlainText("\nProcessing finished.")
 
@@ -631,6 +633,8 @@ def setDefaultParameters(self, parameterNode):
         """
         if not parameterNode.GetParameter("Fast"):
             parameterNode.SetParameter("Fast", "True")
+        if not parameterNode.GetParameter("BodyCrop"):
+            parameterNode.SetParameter("BodyCrop", "False")
         if not parameterNode.GetParameter("Task"):
             parameterNode.SetParameter("Task", "total")
         if not parameterNode.GetParameter("UseStandardSegmentNames"):
@@ -655,14 +659,15 @@ def logProcessOutput(self, proc):
         if retcode != 0:
             raise CalledProcessError(retcode, proc.args, output=proc.stdout, stderr=proc.stderr)
 
-    def process(self, inputVolume, outputSegmentation, fast=True, task=None):
+    def process(self, inputVolume, outputSegmentation, fast=True, bodycrop=False, task=None):
 
         """
         Run the processing algorithm.
         Can be used without GUI widget.
         :param inputVolume: volume to be thresholded
         :param outputVolume: thresholding result
         :param fast: faster and less accurate output
+        :param bodycrop: crop the image to the body region before processing it, saves GPU memory
         :param task: one of self.tasks, default is "total"
         """
 
@@ -757,6 +762,8 @@ def process(self, inputVolume, outputSegmentation, fast=True, task=None):
             options.extend(["--task", task])
         if fast:
             options.append("--fast")
+        if bodycrop:
+            options.append("--body_seg")
         self.log('Creating segmentations with TotalSegmentator AI...')
         self.log(f"Total Segmentator arguments: {options}")
         proc = slicer.util.launchConsoleProcess(totalSegmentatorCommand + options)