VF debugging info and progress callbacks, and mesh recon progress callbacks

src/openlifu/nav/photoscan.py

@@ -289,13 +289,15 @@ def log_stream(stream, handler):
             "Texturing_1",
             "Publish_1" ]
 
-def run_reconstruction(images: list[Path],
-                       pipeline_name: str = "default_pipeline",
-                       input_resize_width: int = 3024,
-                       use_masks: bool = True,
-                       window_radius: int | None = None,
-                       return_durations: bool = False
-                       ) -> Tuple[Photoscan, Path] | Tuple[Photoscan, Path, Dict[str, float]]:
+def run_reconstruction(
+    images: list[Path],
+    pipeline_name: str = "default_pipeline",
+    input_resize_width: int = 3024,
+    use_masks: bool = True,
+    window_radius: int | None = None,
+    return_durations: bool = False,
+    progress_callback : Callable[[int,str],None] | None = None,
+) -> Tuple[Photoscan, Path] | Tuple[Photoscan, Path, Dict[str, float]]:
     """Run Meshroom with the given images and pipeline.
     Args:
         images (list[Path]): List of image file paths.
@@ -306,12 +308,21 @@ def run_reconstruction(images: list[Path],
         window_radius (Optional[int]): number of images forward and backward in the sequence to try and
             match with, if None match each images to all others.
         return_durations (bool): If True, also return a dictionary mapping node names to durations in seconds.
+        progress_callback: An optional function that will be called to report progress. The function should accept two arguments:
+            an integer progress value from 0 to 100 followed by a string message describing the step currently being worked on.
 
     Returns:
         Union[Tuple[Photoscan, Path], Tuple[Photoscan, Path, Dict[str, float]]]:
             - If return_durations is False: returns the Photoscan and the data directory path.
             - If return_durations is True: also returns a dictionary of node execution times.
     """
+
+    if progress_callback is None:
+        def progress_callback(progress_percent : int, step_description : str): # noqa: ARG001
+            pass # Define it to be a no-op if no progress_callback was provided.
+
+    progress_callback(0, "Starting reconstruction")
+
     durations = {}
 
     pipeline_dir = importlib.resources.files("openlifu.nav.meshroom_pipelines")
