adding more nodes

backend/nodes/flatten_base.py

@@ -0,0 +1,67 @@
+"""Flatten base — level the flat base of a surface with raised features."""
+
+from __future__ import annotations
+
+import numpy as np
+from scipy.ndimage import median_filter
+
+from backend.node_registry import register_node
+from backend.data_types import DataField
+
+
+@register_node(display_name="Flatten Base")
+class FlattenBase:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "field": ("DATA_FIELD",),
+                "threshold_percentile": ("FLOAT", {"default": 30.0, "min": 5.0, "max": 80.0, "step": 1.0}),
+                "poly_degree": ("INT", {"default": 2, "min": 0, "max": 5}),
+            }
+        }
+
+    OUTPUTS = (
+        ('DATA_FIELD', 'leveled'),
+    )
+    FUNCTION = "process"
+
+    DESCRIPTION = (
+        "Level the flat base of a surface that has raised features (particles, "
+        "grains). Uses a height percentile threshold to identify base pixels, "
+        "fits a polynomial to those pixels, and subtracts it. Unlike plane level, "
+        "this ignores tall features that would bias the fit. "
+    )
+
+    def process(self, field: DataField, threshold_percentile: float, poly_degree: int) -> tuple:
+        data = np.asarray(field.data, dtype=np.float64)
+        yres, xres = data.shape
+
+        # Identify base pixels: those below the threshold percentile
+        threshold = np.percentile(data, threshold_percentile)
+        base_mask = data <= threshold
+
+        if base_mask.sum() < max(3, (poly_degree + 1) ** 2):
+            # Not enough base pixels, fall back to subtracting the mean
+            return (field.replace(data=data - data.mean()),)
+
+        yy, xx = np.mgrid[:yres, :xres]
+        x_norm = xx.ravel() / max(xres - 1, 1)
+        y_norm = yy.ravel() / max(yres - 1, 1)
+
+        # Build polynomial basis
+        cols = []
+        for py in range(poly_degree + 1):
+            for px in range(poly_degree + 1 - py):
+                cols.append(x_norm**px * y_norm**py)
+        A_full = np.column_stack(cols)
+
+        # Fit on base pixels only
+        base_indices = np.where(base_mask.ravel())[0]
+        A_base = A_full[base_indices]
+        z_base = data.ravel()[base_indices]
+        coeffs, _, _, _ = np.linalg.lstsq(A_base, z_base, rcond=None)
+
+        # Evaluate and subtract
+        background = (A_full @ coeffs).reshape(data.shape)
+        return (field.replace(data=data - background),)