Speed up virtual fitting

pyproject.toml

@@ -49,6 +49,7 @@ dependencies = [
   "OpenEXR",
   "scikit-image",
   "trimesh",
+  "embreex; platform_machine=='x86_64' or platform_machine=='AMD64'",
   "onnxruntime==1.18.0",
 ]
 

src/openlifu/seg/skinseg.py

@@ -18,13 +18,14 @@
 import numpy as np
 import skimage.filters
 import skimage.measure
+import trimesh
 import vtk
 from packaging.version import parse
 from scipy.interpolate import LinearNDInterpolator
 from scipy.ndimage import distance_transform_edt
-from vtk.util.numpy_support import numpy_to_vtk
+from vtk.util.numpy_support import numpy_to_vtk, vtk_to_numpy
 
-from openlifu.geo import cartesian_to_spherical
+from openlifu.geo import cartesian_to_spherical, spherical_to_cartesian_vectorized
 
 
 def apply_affine_to_polydata(affine:np.ndarray, polydata:vtk.vtkPolyData) -> vtk.vtkPolyData:
@@ -282,8 +283,9 @@ def spherical_interpolator_from_mesh(
     surface_mesh: vtk.vtkPolyData,
     origin: Tuple[float, float, float] = (0.,0.,0.),
     xyz_direction_columns: np.ndarray | None = None,
-    dist_tolerance: float = 0.0001
-) -> Callable[[float, float], float]:
+    use_embree: bool|None = None,
+    dist_tolerance: float = 0.0001,
+) -> Callable:
     """Create a spherical interpolator from a vtkPolyData.
 
     Here a "spherical interpolator" is a function that maps angles from a spherical coordinate system
@@ -300,14 +302,18 @@ def spherical_interpolator_from_mesh(
             of how the spherical angles relate to the x, y, and z axes. If not provided, the xyz_direction_columns will
             be an identity matrix, which means that the coordinates in which surface_mesh is given will directly be
             interpreted as the x,y,z upon which a spherical coordinate system will be based.
+        use_embree: Use an alternative algorithm that uses embree CPU raytracing. Defaults to True only if embree is available;
+            it requires x86 architecture.
         dist_tolerance: A vertex of the surface_mesh will only be included if it is the furthest point from the origin
             that is on the mesh along the ray emanating from the origin and passing through the vertex. The
             dist_tolerance is the threshold for determining whether an intersection of the ray with the mesh
-            counts as being a distinct further out point from the vertex.
+            counts as being a distinct further out point from the vertex. This parameter only matters if use_embree is off.
 
     Returns:
         A spherical interpolator, which is a callable that maps (theta,phi) pairs of spherical coordinates (phi being azimuthal)
         to r values (radial spherical coordinate values). The angles are in radians.
+        A spherical interpolator can also run in batch mode, operating on a numpy array of shape (...,2) consisting
+        of theta,phi pairs in the last axis.
 
     Summary of the algorithm:
         - Transform the input mesh based on the desired origin and orientation of the spherical coordinate system.
@@ -327,6 +333,9 @@ def spherical_interpolator_from_mesh(
     if xyz_direction_columns is None:
         xyz_direction_columns = np.eye(3, dtype=float)
 
+    if use_embree is None:
+        use_embree = trimesh.ray.has_embree
+
     xyz_affine = np.eye(4)
     xyz_affine[:3,:3] = xyz_direction_columns
     xyz_affine[:3,3] = origin
@@ -340,20 +349,29 @@ def spherical_interpolator_from_mesh(
     transform_filter = vtk.vtkTransformPolyDataFilter()
     transform_filter.SetTransform(xyz_inverse_transform)
     transform_filter.SetInputData(surface_mesh)
-    transform_filter.Update()
-    surface_mesh_transformed = transform_filter.GetOutput()
+    triangle_filter = vtk.vtkTriangleFilter()
+    triangle_filter.SetInputConnection(transform_filter.GetOutputPort())
+    triangle_filter.Update()
+    surface_mesh_transformed = triangle_filter.GetOutput()
+
+    if use_embree:
+        return _spherical_interpolator_from_mesh_embree(surface_mesh_transformed)
+    else:
+        return _spherical_interpolator_from_mesh_cell_locator(surface_mesh_transformed, dist_tolerance)
+
+def _spherical_interpolator_from_mesh_cell_locator(surface_mesh : vtk.vtkPolyData, dist_tolerance:float) -> Callable:
 
     spherical_coords_on_mesh : List[Tuple[float,float,float]] = []
 
-    points = surface_mesh_transformed.GetPoints()
+    points = surface_mesh.GetPoints()
 
     # The farthest point from the origin is this far out:
     r_max = np.max([np.sqrt(np.sum(np.array(points.GetPoint(i))**2)) for i in range(points.GetNumberOfPoints())])
 
     sqdist_tolerance = dist_tolerance**2
 
     locator = vtk.vtkCellLocator() # Tried vtkOBBTree and it seems vtkCellLocator is much faster for this application
-    locator.SetDataSet(surface_mesh_transformed)
+    locator.SetDataSet(surface_mesh)
     locator.BuildLocator()
 
     for i in range(points.GetNumberOfPoints()):
@@ -420,3 +438,45 @@ def spherical_interpolator_from_mesh(
     )
 
     return interpolator
+
+def _spherical_interpolator_from_mesh_embree(surface_mesh : vtk.vtkPolyData) -> Callable:
+    vtk_points = surface_mesh.GetPoints()
+    points_np = vtk_to_numpy(vtk_points.GetData()).astype(np.float64) # (N,3)
+    polys = surface_mesh.GetPolys()
+    polys_np = vtk_to_numpy(polys.GetData()) # flat array [3,i0,i1,i2,3,i0,i1,i2,...]
+    if polys_np.size == 0:
+        raise RuntimeError("Input mesh has no polygons after transformation/triangulation.")
+    polys_np = polys_np.reshape(-1, 4) # (M, 4)
+    faces_np = polys_np[:, 1:4].astype(np.int64) # (M, 3)
+
+    r_squared = np.sum(points_np**2, axis=1)
+
+    # The farthest point from the origin is this far out:
+    r_max = float(np.sqrt(r_squared.max()))
+
+    tm = trimesh.Trimesh(vertices=points_np, faces=faces_np, process=False)
+    intersector = trimesh.ray.ray_pyembree.RayMeshIntersector(tm)
+
+    def interpolator(*args):
+        if len(args)==2:
+            arr = np.array(args)
+        elif len(args)==1 and isinstance(args[0], np.ndarray):
+            arr = args[0] # expected shape (...,2)
+            if arr.shape[-1] != 2:
+                msg = f"Interpolator expects array of shape (...,2). Got shape {arr.shape}"
+                raise ValueError(msg)
+        else:
+            raise ValueError("Interpolator expects either two args (theta, phi) or a single numpy array arg shaped (...,2)")
+
+        origins = spherical_to_cartesian_vectorized(
+            np.concatenate([np.full(arr.shape[:-1] + (1,), r_max+1),arr], axis=-1) # add r coordinate, giving shape (...,3)
+        )
+
+        # intersects_id will expect shape (N,3), but we want to support (...,3), so reshape if needed:
+        batch_shape = origins.shape[:-1]
+        origins = origins.reshape((-1,3))
+
+        _, _, hit_locations = intersector.intersects_id(origins, -origins, multiple_hits=False, return_locations=True)
+        return np.linalg.norm(hit_locations, axis=-1).reshape(batch_shape)
+
+    return interpolator

src/openlifu/virtual_fit.py

@@ -329,17 +329,19 @@ def progress_callback(progress_percent : int, step_description : str): # noqa: A
     points_asl = np.zeros((num_search_points,3), dtype=float) # search grid points in ASL coordinates
     normals_asl = np.zeros((num_search_points,3), dtype=float) # normal vector of the plane that is fitted at each point, in ASL coordinates
 
+    # Pre-build plane fitting grid that will be used repeatedly in the search loop below.
+    # This grid lives in the spherical coordinate basis theta-phi plane at each point.
+    # Below we will project it onto the skin surface for each search point.
+    dtheta_sequence = np.arange(-planefit_dyaw_extent, planefit_dyaw_extent + planefit_dyaw_step, planefit_dyaw_step)
+    dphi_sequence = np.arange(-planefit_dpitch_extent, planefit_dpitch_extent + planefit_dpitch_step, planefit_dpitch_step)
+    dtheta_grid, dphi_grid = np.meshgrid(dtheta_sequence, dphi_sequence, indexing='ij')
+
     for i in range(num_search_points):
         theta_rad, phi_rad = thetas[i]*np.pi/180, phis[i]*np.pi/180
 
         # Cartesian coordinate location of the point at which we are fitting a plane
         point = np.array(spherical_to_cartesian(skin_interpolator(theta_rad, phi_rad), theta_rad, phi_rad))
 
-        # Build plane fitting grid in the spherical coordinate basis theta-phi plane, which we will later project back onto the skin surface
-        dtheta_sequence = np.arange(-planefit_dyaw_extent, planefit_dyaw_extent + planefit_dyaw_step, planefit_dyaw_step)
-        dphi_sequence = np.arange(-planefit_dpitch_extent, planefit_dpitch_extent + planefit_dpitch_step, planefit_dpitch_step)
-        dtheta_grid, dphi_grid = np.meshgrid(dtheta_sequence, dphi_sequence, indexing='ij')
-
         r_hat, theta_hat, phi_hat = spherical_coordinate_basis(theta_rad,phi_rad)
         planefit_points_unprojected_cartesian = (
             point.reshape((1,1,3))
@@ -350,7 +352,7 @@ def progress_callback(progress_percent : int, step_description : str): # noqa: A
         planefit_points_unprojected_spherical = cartesian_to_spherical_vectorized(
             planefit_points_unprojected_cartesian
         ) # shape (num dthetas, num dphis, 3)
-        skin_projected_r_values = skin_interpolator(planefit_points_unprojected_spherical[...,1:]) # shape (num dthetas, num dphis) # TODO adjust docstrings to demand a *vectorizable* spherical interpolator
+        skin_projected_r_values = skin_interpolator(planefit_points_unprojected_spherical[...,1:]) # shape (num dthetas, num dphis)
         planefit_points_cartesian = spherical_to_cartesian_vectorized( # Could instead renormalize planefit_points_unprojected_cartesian, not sure if it would give a speedup versus this
             np.stack([
                 skin_projected_r_values, # New r values after projection to skin

tests/test_skinseg.py

@@ -4,6 +4,7 @@
 
 import numpy as np
 import pytest
+import trimesh
 import vtk
 from scipy.linalg import expm
 from scipy.stats import skew
@@ -111,7 +112,16 @@ def test_create_closed_surface_from_labelmap():
     # isn't supposed to add a colormap or anything like that
     assert surface.GetPointData().GetScalars() is None
 
-def test_spherical_interpolator_from_mesh():
+@pytest.mark.parametrize(
+        "use_embree",
+        [
+            False,
+            pytest.param(
+                True,
+                marks = pytest.mark.skipif(not trimesh.ray.has_embree, reason="Embree not available (needs x86 architecture)")
+            )
+        ])
+def test_spherical_interpolator_from_mesh(use_embree):
     """Check using a torus that the spherical interpolator behaves reasonably"""
     parametric_torus = vtk.vtkParametricTorus()
     parametric_torus.SetRingRadius(12.)
@@ -131,6 +141,7 @@ def test_spherical_interpolator_from_mesh():
         surface_mesh = torus_polydata,
         origin = origin,
         xyz_direction_columns = xyz_direction_columns,
+        use_embree=use_embree
     )
 
     sphere_source = vtk.vtkSphereSource()