@@ -336,6 +347,7 @@ def run_reconstruction(images: list[Path],
     images_dir.mkdir(parents=True, exist_ok=True)
 
     #resize the images and store in tmp dir
+    progress_callback(0, "Resizing images")
     start_time = time.perf_counter()
     new_paths = []
     for image in images:
@@ -378,6 +390,7 @@ def run_reconstruction(images: list[Path],
     ]
 
     if use_masks:
+        progress_callback(3, "Generating image masks")
         start_time = time.perf_counter()
         masks_dir = temp_dir / "masks"
         masks_dir.mkdir(parents=True, exist_ok=True)
@@ -386,22 +399,32 @@ def run_reconstruction(images: list[Path],
         durations["MaskCreation"] = time.perf_counter() - start_time
 
     #run CameraInit node to set view ids
+    progress_callback(8, "Initializing camera IDs")
     command_camera_init = [*command.copy(), "--toNode", "CameraInit"]
     subprocess_stream_output(command_camera_init, logger_meshroom.info, logger_meshroom.warning)
 
     camera_init_path = next(cache_dir.glob("CameraInit/*/cameraInit.sfm"))
     write_pair_file(new_paths, camera_init_path, pair_file_path, window_radius=window_radius)
 
+    number_of_nodes = len([node for node in _nodes if node in config_nodes])
+    pipeline_progress_start = 10.0
+    pipeline_progress_end = 90.0
+    pipeline_progress_step = (pipeline_progress_end - pipeline_progress_start) / number_of_nodes
+
     #run each node
+    pipeline_progress = pipeline_progress_start
     for node in _nodes:
         if node not in config_nodes:
             continue
         command_i = [*command.copy(), "--toNode", node]
 
+        progress_callback(int(pipeline_progress),f"Meshroom: {node.strip('_1')}")
         start_time = time.perf_counter()
         subprocess_stream_output(command_i, logger_meshroom.info, logger_meshroom.warning)
         durations[node] = time.perf_counter() - start_time
+        pipeline_progress += pipeline_progress_step
 
+    progress_callback(pipeline_progress_end, "Merging texture maps")
 
     #merge texturemaps
     output_dir_merged = temp_dir / "output_dir_merged"
@@ -416,6 +439,7 @@ def run_reconstruction(images: list[Path],
     }
 
     photoscan = Photoscan.from_dict(photoscan_dict)
+    progress_callback(100, "Complete")
     if return_durations:
         return photoscan, output_dir_merged, durations
     return photoscan, output_dir_merged

src/openlifu/virtual_fit.py

@@ -2,17 +2,28 @@
 
 import logging
 from dataclasses import dataclass
-from typing import Annotated, Any, Dict, List, Sequence, Tuple
+from typing import (
+    Annotated,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Sequence,
+    Tuple,
+)
 
 import numpy as np
+import vtk
 
 from openlifu.geo import (
+    cartesian_to_spherical,
     cartesian_to_spherical_vectorized,
     spherical_coordinate_basis,
     spherical_to_cartesian,
     spherical_to_cartesian_vectorized,
 )
 from openlifu.seg.skinseg import (
+    apply_affine_to_polydata,
     compute_foreground_mask,
     create_closed_surface_from_labelmap,
     spherical_interpolator_from_mesh,
@@ -104,32 +115,108 @@ def from_dict(parameter_dict: Dict[str,Any]) -> VirtualFitOptions: # Override Di
         parameter_dict["steering_limits"] = tuple(map(tuple,parameter_dict["steering_limits"]))
         return VirtualFitOptions(**parameter_dict)
 
-
-def virtual_fit(
+def compute_skin_mesh_from_volume(
     volume_array : np.ndarray,
     volume_affine_RAS : np.ndarray,
+) -> vtk.vtkPolyData:
+    log.info("Computing foreground mask...")
+    foreground_mask_array = compute_foreground_mask(volume_array)
+    foreground_mask_vtk_image = vtk_img_from_array_and_affine(foreground_mask_array, volume_affine_RAS)
+    log.info("Creating closed surface from labelmap...")
+    skin_mesh = create_closed_surface_from_labelmap(foreground_mask_vtk_image)
+    return skin_mesh
+
+@dataclass
+class VirtualFitDebugInfo:
+    """Debugging information for the result of running `virtual_fit`."""
+    skin_mesh : vtk.vtkPolyData
+    """The skin mesh that was used for virtual fitting"""
+
+    spherically_interpolated_mesh : vtk.vtkPolyData
+    """A mesh representing the spherical interpolator that was used for virtual fitting"""
+
+    search_points : np.ndarray
+    """Array of shape (N,3) containing the coordinates of the points that were tried for virtual fitting"""
+
+    plane_normals : np.ndarray
+    """Array of shape (N,3) containing the normal vectors of the planes that were fitted at each of `search_points`"""
+
+    steering_dists : np.ndarray
+    """Array of shape (N,) containing the computed steering distance for each point in `search_points`"""
+
+    in_bounds : np.ndarray
+    """Boolean array of shape (N,) giving for each point in `search_points` whether the target was determined to be
+    in bounds for that candidate transducer placement."""
+
+
+def sphere_from_interpolator(
+        interpolator: Callable[[float, float], float],
+        theta_res:int = 50,
+        phi_res:int = 50,
+    ) -> vtk.vtkPolyData:
+    """Create a spherical mesh from a spherical interpolator, to help visualize how the interpolator works.
+    This is intended as a debugging utility."""
+    sphere_source = vtk.vtkSphereSource()
+    sphere_source.SetRadius(1.0)
+    sphere_source.SetThetaResolution(theta_res)
+    sphere_source.SetPhiResolution(phi_res)
+    sphere_source.Update()
+    sphere_polydata = sphere_source.GetOutput()
+    sphere_points = sphere_polydata.GetPoints()
+    for i in range(sphere_points.GetNumberOfPoints()):
+        point = np.array(sphere_points.GetPoint(i))
+        r, theta, phi = cartesian_to_spherical(*point)
+        r = interpolator(theta, phi)
+        sphere_points.SetPoint(i, r * point)
+    normals_filter = vtk.vtkPolyDataNormals()
+    normals_filter.SetInputData(sphere_polydata)
+    normals_filter.Update()
+    return normals_filter.GetOutput()
+
+def virtual_fit(
     units: str,
     target_RAS : Sequence[float],
     standoff_transform : np.ndarray,
     options : VirtualFitOptions,
-) -> List[np.ndarray]:
+    volume_array : np.ndarray | None = None,
+    volume_affine_RAS : np.ndarray | None = None,
+    skin_mesh : vtk.vtkPolyData | None = None,
+    include_debug_info : bool = False,
+    progress_callback : Callable[[int,str],None] | None = None,
+) -> List[np.ndarray] | Tuple[List[np.ndarray], VirtualFitDebugInfo]:
     """Run patient-specific "virtual fitting" algorithm, suggesting a series of candidate transducer
     transforms for optimal sonicaiton of a given target.
 
+    Provide either a `volume_array` and `volume_affine_RAS`, or a `skin_mesh`.
+
     Args:
-        volume_array: A 3D volume MRI
-        volume_affine_RAS: A 4x4 affine transform that maps `volume_array` into RAS space with certain units
         units: The spatial units of the RAS space into which volume_affine_RAS maps
         target_RAS: A 3D point, in the coordinates and units of `volume_affine_RAS` (the `units` argument)
         standoff_transform: See the documentation of `create_standoff_transform` or
             `Transducer.standoff_transform` for the meaning of this. Here it should be provided in the
             units `units`. The method `Transducer.get_standoff_transform_in_units` is useful for getting this.
         options : Virtual fitting algorithm configuration. See the `VirtualFitOptions` documentation.
+        volume_array: A 3D volume MRI
+        volume_affine_RAS: A 4x4 affine transform that maps `volume_array` into RAS space with certain units
+        skin_mesh: Optional pre-computed closed surface mesh. If provided, `volume_array` and
+            `volume_affine_RAS` can be omitted. The provided skin mesh should be in RAS space, with units
+            being the provided `units` arg. The function `compute_skin_mesh_from_volume` can be used to pre-compute
+            a skin mesh.
+        include_debug_info: Whether to include debugging info in the return value. Disabled by default because some of the debugging
+            info takes some time to compute.
+        progress_callback: An optional function that will be called to report progress. The function should accept two arguments:
+            an integer progress value from 0 to 100 followed by a string message describing the step currently being worked on.
 
     Returns: A list of transducer transform candidates sorted starting from the best-scoring one. The transforms map transducer space
         into LPS space, and they are in the same units as the RAS space of `volume_affine_RAS` (aka the `units` argument).
     """
 
+    if progress_callback is None:
+        def progress_callback(progress_percent : int, step_description : str): # noqa: ARG001
+            pass # Define it to be a no-op if no progress_callback was provided.
+
+    progress_callback(0, "Starting virtual fit")
+
     # Express all virtual fit options in the units of volume_affine_RAS, i.e. the physical space of the volume
     options = options.to_units(units)
     pitch_range = options.pitch_range
@@ -143,12 +230,20 @@ def virtual_fit(
     planefit_dpitch_extent = options.planefit_dpitch_extent
     planefit_dpitch_step = options.planefit_dpitch_step
 
-    log.info("Computing foreground mask...")
-    foreground_mask_array = compute_foreground_mask(volume_array)
-    foreground_mask_vtk_image = vtk_img_from_array_and_affine(foreground_mask_array, volume_affine_RAS)
-    log.info("Creating closed surface from labelmap...")
-    skin_mesh = create_closed_surface_from_labelmap(foreground_mask_vtk_image)
+
+    if skin_mesh is None:
+        if volume_array is None or volume_affine_RAS is None:
+            raise ValueError("Both `volume_array` and `volume_affine_RAS` must be provided if `skin_mesh` is None.")
+
+        log.info("Computing skin mesh...")
+        progress_callback(0, "Computing skin mesh")
+        skin_mesh = compute_skin_mesh_from_volume(volume_array, volume_affine_RAS)
+    else:
+        log.info("Using provided skin mesh.")
+
+
     log.info("Building skin interpolator...")
+    progress_callback(5, "Building skin interpolator")
     skin_interpolator = spherical_interpolator_from_mesh(
         surface_mesh = skin_mesh,
         origin = target_RAS,
@@ -168,6 +263,7 @@ def virtual_fit(
     steering_maxs = np.array([sl[1] for sl in steering_limits], dtype=float) # shape (3,). It is the lat,ele,ax steering max
 
     log.info("Searching through candidate transducer poses...")
+    progress_callback(50, "Searching through poses")
 
     # Construct search grid
     theta_sequence = np.arange(90 - yaw_range[-1], 90 - yaw_range[0], yaw_step)
@@ -178,10 +274,15 @@ def virtual_fit(
     thetas = theta_grid.reshape(num_search_points)
     phis = phi_grid.reshape(num_search_points)
 
+    # Things that will be computed over the search grid
     transducer_poses = np.empty((num_search_points,4,4), dtype=float)
     in_bounds = np.zeros(shape=num_search_points, dtype=bool)
     steering_dists = np.zeros(shape=num_search_points, dtype=float)
 
+    # Additional debugging info that will be computed over the search grid
+    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
+
     for i in range(num_search_points):
         theta_rad, phi_rad = thetas[i]*np.pi/180, phis[i]*np.pi/180
 
@@ -265,6 +366,8 @@ def virtual_fit(
         transducer_poses[i] = transducer_transform
         steering_dists[i] = steering_distance
         in_bounds[i] = target_in_bounds
+        points_asl[i] = point
+        normals_asl[i] = transducer_z
 
 
     sorted_transforms = [
@@ -273,4 +376,35 @@ def virtual_fit(
 
     log.info("Virtual fitting complete.")
 
+    if include_debug_info:
+        log.info("Generating debug meshes...")
+        progress_callback(80, "Generating debug meshes")
+        interpolator_mesh : vtk.vtkPolyData = sphere_from_interpolator(skin_interpolator, theta_res=100, phi_res=100)
+
+        # A few things are in ASL coordinates, so we transform it to RAS space so that they are in the same coordinates as skin_mesh.
+        interpolator_mesh = apply_affine_to_polydata(interpolator2ras, interpolator_mesh)
+        points_asl_homogenized = np.concatenate([points_asl.T, np.ones((1,num_search_points))], axis=0) # shape (4,num_search_points)
+        points_ras = (interpolator2ras @ points_asl_homogenized)[:3].T # back to shape (num_search_points,3)
+        normals_ras = (asl2ras_3x3 @ (normals_asl.T)).T
+
+        # After transforming the interpolator_mesh, the normals can end up flipped, so we fix it just in case
+        normals_filter = vtk.vtkPolyDataNormals()
+        normals_filter.SetInputData(interpolator_mesh)
+        normals_filter.Update()
+        interpolator_mesh = normals_filter.GetOutput()
+
+        progress_callback(100, "Complete")
+        return (
+            sorted_transforms,
+            VirtualFitDebugInfo(
+                skin_mesh = skin_mesh,
+                spherically_interpolated_mesh = interpolator_mesh,
+                search_points = points_ras,
+                plane_normals = normals_ras,
+                steering_dists = steering_dists,
+                in_bounds = in_bounds,
+            ),
+        )
+
+    progress_callback(100, "Complete")
     return sorted_transforms