Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions packages/scratch-core/src/mutations/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
from .filter import GausianRegressionFilter, LevelMap

__all__ = ["GausianRegressionFilter", "LevelMap"]
147 changes: 146 additions & 1 deletion packages/scratch-core/src/mutations/filter.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,22 @@
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
from conversion.leveling.solver.grid import get_2d_grid
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):
Expand Down Expand Up @@ -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(
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe you could actually remove these local functions and use the implementation in conversion instead, since you need to change the code anyway later

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:
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think only implementing this function should be sufficient for this PR

"""
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
7 changes: 7 additions & 0 deletions packages/scratch-core/src/utils/constants.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
from enum import Enum, auto


class RegressionOrder(Enum):
GAUSSIAN_WEIGHTED_AVERAGE = auto()
LOCAL_PLANAR = auto()
LOCAL_QUADRATIC = auto()
Loading