CMC Pipeline: add cell coarse registration

packages/scratch-core/src/conversion/surface_comparison/cell_registration.py

@@ -1,52 +0,0 @@
-from collections.abc import Iterable
-
-from container_models.scan_image import ScanImage
-from conversion.surface_comparison.models import (
-    ComparisonParams,
-    Cell,
-    CellMetaData,
-    GridCell,
-    ProcessedMark,
-)
-
-from conversion.surface_comparison.utils import convert_pixels_to_meters
-
-
-def coarse_registration(
-    grid_cells: Iterable[GridCell],
-    comparison_image: ScanImage,
-    params: ComparisonParams,
-    fill_value_reference: float,  # Fill value for NaNs in the grid cell data
-) -> list[Cell]:
-    """TODO: Implement function."""
-
-    # Generate dummy output
-    pixel_size = comparison_image.scale_x  # Assumes isotropic image
-    output = []
-    for grid_cell in grid_cells:
-        dummy = Cell(
-            center_reference=convert_pixels_to_meters(
-                values=grid_cell.center, pixel_size=pixel_size
-            ),
-            cell_data=grid_cell.cell_data,
-            fill_fraction_reference=grid_cell.fill_fraction,
-            best_score=0.0,
-            angle_deg=0.0,
-            center_comparison=convert_pixels_to_meters(
-                values=(0, 0), pixel_size=pixel_size
-            ),
-            is_congruent=False,  # TODO: We shouldn't set this here
-            meta_data=CellMetaData(
-                is_outlier=False, residual_angle_deg=0.0, position_error=(0, 0)
-            ),  # TODO: We shouldn't set this here
-        )
-        output.append(dummy)
-
-    return output
-
-
-def fine_registration(
-    comparison_mark: ProcessedMark, cells: Iterable[Cell]
-) -> list[Cell]:
-    """TODO: Implement function."""
-    return list(cells)

packages/scratch-core/src/conversion/surface_comparison/cell_registration/coarse.py

