Add VF option to specify number of transforms returned

src/openlifu/seg/skinseg.py

@@ -443,11 +443,10 @@ def _spherical_interpolator_from_mesh_embree(surface_mesh : vtk.vtkPolyData) ->
     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:
+    connectivity_np = vtk_to_numpy(polys.GetConnectivityArray())
+    if connectivity_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)
+    faces_np = connectivity_np.reshape(-1, 3).astype(np.int64) # (M, 3)
 
     r_squared = np.sum(points_np**2, axis=1)
 

src/openlifu/virtual_fit.py

@@ -90,6 +90,9 @@ class VirtualFitOptions(DictMixin):
     planefit_dpitch_step: Annotated[float, OpenLIFUFieldData("Plane fit pitch step", "Local pitch axis step size to use when constructing plane fitting grids. In spatial units of `units`")] = 3
     """Local pitch axis step size to use when constructing plane fitting grids. In spatial units of `units`."""
 
+    top_n_candidates: Annotated[int, OpenLIFUFieldData("No. of candidates returned", "Sets the limit for the number of transducer transform candidates returned by the algorithm.")] = 4
+    """Sets the limit for the number of transducer transform candidates returned by the algorithm."""
+
     def __post_init__(self):
         if not isinstance(self.units, str):
             raise TypeError("Units must be a string")
@@ -135,6 +138,10 @@ def __post_init__(self):
             raise ValueError("Plane fit pitch extent must be greater than 0")
         if not isinstance(self.planefit_dpitch_step, int | float):
             raise TypeError("Plane fit pitch step must be a number")
+        if not isinstance(self.top_n_candidates, int ):
+            raise TypeError("Number of transducer transform candidates returned must be an integer")
+        if self.top_n_candidates <= 0:
+            raise ValueError("Number of transducer transform candidates returned must be greater than 0")
 
     def to_units(self, target_units: str) -> VirtualFitOptions:
         """Do unit conversion and return a version of this VirtualFitOptions that uses
@@ -152,6 +159,7 @@ def to_units(self, target_units: str) -> VirtualFitOptions:
             planefit_dyaw_step = conversion_factor * self.planefit_dyaw_step,
             planefit_dpitch_extent = conversion_factor * self.planefit_dpitch_extent,
             planefit_dpitch_step = conversion_factor * self.planefit_dpitch_step,
+            top_n_candidates = self.top_n_candidates,
         )
 
     @staticmethod
@@ -275,6 +283,7 @@ def progress_callback(progress_percent : int, step_description : str): # noqa: A
     planefit_dyaw_step = options.planefit_dyaw_step
     planefit_dpitch_extent = options.planefit_dpitch_extent
     planefit_dpitch_step = options.planefit_dpitch_step
+    top_n_candidates = options.top_n_candidates
 
 
     if skin_mesh is None:
@@ -443,7 +452,7 @@ def progress_callback(progress_percent : int, step_description : str): # noqa: A
 
         progress_callback(100, "Complete")
         return (
-            sorted_transforms,
+            sorted_transforms[:top_n_candidates],
             VirtualFitDebugInfo(
                 skin_mesh = skin_mesh,
                 spherically_interpolated_mesh = interpolator_mesh,
@@ -455,4 +464,4 @@ def progress_callback(progress_percent : int, step_description : str): # noqa: A
         )
 
     progress_callback(100, "Complete")
-    return sorted_transforms
+    return sorted_transforms[:top_n_candidates]