ENH: Add SegmentationGeometryLogic API to set reference image geometry

Added SetReferenceImageGeometryInSegmentationNode function to vtkSlicerSegmentationGeometryLogic.
This feature was only available in the segmentation geometry widget, which was not convenient for scripting.

Co-authored-by: Jean-Christophe Fillion-Robin <jchris.fillionr@kitware.com>

Docs/developer_guide/script_repository/segmentations.md

@@ -635,6 +635,30 @@ sourceSegmentId = segmentation.GetSegmentIdBySegmentName(sourceSegmentName)
 segmentation.CopySegmentFromSegmentation(segmentation, sourceSegmentId)
 ```
 
+### Resample segmentation to higher resolution
+
+This code snippet can be used to resample internal binary labelmap representation of a segmentation to allow segmenting finer details
+
+```python
+# Set inputs
+volumeNode = getNode("MRHead")
+segmentationNode = getNode("Segmentation")
+
+# The higher the oversampling factor is the finer resolution the segmentation will be,
+# at the cost of more memory usage and longer computation times.
+# Note that oversampling by a factor of 2 increases memory usage by a factor of 2 * 2 * 2 = 8.
+oversamplingFactor = 2.0
+
+# Make spacing value uniform for all axes.
+# It is useful for removing staircase artifacts in 3D reconstructions but may increase
+# memory usage and computation time.
+isotropicSpacing = True
+
+# Update geometry of internal binary labelmap representation in segmentation node
+segmentationGeometryLogic.SetReferenceImageGeometryInSegmentationNode()
+segmentationGeometryLogic.ResampleLabelmapsInSegmentationNode()
+```
+
 ### Quantifying segments
 
 #### Get volume of each segment

Modules/Loadable/Segmentations/Logic/vtkSlicerSegmentationGeometryLogic.cxx

@@ -151,26 +151,45 @@ void vtkSlicerSegmentationGeometryLogic::CalculatePaddedOutputGeometry()
     return;
   }
 
-  std::string segmentationGeometryString = this->InputSegmentationNode->GetSegmentation()->DetermineCommonLabelmapGeometry(
-    vtkSegmentation::EXTENT_UNION_OF_EFFECTIVE_SEGMENTS);
-  vtkNew<vtkOrientedImageData> inputGeometryImageData;
-  vtkSegmentationConverter::DeserializeImageGeometry(segmentationGeometryString, inputGeometryImageData, false/*don't allocate*/);
-  vtkNew<vtkTransform> segmentationGeometryToReferenceGeometryTransform;
-  vtkOrientedImageDataResample::GetTransformBetweenOrientedImages(inputGeometryImageData,
-    this->OutputGeometryImageData, segmentationGeometryToReferenceGeometryTransform);
-
-  int transformedSegmentationExtent[6] = { 0, -1, 0, -1, 0, -1 };
-  vtkOrientedImageDataResample::TransformExtent(inputGeometryImageData->GetExtent(),
-    segmentationGeometryToReferenceGeometryTransform, transformedSegmentationExtent);
-
   int outputGeometryExtent[6] = { 0, -1, 0, -1, 0, -1 };
   this->OutputGeometryImageData->GetExtent(outputGeometryExtent);
 
-  for (int i = 0; i < 3; ++i)
-  {
-    outputGeometryExtent[2*i] = std::min(outputGeometryExtent[2*i], transformedSegmentationExtent[2*i]);
-    outputGeometryExtent[2*i+1] = std::max(outputGeometryExtent[2*i+1], transformedSegmentationExtent[2*i+1]);
+  std::string segmentationGeometryString = this->InputSegmentationNode->GetSegmentation()->DetermineCommonLabelmapGeometry(
+    vtkSegmentation::EXTENT_UNION_OF_EFFECTIVE_SEGMENTS);
+  if (!segmentationGeometryString.empty())
+  {
+    // There are non-empty segments, expand the output extent to make sure they are fully included
+    vtkNew<vtkOrientedImageData> segmentationGeometryImageData;
+    vtkSegmentationConverter::DeserializeImageGeometry(segmentationGeometryString, segmentationGeometryImageData, false/*don't allocate*/);
+    vtkNew<vtkTransform> segmentationGeometryToReferenceGeometryTransform;
+    vtkOrientedImageDataResample::GetTransformBetweenOrientedImages(segmentationGeometryImageData,
+      this->OutputGeometryImageData, segmentationGeometryToReferenceGeometryTransform);
+
+    int transformedSegmentationExtent[6] = { 0, -1, 0, -1, 0, -1 };
+    vtkOrientedImageDataResample::TransformExtent(segmentationGeometryImageData->GetExtent(),
+      segmentationGeometryToReferenceGeometryTransform, transformedSegmentationExtent);
+
+    if (outputGeometryExtent[0] > outputGeometryExtent[1]
+      || outputGeometryExtent[2] > outputGeometryExtent[3]
+      || outputGeometryExtent[4] > outputGeometryExtent[5])
+    {
+      // Output geometry extent is empty, use the segmentation geometry
+      for (int i = 0; i < 6; ++i)
+      {
+        outputGeometryExtent[i] = transformedSegmentationExtent[i];
+      }
+    }
+    else
+    {
+      // Both the output and segmentation is non-empty, expand the output geometry to include all the segments
+      for (int i = 0; i < 3; ++i)
+      {
+        outputGeometryExtent[2 * i] = std::min(outputGeometryExtent[2 * i], transformedSegmentationExtent[2 * i]);
+        outputGeometryExtent[2 * i + 1] = std::max(outputGeometryExtent[2 * i + 1], transformedSegmentationExtent[2 * i + 1]);
+      }
+    }
   }
+
   this->OutputGeometryImageData->SetExtent(outputGeometryExtent);
 }
 
@@ -552,7 +571,6 @@ bool vtkSlicerSegmentationGeometryLogic::ResampleLabelmapsInSegmentationNode()
   bool success = true;
 
   MRMLNodeModifyBlocker blocker(this->InputSegmentationNode);
-  vtkOrientedImageData* geometryImageData = this->GetOutputGeometryImageData();
   std::vector< std::string > segmentIDs;
   this->InputSegmentationNode->GetSegmentation()->GetSegmentIDs(segmentIDs);
   for (std::vector< std::string >::const_iterator segmentIdIt = segmentIDs.begin(); segmentIdIt != segmentIDs.end(); ++segmentIdIt)
@@ -572,7 +590,7 @@ bool vtkSlicerSegmentationGeometryLogic::ResampleLabelmapsInSegmentationNode()
 
     // Resample
     if (!vtkOrientedImageDataResample::ResampleOrientedImageToReferenceOrientedImage(
-      currentLabelmap, geometryImageData, currentLabelmap, false, this->PadOutputGeometry))
+      currentLabelmap, this->OutputGeometryImageData, currentLabelmap, false, this->PadOutputGeometry))
     {
       vtkErrorMacro("vtkSlicerSegmentationGeometryLogic::ResampleLabelmapsInSegmentationNode: "
         << "Segment " << this->InputSegmentationNode->GetName() << "/" << currentSegmentID.c_str() << " failed to be resampled");
@@ -584,3 +602,37 @@ bool vtkSlicerSegmentationGeometryLogic::ResampleLabelmapsInSegmentationNode()
   this->InputSegmentationNode->Modified();
   return success;
 }
+
+//-----------------------------------------------------------------------------
+bool vtkSlicerSegmentationGeometryLogic::SetReferenceImageGeometryInSegmentationNode()
+{
+  if (!this->InputSegmentationNode || !this->InputSegmentationNode->GetSegmentation())
+  {
+    vtkErrorMacro("vtkSlicerSegmentationGeometryLogic::SetReferenceImageGeometryInSegmentationNode: invalid input segmentation node");
+    return false;
+  }
+
+  // Save reference geometry
+  std::string geometryString = vtkSegmentationConverter::SerializeImageGeometry(this->OutputGeometryImageData);
+  if (geometryString.empty())
+  {
+    vtkErrorMacro("vtkSlicerSegmentationGeometryLogic::SetReferenceImageGeometryInSegmentationNode: invalid output geometry");
+    return false;
+  }
+  this->InputSegmentationNode->GetSegmentation()->SetConversionParameter(
+    vtkSegmentationConverter::GetReferenceImageGeometryParameterName(), geometryString);
+
+  // Save reference geometry node - but only if it is not the segmentation node itself.
+  // This is shown in Segmentations module to give a hint about which node the current geometry is based on
+  // and may be used when exporting the segmentation.
+  if (this->SourceGeometryNode
+    && this->SourceGeometryNode->GetID()
+    && vtkMRMLSegmentationNode::SafeDownCast(this->SourceGeometryNode) != this->InputSegmentationNode)
+  {
+    this->InputSegmentationNode->SetNodeReferenceID(
+      vtkMRMLSegmentationNode::GetReferenceImageGeometryReferenceRole().c_str(),
+      this->SourceGeometryNode->GetID());
+  }
+
+  return true;
+}

Modules/Loadable/Segmentations/Logic/vtkSlicerSegmentationGeometryLogic.h

@@ -57,8 +57,17 @@ class VTK_SLICER_SEGMENTATIONS_LOGIC_EXPORT vtkSlicerSegmentationGeometryLogic :
   /// Use this calculated permutation for updating spacing widget from geometry and interpreting spacing input
   void ComputeSourceAxisIndexForInputAxis();
 
-  /// Resample labelmap representation the input segmentation node according to the current
+  /// Sets the computed reference image geometry into the input segmentation node
+  /// (in the ReferenceImageGeometry conversion parameter and node reference).
+  /// For example, reference geometry may be used for determining the geometry of the segmentation
+  /// when saving into file, to make the extents of the segmentation match the extents of the
+  /// segmented image (instead of making the segmentation just big enough to store all the segments).
+  bool SetReferenceImageGeometryInSegmentationNode();
+
+  /// Resample labelmap representation of the input segmentation node according to the current
   /// output geometry setting.
+  ///
+  /// \sa SetReferenceImageGeometryInSegmentationNode()
   bool ResampleLabelmapsInSegmentationNode();
 
   /// Oriented image data containing output geometry. This is what the class calculates from the inputs

Modules/Loadable/Segmentations/Testing/Python/SegmentationsModuleTest1.py

@@ -34,6 +34,7 @@ def test_SegmentationsModuleTest1(self):
         self.TestSection_ImportExportSegment()
         self.TestSection_ImportExportSegment2()
         self.TestSection_SubjectHierarchy()
+        self.TestSection_SegmentGeometryLogic()
 
         logging.info("Test finished")
 
@@ -554,3 +555,69 @@ def TestSection_SubjectHierarchy(self):
         slicer.mrmlScene.RemoveNode(self.inputSegmentationNode)
         self.assertEqual(shNode.GetItemName(segmentationShItemID), "")
         self.assertEqual(shNode.GetItemName(sphereItemID), "")
+
+    # ------------------------------------------------------------------------------
+    def TestSection_SegmentGeometryLogic(self):
+        # Test segment geometry logic
+        logging.info("Test section: Segment geometry logic")
+
+        # Use MRHead for testing
+        import SampleData
+        volumeNode = SampleData.downloadSample("MRHead")
+
+        # Convert MRHead to oriented image data
+        import vtkSlicerSegmentationsModuleLogicPython as vtkSlicerSegmentationsModuleLogic
+        orientedImageData = vtkSlicerSegmentationsModuleLogic.vtkSlicerSegmentationsModuleLogic.CreateOrientedImageDataFromVolumeNode(volumeNode)
+        orientedImageData.UnRegister(None)
+
+        # Create segmentation node with binary labelmap master and one segment with MRHead geometry
+        segmentationNode = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLSegmentationNode")
+        segmentationNode.GetSegmentation().SetSourceRepresentationName(self.binaryLabelmapReprName)
+        geometryStr = slicer.vtkSegmentationConverter.SerializeImageGeometry(orientedImageData)
+        segmentationNode.GetSegmentation().SetConversionParameter(
+            slicer.vtkSegmentationConverter.GetReferenceImageGeometryParameterName(), geometryStr)
+
+        threshold = vtk.vtkImageThreshold()
+        threshold.SetInputData(orientedImageData)
+        threshold.ThresholdByUpper(16.0)
+        threshold.SetInValue(1)
+        threshold.SetOutValue(0)
+        threshold.SetOutputScalarType(vtk.VTK_UNSIGNED_CHAR)
+        threshold.Update()
+        segmentOrientedImageData = slicer.vtkOrientedImageData()
+        segmentOrientedImageData.DeepCopy(threshold.GetOutput())
+        mrImageToWorldMatrix = vtk.vtkMatrix4x4()
+        orientedImageData.GetImageToWorldMatrix(mrImageToWorldMatrix)
+        segmentOrientedImageData.SetImageToWorldMatrix(mrImageToWorldMatrix)
+        segment = slicer.vtkSegment()
+        segment.SetName("Brain")
+        segment.SetColor(0.0, 0.0, 1.0)
+        segment.AddRepresentation(self.binaryLabelmapReprName, segmentOrientedImageData)
+        segmentationNode.GetSegmentation().AddSegment(segment)
+
+        # Check resolution before resampling
+        geometryImageData = slicer.vtkOrientedImageData()
+        spacingComparisonDecimal = 3
+        segmentationNode.GetBinaryLabelmapRepresentation("Brain", geometryImageData)
+        actualSpacing = geometryImageData.GetSpacing()
+        expectedSpacing = [1, 1, 1.3]
+        import numpy as np
+        for i in range(3):
+            np.testing.assert_almost_equal(actualSpacing[i], expectedSpacing[i], decimal=3, err_msg=f"Spacing mismatch on axis {i} before resampling")
+
+        # Update geometry of internal binary labelmap representation in segmentation node
+        segmentationGeometryLogic = slicer.vtkSlicerSegmentationGeometryLogic()
+        segmentationGeometryLogic.SetInputSegmentationNode(segmentationNode)
+        segmentationGeometryLogic.SetSourceGeometryNode(volumeNode)
+        segmentationGeometryLogic.SetOversamplingFactor(2.0)
+        segmentationGeometryLogic.SetIsotropicSpacing(True)
+        segmentationGeometryLogic.CalculateOutputGeometry()
+        segmentationGeometryLogic.SetReferenceImageGeometryInSegmentationNode()
+        segmentationGeometryLogic.ResampleLabelmapsInSegmentationNode()
+
+        # Check resolution after resampling
+        segmentationNode.GetBinaryLabelmapRepresentation("Brain", geometryImageData)
+        actualSpacing = geometryImageData.GetSpacing()
+        expectedSpacing = [0.5, 0.5, 0.5]
+        for i in range(3):
+            np.testing.assert_almost_equal(actualSpacing[i], expectedSpacing[i], decimal=3, err_msg=f"Spacing mismatch on axis {i} after resampling")

Modules/Loadable/Segmentations/Widgets/qMRMLSegmentationGeometryWidget.cxx

@@ -479,37 +479,16 @@ void qMRMLSegmentationGeometryWidget::setReferenceImageGeometryForSegmentationNo
   Q_D(qMRMLSegmentationGeometryWidget);
   if (!d->SegmentationNode.GetPointer())
   {
-    qCritical() << Q_FUNC_INFO << "No input segmentation specified";
+    qCritical() << Q_FUNC_INFO << ": No input segmentation specified";
     return;
   }
 
-  // Save reference geometry
+  d->Logic->SetReferenceImageGeometryInSegmentationNode();
+
   vtkOrientedImageData* geometryImageData = d->Logic->GetOutputGeometryImageData();
   std::string geometryString = vtkSegmentationConverter::SerializeImageGeometry(geometryImageData);
-  d->SegmentationNode->GetSegmentation()->SetConversionParameter(
-    vtkSegmentationConverter::GetReferenceImageGeometryParameterName(), geometryString );
-
-  // Save reference geometry node (this is shown in Segmentations module
-  // to gives a hint about which node the current geometry is based on)
-  const char* referenceGeometryNodeID = nullptr;
-  if (d->Logic->GetSourceGeometryNode())
-  {
-    referenceGeometryNodeID = d->Logic->GetSourceGeometryNode()->GetID();
-  }
-
-  // If the reference geometry node is the same as the segmentation node, then don't change the node reference
-  if (vtkMRMLSegmentationNode::SafeDownCast(d->Logic->GetSourceGeometryNode()) != d->SegmentationNode)
-  {
-    d->SegmentationNode->SetNodeReferenceID(
-      vtkMRMLSegmentationNode::GetReferenceImageGeometryReferenceRole().c_str(), referenceGeometryNodeID);
-  }
-
-  // Note: it could be also useful to save oversampling value and isotropic flag,
-  // we could then allow the user to modify these settings instead of always
-  // setting from scratch.
-
-  qDebug() << Q_FUNC_INFO << "Reference image geometry of " << d->SegmentationNode->GetName()
-    << " has been set to '" << geometryString.c_str() << "'";
+  qDebug() << Q_FUNC_INFO << ": Reference image geometry of" << d->SegmentationNode->GetName()
+    << "is set to '" << geometryString.c_str() << "'";
 }
 
 //-----------------------------------------------------------------------------