@@ -0,0 +1,198 @@
+from container_models.base import BinaryMask, FloatArray2D
+from container_models.scan_image import ScanImage
+from conversion.surface_comparison.models import (
+    ComparisonParams,
+    Cell,
+    GridCell,
+)
+import numpy as np
+from skimage.transform import rotate
+from conversion.surface_comparison.cell_registration.utils import (
+    convert_grid_cell_to_cell,
+    pad_image_array,
+)
+
+import cv2
+
+from conversion.surface_comparison.utils import rotate_points
+
+
+def match_cells(
+    grid_cells: list[GridCell], comparison_image: ScanImage, params: ComparisonParams
+) -> list[Cell]:
+    """
+    Find the best-matching position and angle for each grid cell in the comparison image.
+
+    For each angle in the configured sweep, the padded comparison image is rotated and a normalized
+    cross-correlation score map is computed per cell using ``cv2.TM_CCOEFF_NORMED``. Positions whose
+    comparison-patch fill fraction falls below ``params.minimum_fill_fraction`` are masked out.
+    Per rotation angle, the highest score with its corresponding translation is stored.
+    The rotation that yields the highest unmasked score will be stored in each cell's :class:`GridSearchParams`.
+
+    The comparison image is padded by a full cell in each direction before the search so that cells whose
+    reference top-left lies near the image boundary can still be matched. The padding offset is subtracted
+    back when the best position is recorded, so all stored coordinates are in the original (unpadded) pixel space.
+
+    :param grid_cells: Reference grid cells to register; all cells must have the same size.
+    :param comparison_image: Comparison scan image to search over.
+    :param params: Algorithm parameters (angle sweep bounds, step, fill-fraction threshold).
+    :returns: List of :class:`Cell` objects with the best registration result per grid cell.
+    """
+    if not grid_cells:
+        return []
+
+    fill_value_comparison = float(np.nanmean(comparison_image.data))
+    pixel_size = comparison_image.scale_x  # Assumes isotropic image
+    cell_width, cell_height = grid_cells[0].width, grid_cells[0].height
+    pad_width, pad_height = cell_width, cell_height  # Set pad size to cell size
+
+    comparison_data = pad_image_array(
+        comparison_image.data, pad_width=pad_width, pad_height=pad_height
+    )
+    angles = np.arange(
+        params.search_angle_min,
+        params.search_angle_max + params.search_angle_step,
+        params.search_angle_step,
+    )
+
+    for angle in angles:
+        angle = float(angle)
+        # Rotate the comparison image by `-angle` degrees.
+        # This is equivalent to rotating the reference patch by `angle` degrees.
+        rotated = rotate(
+            image=comparison_data,
+            angle=-angle,
+            cval=np.nan,  # type: ignore
+            order=0,
+            resize=False,
+        )
+        # Get the mask of valid pixels for the rotated image
+        valid_mask = ~np.isnan(rotated)
+        # Compute the fill fraction mask based on the valid pixels mask
+        fill_fraction_map = _get_fill_fraction_map(
+            valid_pixel_mask=valid_mask,
+            cell_width=cell_width,
+            cell_height=cell_height,
+        )
+        fill_fraction_mask = fill_fraction_map >= params.minimum_fill_fraction
+        # Now that we computed the fill fraction mask, we can safely replace NaN values in the rotated image
+        rotated[~valid_mask] = fill_value_comparison
+        global_center_x, global_center_y = rotated.shape[1] / 2, rotated.shape[0] / 2  # type: ignore
+
+        for grid_cell in grid_cells:
+            score_map = _get_score_map(
+                comparison_array=rotated,
+                template=grid_cell.cell_data_filled,
+            )
+            score, x, y = _compute_best_score_from_maps(
+                score_map=score_map, fill_fraction_mask=fill_fraction_mask
+            )
+            if score > grid_cell.grid_search_params.score:
+                # Compute the center coordinates of the cell on the (original) unrotated image
+                cell_center = (x + cell_width / 2, y + cell_height / 2)
+                original_center_x, original_center_y = _unrotate_point(
+                    rotated_point=cell_center,
+                    angle=angle,
+                    pad_size=(pad_width, pad_height),
+                    rotation_center=(global_center_x, global_center_y),
+                )
+                grid_cell.grid_search_params.update(
+                    score=score,
+                    angle=angle,
+                    center_x=original_center_x,
+                    center_y=original_center_y,
+                )
+
+    return [
+        convert_grid_cell_to_cell(grid_cell=grid_cell, pixel_size=pixel_size)
+        for grid_cell in grid_cells
+    ]
+
+
+def _get_fill_fraction_map(
+    valid_pixel_mask: BinaryMask,
+    cell_height: int,
+    cell_width: int,
+) -> FloatArray2D:
+    """
+    Compute a 2D map where each entry [y, x] is the fill fraction of a cell-sized window with its
+    **top-left corner** at pixel (x, y), matching the indexing convention of ``cv2.matchTemplate``.
+
+    :param valid_pixel_mask: Boolean array (H, W); True where image data is valid.
+    :param cell_height: Height of the cell window in pixels.
+    :param cell_width: Width of the cell window in pixels.
+    :returns: Float64 array (H, W) with fill fractions in [0, 1], top-left indexed.
+        Entries near the bottom-right boundary are underestimates and will be rejected by the fill-fraction gate.
+        Since the image is padded with NaNs before calling this function, this does not matter.
+    """
+    kernel = np.ones((cell_height, cell_width), dtype=np.float32) / (
+        cell_height * cell_width
+    )
+    filtered = cv2.filter2D(
+        valid_pixel_mask.astype(np.float32),
+        ddepth=-1,
+        kernel=kernel,
+        anchor=(0, 0),
+        borderType=cv2.BORDER_CONSTANT,
+    )
+    return np.asarray(filtered, dtype=np.float64)
+
+
+def _unrotate_point(
+    rotated_point: tuple[float, float],
+    angle: float,
+    pad_size: tuple[int, int],
+    rotation_center: tuple[float, float],
+) -> tuple[float, float]:
+    # Undo the -angle rotation that was applied to the padded comparison image
+    unrotated_x, unrotated_y = rotate_points(
+        points=np.array([rotated_point]),
+        center=rotation_center,
+        angle=np.radians(-angle),
+    )[0]
+    # Compute the original coordinates by removing the padding
+    original_x = unrotated_x - pad_size[0]
+    original_y = unrotated_y - pad_size[1]
+    return original_x, original_y
+
+
+def _compute_best_score_from_maps(
+    score_map: FloatArray2D, fill_fraction_mask: BinaryMask
+) -> tuple[float, int, int]:
+    """
+    Compute the highest correlation score and the corresponding x, y coordinates
+    from the score and fill fraction maps.
+    """
+    # Make sure the shape of `score_map` and the `fill_fraction_mask` match, and
+    # discard irrelevant fill fraction mask positions at the bottom right.
+    valid_positions = fill_fraction_mask[: score_map.shape[0], : score_map.shape[1]]
+    # Replace non-valid values (where fill fraction is below threshold) with -inf
+    masked_scores = np.where(valid_positions, score_map, -np.inf)
+    # Compute the best score and x, y position from the score map
+    best_flat_index = np.argmax(masked_scores)
+    score = masked_scores.flat[best_flat_index]
+    y, x = np.unravel_index(best_flat_index, masked_scores.shape)
+    return float(score), int(x), int(y)
+
+
+def _get_score_map(
+    comparison_array: FloatArray2D, template: FloatArray2D
+) -> FloatArray2D:
+    """
+    Compute a normalized cross-correlation score map for one reference cell.
+
+    Slides the cell template over the comparison array using ``cv2.TM_CCOEFF_NORMED``, which computes
+    the Pearson correlation coefficient between the template and every same-sized patch. NaN values must
+    have been replaced in both arrays before calling this function.
+
+    :param comparison_array: NaN-free float32-compatible comparison image, padded by a full cell on each side.
+    :param template: Reference grid cell whose ``cell_data`` is used as the template; must contain no NaN values.
+    :returns: Float64 score map of shape ``(H - cell_height + 1, W - cell_width + 1)`` with values in ``[-1, 1]``.
+    """
+
+    score_map = cv2.matchTemplate(
+        image=comparison_array.astype(np.float32),
+        templ=template.astype(np.float32),
+        method=cv2.TM_CCOEFF_NORMED,
+    )
+    return np.asarray(score_map, dtype=np.float64)

