Integrate virtual fitting algorithm (#139)

Author: ebrahimebrahim

OpenLIFULib/CMakeLists.txt

@@ -19,6 +19,7 @@ set(MODULE_PYTHON_SCRIPTS
   OpenLIFULib/photoscan.py
   OpenLIFULib/virtual_fit_results.py
   OpenLIFULib//transform_conversion.py
+  OpenLIFULib/skinseg.py
 )
 
 set(MODULE_PYTHON_RESOURCES

OpenLIFULib/OpenLIFULib/coordinate_system_utils.py

@@ -48,15 +48,15 @@ def linear_to_affine(matrix, translation=None):
         axis=0,
     )
 
-def get_RAS2IJK(volume_node: vtkMRMLScalarVolumeNode):
-    """Get the _world_ RAS to volume IJK affine matrix for a given volume node.
+def get_IJK2RAS(volume_node: vtkMRMLScalarVolumeNode):
+    """Get the trasnfrom from IJK to the _world_ RAS for a given volume node.
 
     This takes into account any transforms that the volume node may be subject to.
 
     Returns a numpy array of shape (4,4).
     """
     IJK_to_volumeRAS_vtk = vtk.vtkMatrix4x4()
-    volume_node.GetRASToIJKMatrix(IJK_to_volumeRAS_vtk)
+    volume_node.GetIJKToRASMatrix(IJK_to_volumeRAS_vtk)
     IJK_to_volumeRAS = slicer.util.arrayFromVTKMatrix(IJK_to_volumeRAS_vtk)
     if volume_node.GetParentTransformNode():
         volumeRAS_to_worldRAS_vtk = vtk.vtkMatrix4x4()

OpenLIFULib/OpenLIFULib/simulation.py

@@ -5,7 +5,7 @@
 from vtk.util import numpy_support
 import slicer
 from slicer import vtkMRMLScalarVolumeNode
-from OpenLIFULib.coordinate_system_utils import get_RAS2IJK
+from OpenLIFULib.coordinate_system_utils import get_IJK2RAS
 from OpenLIFULib.lazyimport import xarray_lz
 
 if TYPE_CHECKING:
@@ -60,7 +60,7 @@ def make_xarray_in_transducer_coords_from_volume(volume_node:vtkMRMLScalarVolume
     # IJK : the volume node's underlying data array indices
     ijk2xyz = np.concatenate([np.concatenate([np.diag(spacing),origin.reshape(3,1)], axis=1), np.array([0,0,0,1],dtype=origin.dtype).reshape(1,4)])
     xyz2ras = slicer.util.arrayFromTransformMatrix(transducer.transform_node)
-    ras2IJK = get_RAS2IJK(volume_node)
+    ras2IJK = np.linalg.inv(get_IJK2RAS(volume_node))
     ijk2IJK = ras2IJK @ xyz2ras @ ijk2xyz
     volume_resampled_array = affine_transform(
         slicer.util.arrayFromVolume(volume_node).transpose((2,1,0)), # the array indices come in KJI rather than IJK so we permute them

OpenLIFULib/OpenLIFULib/skinseg.py

@@ -0,0 +1,16 @@
+"""Skin segmentation tools useful mainly for debugging"""
+
+from OpenLIFULib.lazyimport import openlifu_lz
+from slicer import vtkMRMLScalarVolumeNode, vtkMRMLModelNode
+from OpenLIFULib.coordinate_system_utils import get_IJK2RAS
+import slicer
+
+def generate_skin_mesh(volume_node:vtkMRMLScalarVolumeNode) -> vtkMRMLModelNode:
+    volume_array = slicer.util.arrayFromVolume(volume_node).transpose((2,1,0)) # the array indices come in KJI rather than IJK so we permute them
+    volume_affine_RAS = get_IJK2RAS(volume_node)
+    foreground_mask_array = openlifu_lz().seg.skinseg.compute_foreground_mask(volume_array)
+    foreground_mask_vtk_image = openlifu_lz().seg.skinseg.vtk_img_from_array_and_affine(foreground_mask_array, volume_affine_RAS)
+    skin_mesh = openlifu_lz().seg.skinseg.create_closed_surface_from_labelmap(foreground_mask_vtk_image)
+    skin_node = slicer.mrmlScene.AddNewNodeByClass("vtkMRMLModelNode")
+    skin_node.SetAndObservePolyData(skin_mesh)
+    return skin_node

OpenLIFULib/OpenLIFULib/transducer.py

@@ -159,3 +159,11 @@ def observe_transform_modified(self, callback : "Callable[[SlicerOpenLIFUTransdu
 
     def stop_observing_transform_modified(self, tag:int) -> None:
         self.transform_node.RemoveObserver(tag)
+
+    def set_current_transform_to_match_transform_node(self, transform_node : vtkMRMLTransformNode) -> None:
+        """Set the matrix on the current transform node of this transducer to match the matrix of a given transform node.
+        (This is done by a copy not reference, so it's a one-time update -- the tranforms do not become linked in any way.)"""
+
+        transform_matrix = vtk.vtkMatrix4x4()
+        transform_node.GetMatrixTransformToParent(transform_matrix)
+        self.transform_node.SetMatrixTransformToParent(transform_matrix)

OpenLIFUPrePlanning/OpenLIFUPrePlanning.py

@@ -4,6 +4,7 @@
 
 import qt
 import vtk
+import numpy as np
 
 import slicer
 from slicer.i18n import tr as _
@@ -14,13 +15,14 @@
 from slicer import vtkMRMLMarkupsFiducialNode, vtkMRMLScalarVolumeNode, vtkMRMLTransformNode
 
 from OpenLIFULib import (
+    openlifu_lz,
     get_target_candidates,
     get_openlifu_data_parameter_node,
     OpenLIFUAlgorithmInputWidget,
     SlicerOpenLIFUProtocol,
     SlicerOpenLIFUTransducer,
 )
-from OpenLIFULib.util import replace_widget
+from OpenLIFULib.util import replace_widget, BusyCursor
 from OpenLIFULib.virtual_fit_results import (
     add_virtual_fit_result,
     clear_virtual_fit_results,
@@ -32,9 +34,13 @@
     get_target_id_from_virtual_fit_result_node,
 )
 from OpenLIFULib.targets import fiducial_to_openlifu_point_id
+from OpenLIFULib.coordinate_system_utils import get_IJK2RAS
+from OpenLIFULib.transform_conversion import transform_node_from_openlifu
 
 if TYPE_CHECKING:
     from OpenLIFUData.OpenLIFUData import OpenLIFUDataLogic
+    import openlifu
+    import openlifu.geo
 
 PLACE_INTERACTION_MODE_ENUM_VALUE = slicer.vtkMRMLInteractionNode().Place
 
@@ -458,19 +464,19 @@ def updateApprovalStatusLabel(self):
 
     def onVirtualfitClicked(self):
         activeData = self.algorithm_input_widget.get_current_data()
-        virtual_fit_result : Optional[vtkMRMLTransformNode] = self.logic.virtual_fit(
-            activeData["Protocol"],activeData["Transducer"], activeData["Volume"], activeData["Target"]
-        )
+        with BusyCursor():
+            virtual_fit_result : Optional[vtkMRMLTransformNode] = self.logic.virtual_fit(
+                activeData["Protocol"],activeData["Transducer"], activeData["Volume"], activeData["Target"]
+            )
 
-        if virtual_fit_result is None: # Temporary behavior!
-            # None indicates for now that the user activated the transform handles to do the placeholder manual virtual fitting.
-            # It will not be possible to get None once the algorithm is implemented, or at least it wouldn't mean the same thing.
-            target_id = fiducial_to_openlifu_point_id(activeData["Target"])
-            self.algorithm_input_widget.inputs_dict["Target"].disable_with_tooltip(f"VF for {target_id} in progress...") # Disable target selector during manual VF
+        if virtual_fit_result is None:
+            slicer.util.errorDisplay("Virtual fit failed. No viable transducer positions found.")
             return
-        self.algorithm_input_widget.update() # Re-enable target-selector now that manual VF is completed. Again, temporary behavior.
+        else:
+            # TODO: Make the virtual fit button both update the transducer transform and populate in the virtual fit results
+            activeData["Transducer"].set_current_transform_to_match_transform_node(virtual_fit_result)
+            self.watchVirtualFit(virtual_fit_result)
 
-        self.watchVirtualFit(virtual_fit_result)
         self.updateApproveButton()
         self.updateApprovalStatusLabel()
 
@@ -549,39 +555,52 @@ def virtual_fit(
             volume: vtkMRMLScalarVolumeNode,
             target: vtkMRMLMarkupsFiducialNode,
         ) -> Optional[vtkMRMLTransformNode]:
-        # Temporary measure of "manual" virtual fitting. See https://github.com/OpenwaterHealth/SlicerOpenLIFU/issues/153
-        transducer.transform_node.CreateDefaultDisplayNodes()
-        if not transducer.transform_node.GetDisplayNode().GetEditorVisibility():
-            slicer.util.infoDisplay(
-                text=(
-                    "The automatic virtual fitting algorithm is not yet implemented."
-                    " Use the interaction handles on the transducer to manually fit it."
-                    " You can click the Virtual fit button again to remove the interaction handles,"
-                    " completing the manual virtual fit and recording the virtual fit transform."
-                ),
-                windowTitle="Not implemented"
-            )
-            transducer.transform_node.GetDisplayNode().SetEditorVisibility(True)
-            return None # we would also return this in the event of failure to do virtual fitting
-        else:
-            # "Complete" the virtual fit
-            transducer.transform_node.GetDisplayNode().SetEditorVisibility(False)
 
-            session = get_openlifu_data_parameter_node().loaded_session
-            session_id : Optional[str] = session.get_session_id() if session is not None else None
+        # TODO: Many quantities are hard-coded here will not have to be when these two issues are done:
+        # https://github.com/OpenwaterHealth/OpenLIFU-python/issues/166
+        # https://github.com/OpenwaterHealth/OpenLIFU-python/issues/165
+        vf_transforms = openlifu_lz().virtual_fit(
+            standoff_transform = openlifu_lz().geo.create_standoff_transform(
+                z_offset = 13.55,
+                dzdy = 0.15
+            ),
+            volume_array = slicer.util.arrayFromVolume(volume),
+            volume_affine_RAS = get_IJK2RAS(volume),
+            target_RAS = target.GetNthControlPointPosition(0),
+            pitch_range = (-10,150),
+            pitch_step = 5,
+            yaw_range = (-65, 65),
+            yaw_step = 5,
+            transducer_steering_center_distance = 50,
+            steering_limits = (
+                (-50, 50), # lat
+                (-50, 50), # ele
+                (-50, 50), # ax
+            ),
+        )
+        # TODO: add log handler for this
 
-            target_id = fiducial_to_openlifu_point_id(target)
-            clear_virtual_fit_results(target_id=target_id,session_id=session_id)
+        session = get_openlifu_data_parameter_node().loaded_session
+        session_id : Optional[str] = session.get_session_id() if session is not None else None
 
-            # When actually running the real virtual fit algorithm, there will be more virtual fit results to add
-            # but we would only return the best one.
-            return add_virtual_fit_result(
-                transform_node = transducer.transform_node,
+        target_id = fiducial_to_openlifu_point_id(target)
+        clear_virtual_fit_results(target_id=target_id,session_id=session_id)
+
+        vf_result_nodes = []
+
+        for i,vf_transform in zip(range(10), vf_transforms): # We only add the top 10 virtual fit nodes, to not put so many transforms into the scene.
+            node = add_virtual_fit_result(
+                transform_node = transform_node_from_openlifu(vf_transform, transducer.transducer.transducer, "mm"),
                 target_id = target_id,
                 session_id = session_id,
                 approval_status = False,
-                clone_node=True,
+                clone_node=False,
+                rank = i+1,
             )
+            vf_result_nodes.append(node)
+        if len(vf_result_nodes)==0:
+            return None
+        return vf_result_nodes[0]
 
 #
 # OpenLIFUPrePlanningTest