diff --git a/packages/scratch-core/src/mutations/__init__.py b/packages/scratch-core/src/mutations/__init__.py index e69de29b..1837a3ba 100644 --- a/packages/scratch-core/src/mutations/__init__.py +++ b/packages/scratch-core/src/mutations/__init__.py @@ -0,0 +1,3 @@ +from .filter import GausianRegressionFilter, LevelMap + +__all__ = ["GausianRegressionFilter", "LevelMap"] diff --git a/packages/scratch-core/src/mutations/filter.py b/packages/scratch-core/src/mutations/filter.py index c4e92235..06ea8733 100644 --- a/packages/scratch-core/src/mutations/filter.py +++ b/packages/scratch-core/src/mutations/filter.py @@ -1,5 +1,7 @@ +from collections import namedtuple +from itertools import product from typing import NamedTuple -from container_models.base import FloatArray1D, FloatArray2D +from container_models.base import DepthData, FloatArray1D, FloatArray2D from container_models.scan_image import ScanImage from conversion.leveling.data_types import SurfaceTerms from conversion.leveling.solver.design import build_design_matrix @@ -7,6 +9,14 @@ from conversion.leveling.solver.transforms import normalize_coordinates from mutations.base import ImageMutation import numpy as np +from utils.constants import RegressionOrder +from conversion.filter.regression import ( + _build_lhs_matrix, + _solve_pixelwise_regression, + apply_order0_filter, + convolve_2d_separable, + create_normalized_separable_kernels, +) class PointCloud(NamedTuple): @@ -118,3 +128,138 @@ def apply_on_image(self, scan_image: ScanImage) -> ScanImage: scan_image.data = leveled_map_2d return scan_image + + +class GausianRegressionFilter(ImageMutation): + # Constants based on ISO 16610 surface texture standards + # Standard Gaussian alpha for 50% transmission + ALPHA_GAUSSIAN = np.sqrt(np.log(2) / np.pi) + # Adjusted alpha often used for higher-order regression filters to maintain properties + # alpha = Sqrt((-1 - LambertW(-1, -1 / (2 * exp(1)))) / Pi) + ALPHA_REGRESSION = 0.7309134280946760 + _Exponent = namedtuple("Exponent", ["y", "x"]) + + def __init__( + self, cutoff_pixels: FloatArray1D, regression_order: RegressionOrder + ) -> None: + self.cutoff_pixels = cutoff_pixels + self.regression_order = regression_order + + def calculate_polynomial_filter( + self, + data: DepthData, + kernel_x: FloatArray1D, + kernel_y: FloatArray1D, + exponents: list[_Exponent], + ) -> DepthData: + """ + Apply Order-1 or Order-2 Local Polynomial Regression. + This function performs a Weighted Least Squares (WLS) fit of a polynomial surface within a local window + defined by the kernels. For each pixel, it solves the linear system A * c = b, where 'c' contains the + coefficients of the polynomial. The smoothed value is the first coefficient (c0). + The kernels (kernel_x, kernel_y) serve as spatial weight functions. They determine the importance of + neighboring pixels in the regression. A non-uniform kernel (e.g., Gaussian) ensures that points closer + to the target pixel have a higher influence on the fit than points at the window's edge, providing better + localization and noise suppression. + :param data: The 2D input array to be filtered. Can contain NaNs, which are treated as zero-weight during + the regression. + :param kernel_x: 1D array representing the horizontal weight distribution. + :param kernel_y: 1D array representing the vertical weight distribution. + :param exponents: List of (power_y, power_x) tuples defining the polynomial terms. + :returns: The filtered (smoothed) version of the input data. + """ + # 1. Setup Coordinate Systems (Normalized to [-1, 1] for stability) + ny, nx = len(kernel_y), len(kernel_x) + radius_y, radius_x = (ny - 1) // 2, (nx - 1) // 2 + + y_coords = np.arange(-radius_y, radius_y + 1) / (radius_y if ny > 1 else 1.0) + x_coords = np.arange(-radius_x, radius_x + 1) / (radius_x if nx > 1 else 1.0) + + # 2. Construct the Linear System Components (A matrix and b vector) + nan_mask = np.isnan(data) + weights = np.where(nan_mask, 0.0, 1.0) + weighted_data = np.where(nan_mask, 0.0, data * weights) + + # Calculate RHS vector 'b' (Data Moments) + # b_k = Convolution(weighted_data, x^px * y^py * Kernel) + rhs_moments = np.array( + [ + convolve_2d_separable( + weighted_data, (x_coords**px) * kernel_x, (y_coords**py) * kernel_y + ) + for py, px in exponents + ] + ) + + # Calculate LHS Matrix 'A' (Weight Moments) + # A_jk = Convolution(weights, x^(px_j + px_k) * y^(py_j + py_k) * Kernel) + lhs_matrix = _build_lhs_matrix( + weights, + kernel_x, + kernel_y, + x_coords, + y_coords, + exponents, # type: ignore + ) + + # 3. Solve the System (A * c = b) per pixel + return _solve_pixelwise_regression(lhs_matrix, rhs_moments, data) + + def generate_polynomial_exponents(self, order: int) -> list[_Exponent]: + """ + Generate polynomial exponent pairs for 2D polynomial terms up to a given order. + :param order: Maximum total degree (py + px) for the polynomial terms. + :returns: List of (power_y, power_x) tuples representing polynomial terms. + """ + return [ + self._Exponent(x, y) + for y, x in product(range(order + 1), repeat=2) + if y + x <= order + ] + + def apply_on_image(self, scan_image: ScanImage) -> ScanImage: + """ + Apply a 2D Savitzky-Golay filter with Gaussian weighting via local polynomial regression (ISO 16610-21). + This implementation generalizes standard Gaussian filtering to handle missing data (NaNs) using local + regression techniques. It supports 0th order (Gaussian Kernel weighted average), 1st order (planar fit), + and 2nd order (quadratic fit) regression. + Explanation of Regression Orders: + - **Order 0**: Equivalent to the Nadaraya-Watson estimator. It calculates a weighted average where weights + are determined by the Gaussian kernel and the validity (non-NaN status) of neighboring pixels. + - **Order 1 & 2**: Local Weighted Least Squares (LOESS). It fits a polynomial surface (plane or quadratic) to + the local neighborhood weighted by the Gaussian kernel. This acts as a robust 2D Savitzky-Golay filter. + Mathematical basis: + - Approximate a signal s(x, y) from noisy data f(x, y) = s(x, y) + e(x, y) using weighted local regression. + - The approximation b(x, y) is calculated as the fitted value at point (x, y) using a weighted least squares + approach. Weights are non-zero within the neighborhood [x - rx, x + rx] and [y - ry, y + ry], following a + Gaussian distribution with standard deviations proportional to rx and ry. + - Optimization: + For **Order 0**, the operation is mathematically equivalent to a normalized convolution. This implementation + uses FFT-based convolution for performance gains compared to pixel-wise regression. + :param data: 2D input array containing float data. May contain NaNs. + :param cutoff_pixels: The filter cutoff wavelength in pixels as array [cutoff_y, cutoff_x]. + :param regression_order: RegressionOrder enum specifying the polynomial fit order: + GAUSSIAN_WEIGHTED_AVERAGE (0) = Gaussian weighted average. + LOCAL_PLANAR (1) = Local planar fit (corrects for tilt). + LOCAL_QUADRATIC (2) = Local quadratic fit (corrects for quadratic curvature). + :returns: The filtered 2D array of the same shape as input. + """ + alpha = ( + self.ALPHA_REGRESSION + if self.regression_order == RegressionOrder.LOCAL_QUADRATIC + else self.ALPHA_GAUSSIAN + ) + kernel_x, kernel_y = create_normalized_separable_kernels( + alpha, self.cutoff_pixels + ) + + if self.regression_order == RegressionOrder.GAUSSIAN_WEIGHTED_AVERAGE: + scan_image.data = apply_order0_filter(scan_image.data, kernel_x, kernel_y) + return scan_image + scan_image.data = self.calculate_polynomial_filter( + scan_image.data, + kernel_x, + kernel_y, + exponents=self.generate_polynomial_exponents(self.regression_order.value), + ) + return scan_image diff --git a/packages/scratch-core/src/utils/constants.py b/packages/scratch-core/src/utils/constants.py new file mode 100644 index 00000000..1ded9267 --- /dev/null +++ b/packages/scratch-core/src/utils/constants.py @@ -0,0 +1,7 @@ +from enum import Enum, auto + + +class RegressionOrder(Enum): + GAUSSIAN_WEIGHTED_AVERAGE = auto() + LOCAL_PLANAR = auto() + LOCAL_QUADRATIC = auto()