packages/scratch-core/src/conversion/surface_comparison/cell_registration/core.py

@@ -0,0 +1,35 @@
+from container_models.scan_image import ScanImage
+from conversion.surface_comparison.cell_registration.coarse import match_cells
+from conversion.surface_comparison.grid import GridCell
+from conversion.surface_comparison.models import (
+    ComparisonParams,
+    Cell,
+    ProcessedMark,
+)
+
+
+def coarse_registration(
+    grid_cells: list[GridCell],
+    comparison_image: ScanImage,
+    params: ComparisonParams,
+) -> list[Cell]:
+    """
+    Register each reference grid cell against the comparison image.
+
+    Computes the global mean of the reference image as a NaN fill value, then delegates to :func:`match_cells`
+    to find the best-matching position and angle for every cell via a coarse angle sweep.
+
+    :param grid_cells: Reference grid cells to register.
+    :param comparison_image: Comparison scan image to search over.
+    :param params: Algorithm parameters controlling the angle sweep and fill-fraction thresholds.
+    :returns: List of :class:`Cell` objects with the best registration result per grid cell.
+    """
+    matched_cells = match_cells(
+        grid_cells=grid_cells, comparison_image=comparison_image, params=params
+    )
+    return matched_cells
+
+
+def fine_registration(comparison_mark: ProcessedMark, cells: list[Cell]) -> list[Cell]:
+    """TODO: Implement this function."""
+    return cells

packages/scratch-core/src/conversion/surface_comparison/cell_registration/utils.py

@@ -0,0 +1,55 @@
+from container_models.base import FloatArray2D
+from conversion.surface_comparison.models import Cell, GridCell, CellMetaData
+
+import numpy as np
+
+from conversion.surface_comparison.utils import convert_pixels_to_meters
+
+
+def convert_grid_cell_to_cell(grid_cell: GridCell, pixel_size: float) -> Cell:
+    """Convert an instance of `GridCell` to an instance of `Cell`."""
+    cell = Cell(
+        center_reference=convert_pixels_to_meters(
+            values=grid_cell.center, pixel_size=pixel_size
+        ),
+        cell_data=grid_cell.cell_data,
+        # TODO: Do we actually need cell data here?
+        # TODO: Add the cell size in meters as well
+        fill_fraction_reference=grid_cell.fill_fraction,
+        best_score=grid_cell.grid_search_params.score,
+        angle_deg=grid_cell.grid_search_params.angle,
+        center_comparison=convert_pixels_to_meters(
+            values=(
+                grid_cell.grid_search_params.center_x,
+                grid_cell.grid_search_params.center_y,
+            ),
+            pixel_size=pixel_size,
+        ),
+        is_congruent=False,  # TODO: We shouldn't set this here?
+        meta_data=CellMetaData(
+            is_outlier=False, residual_angle_deg=0.0, position_error=(0, 0)
+        ),  # TODO: We shouldn't set this here?
+    )
+    return cell
+
+
+def pad_image_array(
+    array: FloatArray2D, pad_width: int, pad_height: int, fill_value: float = np.nan
+) -> FloatArray2D:
+    """
+    Pad a 2D array symmetrically with a constant fill value.
+
+    Adds ``pad_height`` rows above and below and ``pad_width`` columns to the left and right of the input array.
+    The original data is placed in the center of the output; the border is filled with ``fill_value``.
+
+    :param array: Input 2D array of shape ``(height, width)``.
+    :param pad_width: Number of columns to add on each side (left and right).
+    :param pad_height: Number of rows to add on each side (top and bottom).
+    :param fill_value: Constant value written into the padded border; defaults to NaN.
+    :returns: Padded array of shape ``(height + 2 * pad_height, width + 2 * pad_width)``, same dtype as input.
+    """
+    height, width = array.shape
+    new_shape = height + 2 * pad_height, width + 2 * pad_width
+    output = np.full(shape=new_shape, fill_value=fill_value, dtype=array.dtype)
+    output[pad_height : pad_height + height, pad_width : pad_width + width] = array
+    return output

packages/scratch-core/src/conversion/surface_comparison/cmc_classification.py

@@ -1,12 +1,13 @@
 import numpy as np
 from scipy.stats import t
 
-from container_models.base import Points2D, FloatArray1D, BoolArray1D
+from container_models.base import FloatArray1D, BoolArray1D
 from conversion.surface_comparison.models import (
     Cell,
     ComparisonResult,
     ComparisonParams,
 )
+from conversion.surface_comparison.utils import rotate_points
 
 
 def classify_congruent_cells(
@@ -89,23 +90,6 @@ def _wrap_angles(angles: FloatArray1D) -> FloatArray1D:
     return (angles + np.pi) % (2 * np.pi) - np.pi
 
 
-def _rotate_points(
-    points: Points2D, angle: float, center: tuple[float, float]
-) -> Points2D:
-    """
-    Rotate 2-D points around a center.
-
-    :param points: (N, 2) array of [x, y] coordinates.
-    :param angle: Rotation angle in radians.
-    :param center: Tuple for the center of rotation [x, y].
-    :returns: (N, 2) rotated points.
-    """
-    cos_val, sin_val = np.cos(angle), np.sin(angle)
-    rotation_matrix = np.array([[cos_val, -sin_val], [sin_val, cos_val]])
-    translation = np.array(center)
-    return (points - translation) @ rotation_matrix.T + translation
-
-
 def _get_esd_criterion(values: FloatArray1D) -> BoolArray1D:
     """
     Return a boolean inlier mask for ``values`` using the generalized ESD test.
@@ -201,7 +185,7 @@ def _get_consensus_translation(
     outliers = np.array([c.meta_data.is_outlier for c in cells])
     centers_reference[outliers] = np.nan
     centers_comparison[outliers] = np.nan
-    expected_positions_on_reference = _rotate_points(
+    expected_positions_on_reference = rotate_points(
         points=centers_reference, angle=angle, center=rotation_center
     )
     # Compute residuals with respect to comparison.