matei/tono
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

low pri features

Browse Source
This commit is contained in:
matei jordache
2026-04-04 00:25:53 -07:00
parent 4818c1123c
commit 5de93e6c4d
47 changed files with 3866 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
backend/node_menu.py
View File

@@ -35,6 +35,7 @@ MENU_LAYOUT: dict[str, list[str]] = {
"Save",
"SaveImage",
"Shade",
"PresentationOps",
],
"Overlay": [
"Markup",
@@ -56,6 +57,7 @@ MENU_LAYOUT: dict[str, list[str]] = {
"PixelBinning",
"ExtendPad",
"FieldArithmetic",
"DisplacementField",
],
"Level & Correct": [
"FixZero",
@@ -72,6 +74,8 @@ MENU_LAYOUT: dict[str, list[str]] = {
"ScanLineReorder",
"Tilt",
"WrapValue",
"DistributionCoercion",
"Calibration",
],
"Filter": [
"GaussianFilter",
@@ -98,6 +102,7 @@ MENU_LAYOUT: dict[str, list[str]] = {
"LogPolarPSDF",
"FrequencySplit",
"CrossCorrelate",
"SuperResolution",
],
"Measure": [
"CrossSection",
@@ -116,6 +121,7 @@ MENU_LAYOUT: dict[str, list[str]] = {
"MultipleProfiles",
"StraightenPath",
"RelateFields",
"DWTAnisotropy",
],
"Detect": [
"FeatureDetection",
@@ -130,6 +136,9 @@ MENU_LAYOUT: dict[str, list[str]] = {
"LateralForceSim",
"SEMSimulation",
"SMMAnalysis",
"PixelClassification",
"NeuralClassification",
"LogisticClassification",
],
"Mask": [
"DrawMask",
@@ -139,6 +148,9 @@ MENU_LAYOUT: dict[str, list[str]] = {
"MaskMorphology",
"MaskInvert",
"MaskOperations",
"MarkDisconnected",
"MaskShift",
"MaskNoisify",
],
"Grains": [
"GrainDistanceTransform",
@@ -150,11 +162,14 @@ MENU_LAYOUT: dict[str, list[str]] = {
"LevelGrains",
"GrainEdge",
"GrainCross",
"GrainVisualization",
],
"Tip": [
"TipModel",
"TipDeconvolution",
"BlindTipEstimate",
"TipShapeEstimate",
"PSFEstimation",
],
}

143
backend/nodes/calibration.py Normal file
View File

@@ -0,0 +1,143 @@
"""Calibration — apply lateral and height calibration corrections."""
from __future__ import annotations
import numpy as np
from backend.node_registry import register_node
from backend.data_types import DataField
@register_node(display_name="Calibration")
class Calibration:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"field": ("DATA_FIELD",),
"xy_mode": (["keep", "set_size", "scale"], {"default": "keep"}),
"z_mode": (["keep", "set_range", "scale", "offset"], {"default": "keep"}),
"xreal_new": ("FLOAT", {
"default": 1e-6,
"min": 1e-12,
"max": 1.0,
"step": 1e-9,
"show_when_widget_value": {"xy_mode": ["set_size"]},
}),
"yreal_new": ("FLOAT", {
"default": 1e-6,
"min": 1e-12,
"max": 1.0,
"step": 1e-9,
"show_when_widget_value": {"xy_mode": ["set_size"]},
}),
"xy_scale": ("FLOAT", {
"default": 1.0,
"min": 0.001,
"max": 1000.0,
"step": 0.001,
"show_when_widget_value": {"xy_mode": ["scale"]},
}),
"z_min": ("FLOAT", {
"default": 0.0,
"min": -1e-3,
"max": 1e-3,
"step": 1e-12,
"show_when_widget_value": {"z_mode": ["set_range"]},
}),
"z_max": ("FLOAT", {
"default": 1e-9,
"min": -1e-3,
"max": 1e-3,
"step": 1e-12,
"show_when_widget_value": {"z_mode": ["set_range"]},
}),
"z_scale": ("FLOAT", {
"default": 1.0,
"min": 0.001,
"max": 1000.0,
"step": 0.001,
"show_when_widget_value": {"z_mode": ["scale"]},
}),
"z_offset": ("FLOAT", {
"default": 0.0,
"min": -1e-3,
"max": 1e-3,
"step": 1e-12,
"show_when_widget_value": {"z_mode": ["offset"]},
}),
"xy_unit": ("STRING", {"default": ""}),
"z_unit": ("STRING", {"default": ""}),
}
}
OUTPUTS = (
('DATA_FIELD', 'result'),
)
FUNCTION = "process"
DESCRIPTION = (
"Apply lateral and height calibration corrections to a DATA_FIELD. "
"Lateral mode can set explicit physical size or scale uniformly. "
"Height mode can rescale to a target range, multiply by a factor, "
"or add a constant offset. Optionally override the XY or Z unit strings. "
"Equivalent to Gwyddion's calibrate functionality."
)
def process(
self,
field: DataField,
xy_mode: str,
z_mode: str,
xreal_new: float,
yreal_new: float,
xy_scale: float,
z_min: float,
z_max: float,
z_scale: float,
z_offset: float,
xy_unit: str,
z_unit: str,
) -> tuple:
data = np.asarray(field.data, dtype=np.float64).copy()
xreal = float(field.xreal)
yreal = float(field.yreal)
si_unit_xy = field.si_unit_xy
si_unit_z = field.si_unit_z
# --- lateral calibration ---
if xy_mode == "set_size":
xreal = float(xreal_new)
yreal = float(yreal_new)
elif xy_mode == "scale":
xreal *= float(xy_scale)
yreal *= float(xy_scale)
# "keep" → no change
# --- height calibration ---
if z_mode == "set_range":
cur_min = float(data.min())
cur_max = float(data.max())
if cur_max > cur_min:
data = float(z_min) + (data - cur_min) * (float(z_max) - float(z_min)) / (cur_max - cur_min)
else:
data[:] = float(z_min)
elif z_mode == "scale":
data *= float(z_scale)
elif z_mode == "offset":
data += float(z_offset)
# "keep" → no change
# --- unit overrides ---
if xy_unit:
si_unit_xy = xy_unit
if z_unit:
si_unit_z = z_unit
return (field.replace(
data=data,
xreal=xreal,
yreal=yreal,
si_unit_xy=si_unit_xy,
si_unit_z=si_unit_z,
),)

95
backend/nodes/displacement_field.py Normal file
View File

@@ -0,0 +1,95 @@
"""Displacement field — distort images using displacement fields."""
from __future__ import annotations
import numpy as np
from scipy.ndimage import gaussian_filter, gaussian_filter1d, map_coordinates
from backend.node_registry import register_node
from backend.data_types import DataField
@register_node(display_name="Displacement Field")
class DisplacementField:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"field": ("DATA_FIELD",),
"method": (["gaussian_1d", "gaussian_2d", "tear"], {"default": "gaussian_1d"}),
"sigma": ("FLOAT", {"default": 5.0, "min": 0.1, "max": 100.0, "step": 0.1}),
"tau": ("FLOAT", {"default": 20.0, "min": 1.0, "max": 500.0, "step": 1.0}),
"density": ("FLOAT", {
"default": 0.02,
"min": 0.001,
"max": 0.25,
"step": 0.001,
"show_when_widget_value": {"method": ["tear"]},
}),
"seed": ("INT", {"default": 42, "min": 0, "max": 999999}),
}
}
OUTPUTS = (
('DATA_FIELD', 'result'),
)
FUNCTION = "process"
DESCRIPTION = (
"Distort an image using synthetic displacement fields. "
"Supports 1D Gaussian (row-correlated), 2D Gaussian (fully correlated), "
"and tear (random horizontal tear lines) distortion modes. "
"Equivalent to Gwyddion's displfield.c module."
)
def process(
self,
field: DataField,
method: str,
sigma: float,
tau: float,
density: float,
seed: int,
) -> tuple:
data = np.asarray(field.data, dtype=np.float64)
yres, xres = data.shape
rng = np.random.default_rng(seed)
y_grid, x_grid = np.mgrid[:yres, :xres].astype(np.float64)
if method == "gaussian_1d":
dx_1d = gaussian_filter1d(rng.standard_normal(xres), tau) * sigma
dx = np.tile(dx_1d, (yres, 1))
dy = np.zeros_like(dx)
elif method == "gaussian_2d":
dx = gaussian_filter(rng.standard_normal((yres, xres)), tau) * sigma
dy = gaussian_filter(rng.standard_normal((yres, xres)), tau) * sigma
elif method == "tear":
dx = np.zeros((yres, xres), dtype=np.float64)
dy = np.zeros((yres, xres), dtype=np.float64)
# Select tear rows based on density
tear_mask = rng.random(yres) < density
tear_rows = np.nonzero(tear_mask)[0]
for row in tear_rows:
offset = rng.standard_normal() * sigma
# Apply offset that decays exponentially away from the tear line
for i in range(yres):
dist = abs(i - row)
dx[i] += offset * np.exp(-dist / max(tau, 1.0))
# Smooth the displacement to avoid sharp edges
for i in range(yres):
dx[i] = gaussian_filter1d(dx[i], tau / 4.0)
else:
raise ValueError(f"Unknown method: {method!r}")
coords_y = y_grid + dy
coords_x = x_grid + dx
result = map_coordinates(data, [coords_y, coords_x], mode='reflect', order=3)
return (field.replace(data=result),)

81
backend/nodes/distribution_coercion.py Normal file
View File

@@ -0,0 +1,81 @@
"""Distribution coercion — transform data to match a target distribution."""
from __future__ import annotations
import numpy as np
from math import ceil
from scipy.stats import norm
from backend.node_registry import register_node
from backend.data_types import DataField
def _coerce_block(data: np.ndarray, distribution: str, n_levels: int) -> np.ndarray:
"""Coerce a flat or 2-D block to the target distribution, returning same shape."""
shape = data.shape
flat = data.ravel().astype(np.float64)
n_pixels = flat.size
if n_pixels == 0:
return data.copy()
indices = np.argsort(flat, kind="mergesort")
if distribution == "uniform":
target = np.linspace(float(flat.min()), float(flat.max()), n_pixels)
elif distribution == "gaussian":
eps = 0.5 / n_pixels
quantiles = np.linspace(eps, 1.0 - eps, n_pixels)
target = norm.ppf(quantiles) * float(flat.std()) + float(flat.mean())
elif distribution == "levels":
n_levels = max(2, int(n_levels))
level_values = np.linspace(float(flat.min()), float(flat.max()), n_levels)
target = np.repeat(level_values, ceil(n_pixels / n_levels))[:n_pixels]
else:
raise ValueError(f"Unknown distribution: {distribution}")
result = np.empty_like(flat)
result[indices] = target
return result.reshape(shape)
@register_node(display_name="Distribution Coercion")
class DistributionCoercion:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"field": ("DATA_FIELD",),
"distribution": (["uniform", "gaussian", "levels"], {"default": "uniform"}),
"n_levels": ("INT", {
"default": 4,
"min": 2,
"max": 1000,
"show_when_widget_value": {"distribution": ["levels"]},
}),
"processing": (["field", "rows"], {"default": "field"}),
}
}
OUTPUTS = (
('DATA_FIELD', 'result'),
)
FUNCTION = "process"
DESCRIPTION = (
"Transform pixel values so their distribution matches a target shape "
"(uniform, Gaussian, or discrete levels) using rank-based reassignment. "
"Equivalent to Gwyddion's coerce.c module."
)
def process(self, field: DataField, distribution: str, n_levels: int,
processing: str) -> tuple:
data = np.asarray(field.data, dtype=np.float64)
if processing == "rows":
result = np.empty_like(data)
for i in range(data.shape[0]):
result[i] = _coerce_block(data[i], distribution, n_levels)
else:
result = _coerce_block(data, distribution, n_levels)
return (field.replace(data=result),)

198
backend/nodes/dwt_anisotropy.py Normal file
View File

@@ -0,0 +1,198 @@
"""DWT anisotropy — quantify surface anisotropy using wavelet decomposition."""
from __future__ import annotations
import numpy as np
from backend.data_types import DataField, RecordTable
from backend.node_registry import register_node
def _next_power_of_2(n: int) -> int:
"""Return the smallest power of 2 >= n."""
p = 1
while p < n:
p <<= 1
return p
def _pad_to_pow2(data: np.ndarray) -> np.ndarray:
"""Pad *data* to the next power-of-2 dimensions using edge values."""
rows, cols = data.shape
new_rows = _next_power_of_2(rows)
new_cols = _next_power_of_2(cols)
if new_rows == rows and new_cols == cols:
return data.copy()
out = np.zeros((new_rows, new_cols), dtype=np.float64)
out[:rows, :cols] = data
# edge-pad
if new_cols > cols:
out[:rows, cols:] = data[:, -1:]
if new_rows > rows:
out[rows:, :cols] = data[-1:, :]
if new_rows > rows and new_cols > cols:
out[rows:, cols:] = data[-1, -1]
return out
def _haar_decompose_2d(data: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""
One level of 2-D Haar wavelet decomposition.
Returns (LL, LH, HL, HH) each of shape (rows//2, cols//2).
LL = (a+b+c+d)/2 approximation
LH = (a+b-c-d)/2 horizontal detail (captures vertical features)
HL = (a-b+c-d)/2 vertical detail (captures horizontal features)
HH = (a-b-c+d)/2 diagonal detail
"""
rows, cols = data.shape
a = data[0::2, 0::2] # top-left
b = data[0::2, 1::2] # top-right
c = data[1::2, 0::2] # bottom-left
d = data[1::2, 1::2] # bottom-right
ll = (a + b + c + d) / 2.0
lh = (a + b - c - d) / 2.0
hl = (a - b + c - d) / 2.0
hh = (a - b - c + d) / 2.0
return ll, lh, hl, hh
def _compute_dwt_anisotropy(
data: np.ndarray,
n_levels: int,
) -> tuple[list[float], list[float], list[float], list[np.ndarray]]:
"""
Multi-level 2-D Haar decomposition with per-level energy ratios.
Returns
-------
x_energies : per-level sum(HL**2)
y_energies : per-level sum(LH**2)
ratios : per-level x_energy / y_energy
ratio_maps : per-level ratio arrays (at decomposition resolution)
"""
padded = _pad_to_pow2(np.asarray(data, dtype=np.float64))
current = padded
x_energies: list[float] = []
y_energies: list[float] = []
ratios: list[float] = []
ratio_maps: list[np.ndarray] = []
for _ in range(n_levels):
if current.shape[0] < 2 or current.shape[1] < 2:
break
ll, lh, hl, hh = _haar_decompose_2d(current)
x_energy = float(np.sum(hl ** 2))
y_energy = float(np.sum(lh ** 2))
ratio = x_energy / (y_energy + 1e-30)
x_energies.append(x_energy)
y_energies.append(y_energy)
ratios.append(ratio)
# Build a per-pixel ratio map at this level's resolution
hl_sq = hl ** 2
lh_sq = lh ** 2
level_map = hl_sq / (lh_sq + 1e-30)
ratio_maps.append(level_map)
current = ll
return x_energies, y_energies, ratios, ratio_maps
def _build_anisotropy_map(
ratio_maps: list[np.ndarray],
orig_rows: int,
orig_cols: int,
) -> np.ndarray:
"""
Combine per-level ratio maps into a single anisotropy map at
the original field resolution by upsampling and averaging.
"""
if not ratio_maps:
return np.ones((orig_rows, orig_cols), dtype=np.float64)
from scipy.ndimage import zoom
target = np.zeros((orig_rows, orig_cols), dtype=np.float64)
for level_map in ratio_maps:
zy = orig_rows / level_map.shape[0]
zx = orig_cols / level_map.shape[1]
upsampled = zoom(level_map, (zy, zx), order=1)
# zoom may produce shape off by one — trim or pad
upsampled = upsampled[:orig_rows, :orig_cols]
if upsampled.shape[0] < orig_rows or upsampled.shape[1] < orig_cols:
tmp = np.ones((orig_rows, orig_cols), dtype=np.float64)
tmp[: upsampled.shape[0], : upsampled.shape[1]] = upsampled
upsampled = tmp
target += upsampled
target /= len(ratio_maps)
return target
@register_node(display_name="DWT Anisotropy")
class DWTAnisotropy:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"field": ("DATA_FIELD",),
"n_levels": (
"INT",
{"default": 4, "min": 1, "max": 10},
),
"ratio_threshold": (
"FLOAT",
{"default": 0.2, "min": 0.001, "max": 10.0, "step": 0.01},
),
}
}
OUTPUTS = (
('DATA_FIELD', 'anisotropy_map'),
('RECORD_TABLE', 'statistics'),
)
FUNCTION = "process"
DESCRIPTION = (
"Quantify surface anisotropy using a multi-level 2-D Haar wavelet decomposition. "
"At each level, horizontal (HL) and vertical (LH) detail energies are compared to "
"produce an X/Y energy ratio. Ratio > 1 indicates more horizontal features; "
"ratio < 1 indicates more vertical features. Equivalent to Gwyddion's dwtanisotropy.c."
)
def process(
self,
field: DataField,
n_levels: int,
ratio_threshold: float,
) -> tuple:
data = np.asarray(field.data, dtype=np.float64)
orig_rows, orig_cols = data.shape
x_energies, y_energies, ratios, ratio_maps = _compute_dwt_anisotropy(
data, int(n_levels),
)
# Build per-pixel anisotropy map
aniso_map = _build_anisotropy_map(ratio_maps, orig_rows, orig_cols)
aniso_field = field.replace(data=aniso_map)
# Build statistics table
rows = []
for i, (xe, ye, r) in enumerate(zip(x_energies, y_energies, ratios)):
rows.append({
"level": i + 1,
"x_energy": float(xe),
"y_energy": float(ye),
"ratio": float(r),
"anisotropic": abs(r - 1.0) > float(ratio_threshold),
})
stats = RecordTable(rows)
return (aniso_field, stats)

238
backend/nodes/grain_visualization.py Normal file
View File

@@ -0,0 +1,238 @@
"""Grain visualization — visualize grains as geometric shapes."""
from __future__ import annotations
import numpy as np
from scipy.ndimage import label, find_objects, distance_transform_edt
from backend.node_registry import register_node
from backend.data_types import DataField
from backend.nodes.helpers import mask_to_bool, bool_to_mask
def _grain_centroid(grain_mask: np.ndarray, slc: tuple[slice, slice]) -> tuple[float, float]:
"""Return (cy, cx) centroid of a grain within its bounding slice."""
ys, xs = np.where(grain_mask[slc])
cy = float(ys.mean()) + slc[0].start
cx = float(xs.mean()) + slc[1].start
return cy, cx
def _grain_inscribed_radius(grain_mask: np.ndarray, slc: tuple[slice, slice]) -> float:
"""Return the inscribed disc radius for a grain region."""
region = grain_mask[slc]
if not np.any(region):
return 0.0
dt = distance_transform_edt(region)
return float(dt.max())
def _grain_inertia(grain_mask: np.ndarray, slc: tuple[slice, slice]) -> tuple[float, float, float]:
"""Return (semi_major, semi_minor, angle_rad) from the inertia tensor."""
ys, xs = np.where(grain_mask[slc])
cy_local = ys.mean()
cx_local = xs.mean()
dy = ys - cy_local
dx = xs - cx_local
n = len(ys)
if n < 2:
return 1.0, 1.0, 0.0
Ixx = np.sum(dy * dy) / n
Iyy = np.sum(dx * dx) / n
Ixy = -np.sum(dx * dy) / n
# Eigenvalues of the 2x2 inertia tensor
mean_I = (Ixx + Iyy) / 2.0
diff_I = (Ixx - Iyy) / 2.0
discriminant = max(0.0, diff_I ** 2 + Ixy ** 2)
sqrt_disc = np.sqrt(discriminant)
lambda1 = mean_I + sqrt_disc
lambda2 = mean_I - sqrt_disc
# Semi-axes proportional to sqrt of eigenvalues, scaled by 2 for visual size
semi_major = 2.0 * np.sqrt(max(lambda1, 0.0))
semi_minor = 2.0 * np.sqrt(max(lambda2, 0.0))
# Angle of the major axis
angle = 0.5 * np.arctan2(2.0 * Ixy, Iyy - Ixx)
return float(semi_major), float(semi_minor), float(angle)
def _draw_circle_filled(canvas: np.ndarray, cy: float, cx: float, r: float) -> None:
h, w = canvas.shape
y_lo = max(0, int(cy - r - 1))
y_hi = min(h, int(cy + r + 2))
x_lo = max(0, int(cx - r - 1))
x_hi = min(w, int(cx + r + 2))
yy, xx = np.ogrid[y_lo:y_hi, x_lo:x_hi]
dist_sq = (yy - cy) ** 2 + (xx - cx) ** 2
canvas[y_lo:y_hi, x_lo:x_hi] |= (dist_sq <= r * r)
def _draw_circle_outline(canvas: np.ndarray, cy: float, cx: float, r: float, thickness: float = 1.5) -> None:
h, w = canvas.shape
y_lo = max(0, int(cy - r - thickness - 1))
y_hi = min(h, int(cy + r + thickness + 2))
x_lo = max(0, int(cx - r - thickness - 1))
x_hi = min(w, int(cx + r + thickness + 2))
yy, xx = np.ogrid[y_lo:y_hi, x_lo:x_hi]
dist = np.sqrt((yy - cy) ** 2 + (xx - cx) ** 2)
canvas[y_lo:y_hi, x_lo:x_hi] |= (np.abs(dist - r) < thickness)
def _draw_rect_filled(canvas: np.ndarray, y0: int, y1: int, x0: int, x1: int) -> None:
h, w = canvas.shape
y0c, y1c = max(0, y0), min(h, y1)
x0c, x1c = max(0, x0), min(w, x1)
canvas[y0c:y1c, x0c:x1c] = True
def _draw_rect_outline(canvas: np.ndarray, y0: int, y1: int, x0: int, x1: int, thickness: int = 1) -> None:
h, w = canvas.shape
y0c, y1c = max(0, y0), min(h, y1)
x0c, x1c = max(0, x0), min(w, x1)
# Top edge
canvas[y0c:min(h, y0c + thickness), x0c:x1c] = True
# Bottom edge
canvas[max(0, y1c - thickness):y1c, x0c:x1c] = True
# Left edge
canvas[y0c:y1c, x0c:min(w, x0c + thickness)] = True
# Right edge
canvas[y0c:y1c, max(0, x1c - thickness):x1c] = True
def _draw_cross(canvas: np.ndarray, cy: float, cx: float, arm: int = 3) -> None:
h, w = canvas.shape
iy, ix = int(round(cy)), int(round(cx))
for d in range(-arm, arm + 1):
if 0 <= iy + d < h and 0 <= ix < w:
canvas[iy + d, ix] = True
if 0 <= iy < h and 0 <= ix + d < w:
canvas[iy, ix + d] = True
def _draw_ellipse_filled(canvas: np.ndarray, cy: float, cx: float,
semi_major: float, semi_minor: float, angle: float) -> None:
h, w = canvas.shape
r_max = max(semi_major, semi_minor, 1.0)
y_lo = max(0, int(cy - r_max - 1))
y_hi = min(h, int(cy + r_max + 2))
x_lo = max(0, int(cx - r_max - 1))
x_hi = min(w, int(cx + r_max + 2))
yy, xx = np.ogrid[y_lo:y_hi, x_lo:x_hi]
cos_a, sin_a = np.cos(angle), np.sin(angle)
dy = yy - cy
dx = xx - cx
# Rotate into ellipse-aligned coordinates
u = cos_a * dx + sin_a * dy
v = -sin_a * dx + cos_a * dy
a = max(semi_major, 0.5)
b = max(semi_minor, 0.5)
canvas[y_lo:y_hi, x_lo:x_hi] |= ((u / a) ** 2 + (v / b) ** 2 <= 1.0)
def _draw_ellipse_outline(canvas: np.ndarray, cy: float, cx: float,
semi_major: float, semi_minor: float, angle: float,
thickness: float = 1.5) -> None:
h, w = canvas.shape
r_max = max(semi_major, semi_minor, 1.0)
y_lo = max(0, int(cy - r_max - thickness - 1))
y_hi = min(h, int(cy + r_max + thickness + 2))
x_lo = max(0, int(cx - r_max - thickness - 1))
x_hi = min(w, int(cx + r_max + thickness + 2))
yy, xx = np.ogrid[y_lo:y_hi, x_lo:x_hi]
cos_a, sin_a = np.cos(angle), np.sin(angle)
dy = yy - cy
dx = xx - cx
u = cos_a * dx + sin_a * dy
v = -sin_a * dx + cos_a * dy
a = max(semi_major, 0.5)
b = max(semi_minor, 0.5)
ellipse_val = (u / a) ** 2 + (v / b) ** 2
canvas[y_lo:y_hi, x_lo:x_hi] |= (np.abs(np.sqrt(ellipse_val) - 1.0) < thickness / max(a, b))
@register_node(display_name="Grain Visualization")
class GrainVisualization:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"field": ("DATA_FIELD",),
"mask": ("IMAGE",),
"style": (["inscribed_disc", "bounding_box", "centroid", "ellipse"], {"default": "inscribed_disc"}),
"fill": ("BOOLEAN", {"default": False}),
}
}
OUTPUTS = (
('IMAGE', 'result'),
('DATA_FIELD', 'labeled'),
)
FUNCTION = "process"
DESCRIPTION = (
"Visualize labeled grains as geometric shapes — inscribed discs, bounding boxes, "
"centroid markers, or fitted ellipses. Produces a mask image with the chosen shapes "
"and a labeled field where each grain has a unique integer value. "
"Equivalent to Gwyddion's grain selection visualization (grain_makesel)."
)
def process(self, field: DataField, mask: np.ndarray, style: str, fill: bool) -> tuple:
mask_bool = mask_to_bool(mask)
labels, n_grains = label(mask_bool.astype(np.int32))
slices = find_objects(labels)
h, w = mask_bool.shape[:2]
canvas = np.zeros((h, w), dtype=bool)
for gid in range(1, n_grains + 1):
slc = slices[gid - 1]
if slc is None:
continue
grain_mask = labels == gid
cy, cx = _grain_centroid(grain_mask, slc)
if style == "inscribed_disc":
r = _grain_inscribed_radius(grain_mask, slc)
if r < 0.5:
r = 0.5
if fill:
_draw_circle_filled(canvas, cy, cx, r)
else:
_draw_circle_outline(canvas, cy, cx, r)
elif style == "bounding_box":
y0, y1 = slc[0].start, slc[0].stop
x0, x1 = slc[1].start, slc[1].stop
if fill:
_draw_rect_filled(canvas, y0, y1, x0, x1)
else:
_draw_rect_outline(canvas, y0, y1, x0, x1)
elif style == "centroid":
arm = max(3, int(round(min(h, w) * 0.01)))
_draw_cross(canvas, cy, cx, arm)
elif style == "ellipse":
semi_major, semi_minor, angle = _grain_inertia(grain_mask, slc)
if semi_major < 0.5:
semi_major = 0.5
if semi_minor < 0.5:
semi_minor = 0.5
if fill:
_draw_ellipse_filled(canvas, cy, cx, semi_major, semi_minor, angle)
else:
_draw_ellipse_outline(canvas, cy, cx, semi_major, semi_minor, angle)
else:
raise ValueError(f"Unknown visualization style: {style!r}")
result = bool_to_mask(canvas)
labeled_field = field.replace(data=labels.astype(np.float64))
return (result, labeled_field)

229
backend/nodes/logistic_classification.py Normal file
View File

@@ -0,0 +1,229 @@
"""Logistic classification — classify features using logistic regression."""
from __future__ import annotations
import numpy as np
from scipy.ndimage import gaussian_filter, sobel
from backend.node_registry import register_node
from backend.data_types import DataField
from backend.nodes.helpers import mask_to_bool, bool_to_mask
def _build_features(data: np.ndarray, use_gaussians: bool, n_gaussians: int,
use_sobel: bool, use_laplacian: bool) -> np.ndarray:
"""Build a feature matrix from the height field.
Each feature is normalized to zero mean, unit variance. The raw
(normalized) height is always included as the first feature.
"""
h, w = data.shape
features: list[np.ndarray] = []
# Always include raw height (normalized)
features.append(data.ravel().copy())
# Gaussian blur features at increasing scales
if use_gaussians:
for i in range(int(n_gaussians)):
sigma = float(2 ** i)
features.append(gaussian_filter(data, sigma).ravel())
# Sobel gradient features
if use_sobel:
features.append(sobel(data, axis=0).ravel())
features.append(sobel(data, axis=1).ravel())
# Laplacian feature (sum of second differences)
if use_laplacian:
lap = np.zeros_like(data)
lap[1:-1, :] += data[:-2, :] - 2 * data[1:-1, :] + data[2:, :]
lap[:, 1:-1] += data[:, :-2] - 2 * data[:, 1:-1] + data[:, 2:]
features.append(lap.ravel())
# Stack into (n_pixels, n_features) matrix
X = np.column_stack(features)
# Normalize each feature to zero mean, unit variance
means = X.mean(axis=0)
stds = X.std(axis=0)
stds[stds == 0] = 1.0
X = (X - means) / stds
# Add bias column
X = np.column_stack([np.ones(X.shape[0]), X])
return X
def _sigmoid(z: np.ndarray) -> np.ndarray:
z = np.clip(z, -500, 500)
return 1.0 / (1.0 + np.exp(-z))
def _otsu_threshold(data: np.ndarray) -> float:
"""Simple Otsu threshold on flattened data."""
flat = data.ravel()
counts, bin_edges = np.histogram(flat, bins=256)
centers = 0.5 * (bin_edges[:-1] + bin_edges[1:])
total = counts.sum()
if total == 0:
return float(np.median(flat))
sum_total = (counts * centers).sum()
sum_bg = 0.0
weight_bg = 0.0
best_var = -1.0
best_thresh = float(centers[0])
for i in range(len(counts)):
weight_bg += counts[i]
if weight_bg == 0:
continue
weight_fg = total - weight_bg
if weight_fg == 0:
break
sum_bg += counts[i] * centers[i]
mean_bg = sum_bg / weight_bg
mean_fg = (sum_total - sum_bg) / weight_fg
var_between = weight_bg * weight_fg * (mean_bg - mean_fg) ** 2
if var_between > best_var:
best_var = var_between
best_thresh = float(centers[i])
return best_thresh
def _train_logistic(X: np.ndarray, y: np.ndarray, regularization: float,
max_iter: int, seed: int) -> np.ndarray:
"""Train logistic regression via gradient descent.
Parameters
----------
X : (m, n_features+1) array with bias column already included.
y : (m,) binary labels (0 or 1).
regularization : L2 penalty lambda.
max_iter : maximum gradient descent iterations.
seed : random seed (unused here; theta starts at zeros).
Returns
-------
theta : (n_features+1,) weight vector.
"""
rng = np.random.default_rng(seed)
n = X.shape[1]
theta = np.zeros(n)
m = len(y)
lr = 0.1
for _ in range(max_iter):
h = _sigmoid(X @ theta)
error = h - y
grad = X.T @ error / m
# L2 regularization (don't regularize bias at index 0)
reg_term = (regularization / m) * theta
reg_term[0] = 0.0
grad += reg_term
theta -= lr * grad
if np.linalg.norm(grad) < 1e-6:
break
return theta
@register_node(display_name="Logistic Classification")
class LogisticClassification:
_CUSTOM_PREVIEW = True
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"field": ("DATA_FIELD",),
"use_gaussians": ("BOOLEAN", {"default": True}),
"n_gaussians": ("INT", {
"default": 4, "min": 1, "max": 10,
"show_when_widget_value": {"use_gaussians": [True]},
}),
"use_sobel": ("BOOLEAN", {"default": True}),
"use_laplacian": ("BOOLEAN", {"default": True}),
"regularization": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1}),
"max_iter": ("INT", {"default": 500, "min": 10, "max": 5000}),
"seed": ("INT", {"default": 42, "min": 0, "max": 999999}),
},
"optional": {
"training_mask": ("IMAGE",),
},
}
OUTPUTS = (
('IMAGE', 'mask'),
('DATA_FIELD', 'probability'),
)
FUNCTION = "process"
DESCRIPTION = (
"Classify surface features using logistic regression on engineered "
"height-derived features (Gaussian blurs, Sobel gradients, Laplacian). "
"Optionally accepts a training mask; otherwise an Otsu-based threshold "
"generates pseudo-labels automatically."
)
def process(
self,
field: DataField,
use_gaussians: bool,
n_gaussians: int,
use_sobel: bool,
use_laplacian: bool,
regularization: float,
max_iter: int,
seed: int,
training_mask: np.ndarray | None = None,
) -> tuple:
data = np.asarray(field.data, dtype=np.float64)
h, w = data.shape
# Build feature matrix for all pixels
X_all = _build_features(data, use_gaussians, n_gaussians, use_sobel, use_laplacian)
if training_mask is not None:
# Extract training labels from the mask
mask_bool = mask_to_bool(training_mask)
if mask_bool.shape[:2] != (h, w):
raise ValueError(
f"Training mask shape {mask_bool.shape} does not match "
f"field shape {(h, w)}."
)
labeled_pixels = mask_bool.ravel()
# Use masked pixels as positive class, unmasked as negative
y_train = labeled_pixels.astype(np.float64)
X_train = X_all
else:
# No training mask: use Otsu threshold to create pseudo-labels
threshold = _otsu_threshold(data)
y_train = (data.ravel() >= threshold).astype(np.float64)
X_train = X_all
# Train logistic regression
theta = _train_logistic(X_train, y_train, regularization, max_iter, seed)
# Apply to all pixels
probability = _sigmoid(X_all @ theta).reshape(h, w)
# Create binary mask
mask = bool_to_mask(probability > 0.5)
# Emit preview
from backend.execution_context import emit_preview
from backend.data_types import encode_preview
from backend.nodes.helpers import _mask_overlay
emit_preview(encode_preview(_mask_overlay(field, mask)))
# Build probability output as a DataField
prob_field = field.replace(data=probability, si_unit_z="")
return (mask, prob_field)

71
backend/nodes/mark_disconnected.py Normal file
View File

@@ -0,0 +1,71 @@
"""Mark disconnected regions — mask topologically isolated surface regions."""
from __future__ import annotations
import numpy as np
from scipy.ndimage import grey_opening, grey_closing
from backend.node_registry import register_node
from backend.data_types import DataField
from backend.nodes.helpers import bool_to_mask, _mask_structure, emit_mask_preview
@register_node(display_name="Mark Disconnected")
class MarkDisconnected:
"""
Detect topologically disconnected (isolated) surface regions using
morphological opening/closing to build a defect-free reference, then
thresholding the residual difference.
"""
_CUSTOM_PREVIEW = True
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"field": ("DATA_FIELD",),
"defect_type": (["positive", "negative", "both"],),
"radius": ("INT", {"default": 5, "min": 1, "max": 100, "step": 1}),
"threshold": ("FLOAT", {"default": 0.1, "min": 0.001, "max": 1.0, "step": 0.001}),
}
}
OUTPUTS = (
('IMAGE', 'mask'),
)
FUNCTION = "process"
DESCRIPTION = (
"Mark topologically disconnected (isolated) surface regions. "
"A morphological opening followed by closing builds a smooth "
"defect-free reference surface; pixels whose deviation from that "
"reference exceeds the sensitivity threshold are flagged. "
"Equivalent to Gwyddion's mark_disconn module."
)
def process(self, field: DataField, defect_type: str, radius: int, threshold: float) -> tuple:
data = field.data.astype(np.float64)
# Build a disk structuring element for grey-scale morphology.
struct = _mask_structure(radius, "disk")
# Morphological opening then closing produces a defect-free reference.
reference = grey_opening(data, footprint=struct)
reference = grey_closing(reference, footprint=struct)
difference = data - reference
diff_range = difference.max() - difference.min()
# Avoid division-by-zero on perfectly flat surfaces.
if diff_range == 0:
mask = np.zeros(data.shape, dtype=bool)
else:
abs_threshold = threshold * diff_range
if defect_type == "positive":
mask = difference > abs_threshold
elif defect_type == "negative":
mask = difference < -abs_threshold
else: # "both"
mask = np.abs(difference) > abs_threshold
out = bool_to_mask(mask)
emit_mask_preview(field, out)
return (out,)

84
backend/nodes/mask_noisify.py Normal file
View File

@@ -0,0 +1,84 @@
"""Mask noisify -- add random perturbation to mask boundaries."""
from __future__ import annotations
import numpy as np
from backend.node_registry import register_node
from backend.data_types import DataField
from backend.nodes.helpers import mask_to_bool, bool_to_mask, emit_mask_preview
@register_node(display_name="Mask Noisify")
class MaskNoisify:
"""
Add random perturbation to mask boundaries.
"""
_CUSTOM_PREVIEW = True
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"mask": ("IMAGE",),
"density": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01}),
"direction": (["both", "add", "remove"],),
"boundaries_only": ("BOOLEAN", {"default": True}),
"seed": ("INT", {"default": 42, "min": 0, "max": 999999}),
},
"optional": {
"field": ("DATA_FIELD",),
}
}
OUTPUTS = (
('IMAGE', 'mask'),
)
FUNCTION = "process"
DESCRIPTION = (
"Add random noise to a binary mask by flipping pixels near boundaries. "
"Control the fraction of affected pixels with density, restrict changes "
"to boundary pixels, and choose whether to add, remove, or both. "
"Use a fixed seed for reproducible results."
)
def process(self, mask: np.ndarray, density: float, direction: str,
boundaries_only: bool, seed: int,
field: DataField | None = None) -> tuple:
binary = mask_to_bool(mask)
# Identify boundary pixels: pixels that differ from at least one neighbour
if boundaries_only:
boundary = np.zeros_like(binary)
for shift_axis, shift_dir in [(0, 1), (0, -1), (1, 1), (1, -1)]:
shifted = np.roll(binary, shift_dir, axis=shift_axis)
boundary |= (binary != shifted)
else:
boundary = np.ones_like(binary)
# Select candidate pixels based on direction
if direction == "add":
candidates = boundary & ~binary
elif direction == "remove":
candidates = boundary & binary
else: # "both"
candidates = boundary
# Randomly flip density fraction of candidates
candidate_indices = np.argwhere(candidates)
n_candidates = len(candidate_indices)
if n_candidates > 0 and density > 0:
rng = np.random.default_rng(seed)
n_flip = int(round(density * n_candidates))
n_flip = max(0, min(n_flip, n_candidates))
if n_flip > 0:
chosen = rng.choice(n_candidates, size=n_flip, replace=False)
for idx in chosen:
r, c = candidate_indices[idx]
binary[r, c] = ~binary[r, c]
out = bool_to_mask(binary)
emit_mask_preview(field, out)
return (out,)

98
backend/nodes/mask_shift.py Normal file
View File

@@ -0,0 +1,98 @@
"""Mask shift — translate mask by pixel offset."""
from __future__ import annotations
import numpy as np
from backend.node_registry import register_node
from backend.data_types import DataField
from backend.nodes.helpers import mask_to_bool, bool_to_mask, emit_mask_preview
@register_node(display_name="Mask Shift")
class MaskShift:
"""Translate a binary mask by an integer pixel offset."""
_CUSTOM_PREVIEW = True
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"mask": ("IMAGE",),
"shift_x": ("INT", {"default": 0, "min": -1000, "max": 1000, "step": 1}),
"shift_y": ("INT", {"default": 0, "min": -1000, "max": 1000, "step": 1}),
"border_mode": (["zero", "wrap", "mirror"],),
},
"optional": {
"field": ("DATA_FIELD",),
},
}
OUTPUTS = (
('IMAGE', 'mask'),
)
FUNCTION = "process"
DESCRIPTION = (
"Translate a binary mask by an integer pixel offset. "
"Choose how out-of-bounds regions are filled: zero (empty), "
"wrap (periodic roll), or mirror (reflected padding)."
)
def process(self, mask: np.ndarray, shift_x: int, shift_y: int,
border_mode: str, field: DataField | None = None) -> tuple:
binary = mask_to_bool(mask)
if border_mode == "wrap":
result = self._shift_wrap(binary, shift_x, shift_y)
elif border_mode == "zero":
result = self._shift_zero(binary, shift_x, shift_y)
elif border_mode == "mirror":
result = self._shift_mirror(binary, shift_x, shift_y)
else:
raise ValueError(f"Unknown border mode: {border_mode}")
out = bool_to_mask(result)
emit_mask_preview(field, out)
return (out,)
@staticmethod
def _shift_wrap(binary: np.ndarray, sx: int, sy: int) -> np.ndarray:
"""Shift with periodic wrapping (np.roll)."""
return np.roll(np.roll(binary, sx, axis=1), sy, axis=0)
@staticmethod
def _shift_zero(binary: np.ndarray, sx: int, sy: int) -> np.ndarray:
"""Shift then zero-fill the wrapped region."""
result = np.roll(np.roll(binary, sx, axis=1), sy, axis=0)
h, w = result.shape[:2]
# Zero-fill columns wrapped by horizontal shift
if sx > 0:
result[:, :sx] = False
elif sx < 0:
result[:, w + sx:] = False
# Zero-fill rows wrapped by vertical shift
if sy > 0:
result[:sy, :] = False
elif sy < 0:
result[h + sy:, :] = False
return result
@staticmethod
def _shift_mirror(binary: np.ndarray, sx: int, sy: int) -> np.ndarray:
"""Shift using reflected padding then crop back to original size."""
h, w = binary.shape[:2]
abs_sx = abs(sx)
abs_sy = abs(sy)
# Pad with reflect mode
padded = np.pad(binary, ((abs_sy, abs_sy), (abs_sx, abs_sx)), mode="reflect")
# Crop with offset to achieve the shift
row_start = abs_sy - sy
col_start = abs_sx - sx
return padded[row_start:row_start + h, col_start:col_start + w]

204
backend/nodes/neural_classification.py Normal file
View File

@@ -0,0 +1,204 @@
"""Neural network classification — classify pixels using a simple feedforward network."""
from __future__ import annotations
import numpy as np
from scipy.ndimage import gaussian_filter
from backend.node_registry import register_node
from backend.data_types import DataField
from backend.nodes.helpers import mask_to_bool, bool_to_mask
def _sigmoid(x: np.ndarray) -> np.ndarray:
"""Numerically stable sigmoid."""
return np.where(
x >= 0,
1.0 / (1.0 + np.exp(-x)),
np.exp(x) / (1.0 + np.exp(x)),
)
def _extract_features(data: np.ndarray, n_gaussians: int) -> np.ndarray:
"""Build multi-scale Gaussian feature matrix from 2-D data.
For each scale sigma = 2^i (i = 0 .. n_gaussians-1), compute
gaussian_filter(data, sigma) and stack as feature columns.
Each feature is normalised to zero mean and unit variance.
"""
rows, cols = data.shape
features = np.empty((rows * cols, n_gaussians), dtype=np.float64)
for i in range(n_gaussians):
sigma = 2.0 ** i
blurred = gaussian_filter(data, sigma).ravel()
mean = blurred.mean()
std = blurred.std()
if std > 0:
blurred = (blurred - mean) / std
else:
blurred = blurred - mean
features[:, i] = blurred
return features
def _forward(X: np.ndarray, W1: np.ndarray, b1: np.ndarray,
W2: np.ndarray, b2: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
"""Forward pass through a 2-layer sigmoid network."""
h = _sigmoid(X @ W1 + b1)
y = _sigmoid(h @ W2 + b2)
return h, y
def _train_network(
X: np.ndarray,
targets: np.ndarray,
W1: np.ndarray,
b1: np.ndarray,
W2: np.ndarray,
b2: np.ndarray,
train_steps: int,
lr: float = 0.1,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""Train via gradient descent with binary cross-entropy loss."""
eps = 1e-7
n = X.shape[0]
t = targets.reshape(-1, 1)
for _ in range(train_steps):
# Forward
h, y = _forward(X, W1, b1, W2, b2)
# Clamp to avoid log(0)
y_clamped = np.clip(y, eps, 1.0 - eps)
# Backward — output layer
dy = (y_clamped - t) / (y_clamped * (1.0 - y_clamped) + eps)
dy *= y * (1.0 - y) # sigmoid derivative
dW2 = (h.T @ dy) / n
db2 = dy.mean(axis=0)
# Backward — hidden layer
dh = (dy @ W2.T) * h * (1.0 - h)
dW1 = (X.T @ dh) / n
db1 = dh.mean(axis=0)
# Update
W1 -= lr * dW1
b1 -= lr * db1
W2 -= lr * dW2
b2 -= lr * db2
return W1, b1, W2, b2
@register_node(display_name="Neural Classification")
class NeuralClassification:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"field": ("DATA_FIELD",),
"n_gaussians": ("INT", {"default": 4, "min": 1, "max": 10, "step": 1}),
"n_hidden": ("INT", {"default": 16, "min": 4, "max": 128, "step": 1}),
"train_steps": ("INT", {"default": 200, "min": 10, "max": 5000, "step": 1}),
"seed": ("INT", {"default": 42, "min": 0, "max": 999999, "step": 1}),
},
"optional": {
"training_mask": ("IMAGE",),
},
}
OUTPUTS = (
('IMAGE', 'mask'),
('DATA_FIELD', 'probability'),
)
FUNCTION = "process"
DESCRIPTION = (
"Classify surface pixels into two classes using a simple two-layer "
"feedforward neural network with sigmoid activations. Features are "
"extracted via multi-scale Gaussian filtering. When a training mask "
"is provided the network learns from labelled pixels; otherwise it "
"uses unsupervised self-labelling from the initial random projection. "
"Equivalent in purpose to Gwyddion's neural.c classifier."
)
def process(
self,
field: DataField,
n_gaussians: int,
n_hidden: int,
train_steps: int,
seed: int,
training_mask: np.ndarray | None = None,
) -> tuple:
data = np.asarray(field.data, dtype=np.float64)
yres, xres = data.shape
n_features = int(n_gaussians)
n_hidden = int(n_hidden)
train_steps = int(train_steps)
# 1. Feature extraction
X_all = _extract_features(data, n_features)
# 2. Initialise weights
rng = np.random.default_rng(int(seed))
scale1 = np.sqrt(2.0 / n_features)
W1 = rng.standard_normal((n_features, n_hidden)) * scale1
b1 = np.zeros(n_hidden)
scale2 = np.sqrt(2.0 / n_hidden)
W2 = rng.standard_normal((n_hidden, 1)) * scale2
b2 = np.zeros(1)
# 3/4. Training
if training_mask is not None:
# Supervised — use labelled pixels
mask_bool = mask_to_bool(training_mask)
if mask_bool.shape != data.shape:
raise ValueError(
f"Training mask shape {mask_bool.shape} does not match "
f"field shape {data.shape}."
)
# Class B = masked (255), class A = unmasked but we need both labels.
# Pixels that are 0 are class A, pixels that are 255 are class B.
# We train on ALL pixels that have a definitive label.
labels_flat = training_mask.ravel().astype(np.float64) / 255.0
# Use all pixels as training data (0 = class A, 1 = class B)
X_train = X_all
targets = labels_flat
W1, b1, W2, b2 = _train_network(
X_train, targets, W1, b1, W2, b2, train_steps,
)
else:
# Unsupervised — use random projection to create initial labels,
# then refine with self-training.
_, y_init = _forward(X_all, W1, b1, W2, b2)
self_labels = (y_init.ravel() > 0.5).astype(np.float64)
# Train on the self-assigned labels for a few iterations
steps = min(train_steps, 50)
W1, b1, W2, b2 = _train_network(
X_all, self_labels, W1, b1, W2, b2, steps,
)
# 5. Apply trained network to all pixels
_, prob_flat = _forward(X_all, W1, b1, W2, b2)
probability = prob_flat.reshape(yres, xres)
# 6. Build outputs
mask = bool_to_mask(probability > 0.5)
prob_field = DataField(
data=probability,
xreal=field.xreal,
yreal=field.yreal,
xoff=field.xoff,
yoff=field.yoff,
si_unit_xy=field.si_unit_xy,
si_unit_z="",
domain="spatial",
)
return (mask, prob_field)

209
backend/nodes/pixel_classification.py Normal file
View File

@@ -0,0 +1,209 @@
"""Pixel classification — classify pixels using decision tree on height, slope, and curvature."""
from __future__ import annotations
import numpy as np
from backend.node_registry import register_node
from backend.data_types import DataField
from backend.nodes.helpers import bool_to_mask
def _compute_slope(data: np.ndarray) -> np.ndarray:
"""Gradient magnitude via np.gradient."""
gy, gx = np.gradient(data.astype(np.float64))
return np.sqrt(gx**2 + gy**2)
def _compute_curvature(data: np.ndarray) -> np.ndarray:
"""Laplacian (sum of second derivatives)."""
d = data.astype(np.float64)
gy, gx = np.gradient(d)
gyy, _ = np.gradient(gy)
_, gxx = np.gradient(gx)
return np.abs(gxx + gyy)
def _feature_maps(data: np.ndarray, feature: str) -> list[np.ndarray]:
"""Return a list of 2-D feature arrays based on the feature selector."""
height = data.astype(np.float64)
if feature == "height":
return [height]
slope = _compute_slope(data)
if feature == "slope":
return [slope]
curvature = _compute_curvature(data)
if feature == "curvature":
return [curvature]
if feature == "height_slope":
return [height, slope]
# "all"
return [height, slope, curvature]
def _normalize_01(arr: np.ndarray) -> np.ndarray:
vmin, vmax = arr.min(), arr.max()
if vmax > vmin:
return (arr - vmin) / (vmax - vmin)
return np.zeros_like(arr)
def _classify_single(values: np.ndarray, n_classes: int, method: str) -> np.ndarray:
"""Classify a single feature map into n_classes using the chosen method."""
labels = np.zeros(values.shape, dtype=np.int32)
if method == "equal_range":
vmin, vmax = values.min(), values.max()
if vmax <= vmin:
return labels
edges = np.linspace(vmin, vmax, n_classes + 1)
for i in range(n_classes - 1):
labels[values >= edges[i + 1]] = i + 1
elif method == "quantile":
percentiles = np.linspace(0, 100, n_classes + 1)
edges = np.percentile(values, percentiles)
for i in range(n_classes - 1):
labels[values >= edges[i + 1]] = i + 1
elif method == "otsu":
# Multi-Otsu: find n_classes-1 thresholds via histogram analysis
flat = values.ravel()
n_bins = min(256, max(32, len(flat) // 10))
counts, bin_edges = np.histogram(flat, bins=n_bins)
centers = 0.5 * (bin_edges[:-1] + bin_edges[1:])
total = counts.sum()
if total == 0 or n_classes < 2:
return labels
# For multi-Otsu, find thresholds that minimise intra-class variance
# Use quantile-based initial thresholds then refine with exhaustive
# search over histogram bins for each threshold
thresholds = []
if n_classes == 2:
# Standard single-threshold Otsu
best_var = -1.0
best_t = 0
cum_sum = 0.0
cum_count = 0
total_sum = float(np.sum(counts * centers))
for i in range(n_bins - 1):
cum_count += counts[i]
cum_sum += counts[i] * centers[i]
if cum_count == 0 or cum_count == total:
continue
w0 = cum_count / total
w1 = 1.0 - w0
mu0 = cum_sum / cum_count
mu1 = (total_sum - cum_sum) / (total - cum_count)
between_var = w0 * w1 * (mu0 - mu1) ** 2
if between_var > best_var:
best_var = between_var
best_t = i
thresholds = [0.5 * (bin_edges[best_t + 1] + bin_edges[best_t + 2])]
else:
# Multi-threshold: use quantile splits as a good approximation
percentiles = np.linspace(0, 100, n_classes + 1)[1:-1]
thresholds = list(np.percentile(flat, percentiles))
thresholds = sorted(thresholds)
for i, t in enumerate(thresholds):
labels[values >= t] = i + 1
else:
raise ValueError(f"Unknown classification method: {method!r}")
return labels
def _kmeans_classify(features: np.ndarray, n_classes: int, max_iter: int = 20) -> np.ndarray:
"""Simple k-means on stacked normalised features.
Parameters
----------
features : (n_pixels, n_features) array
n_classes : number of clusters
max_iter : maximum iterations
Returns
-------
labels : (n_pixels,) int32 array with values in [0, n_classes-1]
"""
rng = np.random.RandomState(42)
n_pixels = features.shape[0]
# Initialise centroids by choosing random data points
indices = rng.choice(n_pixels, size=min(n_classes, n_pixels), replace=False)
centroids = features[indices].copy()
labels = np.zeros(n_pixels, dtype=np.int32)
for _ in range(max_iter):
# Assign each pixel to nearest centroid
dists = np.stack([
np.sum((features - c) ** 2, axis=1) for c in centroids
], axis=1) # (n_pixels, n_classes)
new_labels = np.argmin(dists, axis=1).astype(np.int32)
if np.array_equal(new_labels, labels):
break
labels = new_labels
# Update centroids
for k in range(n_classes):
members = features[labels == k]
if len(members) > 0:
centroids[k] = members.mean(axis=0)
return labels
@register_node(display_name="Pixel Classification")
class PixelClassification:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"field": ("DATA_FIELD",),
"n_classes": ("INT", {"default": 3, "min": 2, "max": 10, "step": 1}),
"feature": (["height", "slope", "curvature", "height_slope", "all"],),
"method": (["otsu", "equal_range", "quantile"],),
}
}
OUTPUTS = (
('DATA_FIELD', 'classified'),
('IMAGE', 'mask'),
)
FUNCTION = "process"
DESCRIPTION = (
"Classify pixels into discrete classes based on height, slope, and/or curvature. "
"Single-feature modes use threshold-based classification (Otsu, equal range, or quantile). "
"Multi-feature modes (height_slope, all) use k-means clustering. "
"Equivalent to Gwyddion's classify.c module."
)
def process(self, field: DataField, n_classes: int, feature: str, method: str) -> tuple:
data = np.asarray(field.data, dtype=np.float64)
maps = _feature_maps(data, feature)
if len(maps) == 1:
# Single-feature: use threshold-based classification
labels = _classify_single(maps[0], int(n_classes), method)
else:
# Multi-feature: normalise and use k-means
normed = [_normalize_01(m) for m in maps]
stacked = np.stack([m.ravel() for m in normed], axis=1) # (n_pixels, n_features)
labels = _kmeans_classify(stacked, int(n_classes)).reshape(data.shape)
# Build output DataField with integer class labels
classified = DataField(
data=labels.astype(np.float64),
xreal=field.xreal,
yreal=field.yreal,
si_unit_xy=field.si_unit_xy,
si_unit_z="",
)
# Mask for class 0
mask = bool_to_mask(labels == 0)
return (classified, mask)

88
backend/nodes/presentation_ops.py Normal file
View File

@@ -0,0 +1,88 @@
"""Presentation operations -- manage presentation overlays on data fields."""
from __future__ import annotations
import numpy as np
from backend.node_registry import register_node
from backend.data_types import DataField
@register_node(display_name="Presentation Ops")
class PresentationOps:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"field": ("DATA_FIELD",),
"operation": (["logscale", "extract_presentation", "attach", "blend"],),
"blend_factor": ("FLOAT", {
"default": 0.5,
"min": 0.0,
"max": 1.0,
"step": 0.01,
"show_when_widget_value": {"operation": ["blend"]},
}),
},
"optional": {
"overlay": ("DATA_FIELD", {
"show_when_widget_value": {"operation": ["attach", "blend"]},
}),
},
}
OUTPUTS = (
('DATA_FIELD', 'result'),
)
FUNCTION = "process"
DESCRIPTION = (
"Manage presentation overlays on data fields. "
"logscale applies logarithmic scaling for visualising data with large dynamic range. "
"extract_presentation normalises the field to [0, 1]. "
"attach replaces the field data with an overlay (resampled if needed). "
"blend linearly mixes the field and overlay by a configurable factor. "
"Equivalent to Gwyddion's presentationops.c module."
)
def process(self, field: DataField, operation: str, blend_factor: float,
overlay: DataField | None = None) -> tuple:
data = np.asarray(field.data, dtype=np.float64)
if operation == "logscale":
data_pos = data - data.min() + 1e-30
result = np.log10(data_pos)
elif operation == "extract_presentation":
dmin, dmax = data.min(), data.max()
if dmax > dmin:
result = (data - dmin) / (dmax - dmin)
else:
result = np.zeros_like(data)
elif operation == "attach":
if overlay is None:
raise ValueError("'attach' operation requires an overlay field.")
overlay_data = np.asarray(overlay.data, dtype=np.float64)
result = self._match_shape(overlay_data, data.shape)
elif operation == "blend":
if overlay is None:
raise ValueError("'blend' operation requires an overlay field.")
overlay_data = np.asarray(overlay.data, dtype=np.float64)
overlay_matched = self._match_shape(overlay_data, data.shape)
result = (1.0 - blend_factor) * data + blend_factor * overlay_matched
else:
raise ValueError(f"Unknown operation: {operation!r}")
return (field.replace(data=result),)
@staticmethod
def _match_shape(source: np.ndarray, target_shape: tuple[int, ...]) -> np.ndarray:
"""Resample *source* to *target_shape* using scipy zoom if shapes differ."""
if source.shape == target_shape:
return source
from scipy.ndimage import zoom
factors = tuple(t / s for t, s in zip(target_shape, source.shape))
return zoom(source, factors, order=3)

177
backend/nodes/psf_estimation.py Normal file
View File

@@ -0,0 +1,177 @@
"""PSF estimation — estimate and fit point spread functions for deconvolution."""
from __future__ import annotations
import numpy as np
from backend.node_registry import register_node
from backend.data_types import DataField, RecordTable
@register_node(display_name="PSF Estimation")
class PSFEstimation:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"measured": ("DATA_FIELD",),
"ideal": ("DATA_FIELD",),
"method": (["wiener", "least_squares", "gaussian_fit"], {"default": "wiener"}),
"regularization": ("FLOAT", {"default": 0.01, "min": 1e-6, "max": 1.0, "step": 0.001}),
"psf_size": ("INT", {"default": 32, "min": 4, "max": 128}),
}
}
OUTPUTS = (
('DATA_FIELD', 'psf'),
('RECORD_TABLE', 'parameters'),
)
FUNCTION = "process"
DESCRIPTION = (
"Estimate a point spread function (PSF) from a measured (blurred) image "
"and an ideal (sharp) reference. The PSF can then be used with the "
"Deconvolution node to restore other images. Three methods are available: "
"pseudo-Wiener deconvolution, regularised least-squares, and Gaussian fit. "
"Equivalent to Gwyddion's psf.c / psf-fit.c modules."
)
# ------------------------------------------------------------------
# helpers
# ------------------------------------------------------------------
@staticmethod
def _crop_centre(arr: np.ndarray, size: int) -> np.ndarray:
"""Crop the central *size x size* region from *arr*."""
yc, xc = arr.shape[0] // 2, arr.shape[1] // 2
half = size // 2
return arr[yc - half : yc - half + size, xc - half : xc - half + size]
@staticmethod
def _normalise(psf: np.ndarray) -> np.ndarray:
"""Normalise so that the PSF sums to 1."""
s = psf.sum()
if abs(s) > 1e-30:
psf = psf / s
return psf
@staticmethod
def _fit_gaussian_2d(psf: np.ndarray):
"""Fit a 2-D Gaussian to *psf* using moment analysis.
Returns (gaussian_array, sigma_x, sigma_y, amplitude).
"""
h, w = psf.shape
psf_pos = np.maximum(psf, 0.0)
total = psf_pos.sum()
if total < 1e-30:
return np.zeros_like(psf), 0.0, 0.0, 0.0
y_idx, x_idx = np.mgrid[0:h, 0:w]
# centroid
cx = float(np.sum(x_idx * psf_pos) / total)
cy = float(np.sum(y_idx * psf_pos) / total)
# second moments → sigma
sx = float(np.sqrt(np.sum(psf_pos * (x_idx - cx) ** 2) / total))
sy = float(np.sqrt(np.sum(psf_pos * (y_idx - cy) ** 2) / total))
sx = max(sx, 1e-6)
sy = max(sy, 1e-6)
amplitude = float(psf_pos.max())
gauss = amplitude * np.exp(
-((x_idx - cx) ** 2 / (2 * sx ** 2) + (y_idx - cy) ** 2 / (2 * sy ** 2))
)
gauss = PSFEstimation._normalise(gauss)
return gauss, sx, sy, amplitude
# ------------------------------------------------------------------
# methods
# ------------------------------------------------------------------
def _wiener(
self,
F_measured: np.ndarray,
F_ideal: np.ndarray,
regularization: float,
psf_size: int,
) -> np.ndarray:
"""Pseudo-Wiener PSF estimation."""
F_psf = np.conj(F_ideal) * F_measured / (np.abs(F_ideal) ** 2 + regularization)
psf = np.real(np.fft.ifft2(F_psf))
psf = np.fft.fftshift(psf)
psf = self._crop_centre(psf, psf_size)
return self._normalise(psf)
def _least_squares(
self,
F_measured: np.ndarray,
F_ideal: np.ndarray,
regularization: float,
psf_size: int,
) -> np.ndarray:
"""Regularised least-squares PSF estimation."""
abs_ideal = np.abs(F_ideal)
F_psf = np.where(
abs_ideal < regularization,
0.0,
F_measured / (F_ideal + regularization * np.sign(F_ideal)),
)
psf = np.real(np.fft.ifft2(F_psf))
psf = np.fft.fftshift(psf)
psf = self._crop_centre(psf, psf_size)
return self._normalise(psf)
# ------------------------------------------------------------------
# main entry
# ------------------------------------------------------------------
def process(
self,
measured: DataField,
ideal: DataField,
method: str,
regularization: float,
psf_size: int,
) -> tuple:
measured_data = np.asarray(measured.data, dtype=np.float64)
ideal_data = np.asarray(ideal.data, dtype=np.float64)
F_measured = np.fft.fft2(measured_data)
F_ideal = np.fft.fft2(ideal_data)
parameters = RecordTable()
if method == "wiener":
psf = self._wiener(F_measured, F_ideal, regularization, psf_size)
elif method == "least_squares":
psf = self._least_squares(F_measured, F_ideal, regularization, psf_size)
elif method == "gaussian_fit":
raw_psf = self._wiener(F_measured, F_ideal, regularization, psf_size)
psf, sigma_x, sigma_y, amplitude = self._fit_gaussian_2d(raw_psf)
parameters = RecordTable([
{"quantity": "sigma_x", "value": sigma_x, "unit": "px"},
{"quantity": "sigma_y", "value": sigma_y, "unit": "px"},
{"quantity": "amplitude", "value": amplitude, "unit": ""},
])
else:
raise ValueError(f"Unknown PSF estimation method: {method!r}")
# Build output DataField — inherit spatial metadata, adjust for psf_size
yres, xres = measured_data.shape
psf_xreal = measured.xreal * psf_size / xres
psf_yreal = measured.yreal * psf_size / yres
psf_field = measured.replace(
data=psf,
xreal=psf_xreal,
yreal=psf_yreal,
xoff=0.0,
yoff=0.0,
)
return (psf_field, parameters)

126
backend/nodes/super_resolution.py Normal file
View File

@@ -0,0 +1,126 @@
"""Super-resolution -- combine multiple aligned scans for resolution enhancement."""
from __future__ import annotations
import numpy as np
from scipy.ndimage import shift as ndimage_shift, zoom as ndimage_zoom
from backend.node_registry import register_node
from backend.data_types import DataField
def _find_subpixel_shift(ref: np.ndarray, img: np.ndarray) -> tuple[float, float]:
"""Estimate the (dy, dx) sub-pixel shift of *img* relative to *ref* via cross-correlation.
Uses FFT-based cross-correlation with parabolic peak refinement.
"""
fa = np.fft.fft2(ref - ref.mean())
fb = np.fft.fft2(img - img.mean())
cross = np.fft.ifft2(fa * np.conj(fb))
cc = np.abs(np.fft.fftshift(cross))
cy, cx = np.array(cc.shape) // 2
peak_y, peak_x = np.unravel_index(np.argmax(cc), cc.shape)
# Integer shift relative to centre
dy = peak_y - cy
dx = peak_x - cx
# Parabolic sub-pixel refinement around peak
h, w = cc.shape
if 1 <= peak_y <= h - 2:
num = float(cc[peak_y - 1, peak_x] - cc[peak_y + 1, peak_x])
den = float(
cc[peak_y - 1, peak_x] - 2.0 * cc[peak_y, peak_x] + cc[peak_y + 1, peak_x]
)
if abs(den) > 1e-12:
dy += 0.5 * num / den
if 1 <= peak_x <= w - 2:
num = float(cc[peak_y, peak_x - 1] - cc[peak_y, peak_x + 1])
den = float(
cc[peak_y, peak_x - 1] - 2.0 * cc[peak_y, peak_x] + cc[peak_y, peak_x + 1]
)
if abs(den) > 1e-12:
dx += 0.5 * num / den
return float(dy), float(dx)
@register_node(display_name="Super Resolution")
class SuperResolution:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"field1": ("DATA_FIELD",),
"upscale": ("INT", {"default": 2, "min": 2, "max": 4, "step": 1}),
},
"optional": {
"field2": ("DATA_FIELD",),
"field3": ("DATA_FIELD",),
"field4": ("DATA_FIELD",),
},
}
OUTPUTS = (
('DATA_FIELD', 'result'),
)
FUNCTION = "process"
DESCRIPTION = (
"Combine multiple aligned scans to produce a super-resolved image with higher "
"spatial resolution. Sub-pixel shifts between inputs are estimated via FFT "
"cross-correlation and used to reconstruct a finer grid. When only one field "
"is provided the image is upsampled using cubic interpolation."
)
def process(
self,
field1: DataField,
upscale: int,
field2: DataField | None = None,
field3: DataField | None = None,
field4: DataField | None = None,
) -> tuple:
fields = [field1]
for f in (field2, field3, field4):
if f is not None:
fields.append(f)
ref = np.asarray(field1.data, dtype=np.float64)
# Upsample reference to target resolution
high_res = ndimage_zoom(ref, upscale, order=3)
weight = np.ones_like(high_res)
if len(fields) == 1:
# Single input -- just return the upsampled reference
return (field1.replace(
data=high_res,
xreal=field1.xreal,
yreal=field1.yreal,
),)
# Multiple inputs -- align, upsample, and average
for extra in fields[1:]:
img = np.asarray(extra.data, dtype=np.float64)
# Find sub-pixel shift relative to reference
dy, dx = _find_subpixel_shift(ref, img)
# Shift in high-res coordinates
shifted = ndimage_shift(img.astype(np.float64), (-dy, -dx), order=3)
upsampled = ndimage_zoom(shifted, upscale, order=3)
# Accumulate
high_res += upsampled
weight += 1.0
high_res /= weight
return (field1.replace(
data=high_res,
xreal=field1.xreal,
yreal=field1.yreal,
),)

351
backend/nodes/tip_shape_estimate.py Normal file
View File

@@ -0,0 +1,351 @@
"""Tip shape estimation — estimate SPM tip geometry from known calibration features."""
from __future__ import annotations
import numpy as np
from backend.node_registry import register_node
from backend.data_types import DataField, RecordTable
@register_node(display_name="Tip Shape Estimate")
class TipShapeEstimate:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"field": ("DATA_FIELD",),
"feature_type": (["edge", "sphere", "cylinder"], {"default": "edge"}),
"feature_radius": ("FLOAT", {
"default": 100e-9, "min": 1e-9, "max": 100e-6, "step": 1e-9,
}),
"n_points": ("INT", {"default": 100, "min": 10, "max": 1000}),
}
}
OUTPUTS = (
('DATA_FIELD', 'tip_shape'),
('RECORD_TABLE', 'parameters'),
)
FUNCTION = "process"
DESCRIPTION = (
"Estimate SPM tip geometry from a known calibration feature. "
"Supported features: edge (sharp step), sphere (calibration ball), "
"cylinder (calibration wire). The image of a known feature is a "
"dilation of the feature with the tip; by subtracting the known "
"feature contribution the tip shape can be recovered. "
"The 2D tip is built by revolving the 1D radial profile (axial "
"symmetry assumption). Output parameters include estimated tip "
"radius of curvature at the apex and half-cone angle. "
"Equivalent to Gwyddion's tipshape.c analysis. "
)
def process(
self,
field: DataField,
feature_type: str,
feature_radius: float,
n_points: int,
) -> tuple:
data = field.data.astype(np.float64)
ny, nx = data.shape
pixel_size = (field.dx + field.dy) * 0.5
# ── Step 1: Extract 1D tip profile depending on feature type ──────
if feature_type == "edge":
tip_profile_1d = self._estimate_from_edge(data, pixel_size, n_points)
elif feature_type == "sphere":
tip_profile_1d = self._estimate_from_sphere(
data, pixel_size, feature_radius, n_points,
)
elif feature_type == "cylinder":
tip_profile_1d = self._estimate_from_cylinder(
data, pixel_size, feature_radius, n_points,
)
else:
raise ValueError(
f"Unknown feature_type {feature_type!r}. "
"Choose: edge, sphere, cylinder."
)
# ── Step 2: Build 2D tip by revolution (axial symmetry) ───────────
n_tip = n_points if n_points % 2 == 1 else n_points + 1
ci = n_tip // 2
offsets = np.arange(n_tip) - ci
gx, gy = np.meshgrid(offsets, offsets)
r_grid = np.sqrt(gx ** 2 + gy ** 2)
# The 1D profile goes from r = 0 to r = max_r.
# Interpolate for every pixel in the 2D grid.
r_1d = np.linspace(0, ci, len(tip_profile_1d))
tip_2d = np.interp(r_grid, r_1d, tip_profile_1d, right=0.0)
# Convention: apex (centre) is the maximum; minimum is 0.
tip_2d -= tip_2d.min()
xreal = n_tip * pixel_size
tip_field = DataField(
data=tip_2d,
xreal=xreal,
yreal=xreal,
si_unit_xy=field.si_unit_xy,
si_unit_z=field.si_unit_z,
)
# ── Step 3: Estimate tip parameters ───────────────────────────────
tip_radius, half_angle = self._estimate_parameters(
tip_profile_1d, pixel_size,
)
table = RecordTable([
{"quantity": "tip_radius", "value": tip_radius, "unit": field.si_unit_z},
{"quantity": "half_angle", "value": half_angle, "unit": "deg"},
])
return (tip_field, table)
# ── Feature-specific estimation methods ───────────────────────────────
@staticmethod
def _estimate_from_edge(
data: np.ndarray,
pixel_size: float,
n_points: int,
) -> np.ndarray:
"""
Edge feature: the tip shape is the mirror of the steepest edge
cross-section in the image.
Find the row with the maximum gradient magnitude, extract the
cross-section at that location, mirror and normalise.
Returns a radial profile: index 0 = apex (maximum), decreasing
outward.
"""
ny, nx = data.shape
# Compute row-wise gradient magnitude to find the steepest edge.
grad = np.abs(np.diff(data, axis=1))
row_grad = grad.sum(axis=1)
best_row = int(np.argmax(row_grad))
profile = data[best_row, :]
# Mirror: the tip is the complement of the edge profile.
tip_raw = np.max(profile) - profile[::-1]
# Find the peak of the mirrored profile and take the radial
# (half) profile from apex outward.
peak = int(np.argmax(tip_raw))
# Use the longer side from the peak to preserve resolution.
left = tip_raw[:peak + 1][::-1] # apex to left edge, reversed
right = tip_raw[peak:] # apex to right edge
half = left if len(left) >= len(right) else right
half = half.copy()
# Ensure monotonically decreasing from apex.
for i in range(1, len(half)):
if half[i] > half[i - 1]:
half[i] = half[i - 1]
# Resample to n_points.
x_raw = np.linspace(0, 1, len(half))
x_out = np.linspace(0, 1, n_points)
tip_profile = np.interp(x_out, x_raw, half)
return tip_profile
@staticmethod
def _estimate_from_sphere(
data: np.ndarray,
pixel_size: float,
radius: float,
n_points: int,
) -> np.ndarray:
"""
Sphere feature: dilation model z_measured = z_sphere (+) z_tip.
Extract radial profile from the highest point and subtract the
ideal sphere contribution to recover the tip profile.
Returns a radial profile: index 0 = apex (maximum), decreasing
outward.
"""
ny, nx = data.shape
# Find the highest point (apex of the imaged sphere).
peak_idx = np.unravel_index(np.argmax(data), data.shape)
cy, cx = peak_idx
# Extract radial profile by averaging azimuthally.
max_r = min(cy, ny - 1 - cy, cx, nx - 1 - cx)
if max_r < 2:
max_r = min(ny, nx) // 2
Y, X = np.ogrid[:ny, :nx]
r_map = np.sqrt(((X - cx) * pixel_size) ** 2 + ((Y - cy) * pixel_size) ** 2)
n_bins = min(max_r, n_points)
r_edges = np.linspace(0, max_r * pixel_size, n_bins + 1)
radial_profile = np.zeros(n_bins)
for i in range(n_bins):
mask = (r_map >= r_edges[i]) & (r_map < r_edges[i + 1])
if mask.any():
radial_profile[i] = data[mask].mean()
elif i > 0:
radial_profile[i] = radial_profile[i - 1]
r_centres = 0.5 * (r_edges[:-1] + r_edges[1:])
# Ideal sphere profile: z_sphere(r) = sqrt(R^2 - r^2) for r < R, else 0.
sphere_profile = np.where(
r_centres < radius,
np.sqrt(np.maximum(radius ** 2 - r_centres ** 2, 0.0)),
0.0,
)
# Tip profile = measured - sphere (dilation subtraction).
tip_raw = radial_profile - sphere_profile
# Shift so that the apex (index 0) is the maximum.
tip_raw = tip_raw - tip_raw.min()
# Ensure monotonically decreasing from apex outward by clamping.
for i in range(1, len(tip_raw)):
if tip_raw[i] > tip_raw[i - 1]:
tip_raw[i] = tip_raw[i - 1]
# Resample to n_points.
x_raw = np.linspace(0, 1, len(tip_raw))
x_out = np.linspace(0, 1, n_points)
tip_profile = np.interp(x_out, x_raw, tip_raw)
return tip_profile
@staticmethod
def _estimate_from_cylinder(
data: np.ndarray,
pixel_size: float,
radius: float,
n_points: int,
) -> np.ndarray:
"""
Cylinder feature: extract cross-section perpendicular to the
cylinder axis and subtract ideal cylinder profile.
The cylinder axis is assumed to run along the direction with the
least height variation.
Returns a radial profile: index 0 = apex (maximum), decreasing
outward.
"""
ny, nx = data.shape
# Determine cylinder axis: compare row-wise vs column-wise variance.
row_var = np.var(np.diff(data, axis=1))
col_var = np.var(np.diff(data, axis=0))
if row_var > col_var:
# Cylinder axis along columns -> cross-section along rows.
profile = data.mean(axis=0)
else:
# Cylinder axis along rows -> cross-section along columns.
profile = data.mean(axis=1)
n_prof = len(profile)
peak = int(np.argmax(profile))
# Ideal cylinder cross-section: z = sqrt(R^2 - x^2) for |x| < R.
x_phys = (np.arange(n_prof) - peak) * pixel_size
cyl_profile = np.where(
np.abs(x_phys) < radius,
np.sqrt(np.maximum(radius ** 2 - x_phys ** 2, 0.0)),
0.0,
)
# Tip = measured - cylinder.
tip_raw = profile - cyl_profile
tip_raw -= tip_raw.min()
# Take the radial (half) profile from the peak outward.
left = tip_raw[:peak + 1][::-1]
right = tip_raw[peak:]
half = left if len(left) >= len(right) else right
# Ensure monotonically decreasing from apex.
for i in range(1, len(half)):
if half[i] > half[i - 1]:
half[i] = half[i - 1]
# Resample to n_points.
x_raw = np.linspace(0, 1, len(half))
x_out = np.linspace(0, 1, n_points)
tip_profile = np.interp(x_out, x_raw, half)
return tip_profile
# ── Parameter estimation ──────────────────────────────────────────────
@staticmethod
def _estimate_parameters(
tip_profile: np.ndarray,
pixel_size: float,
) -> tuple[float, float]:
"""
Estimate tip radius of curvature at the apex and half-cone angle
from the 1D radial profile.
tip_radius: fitted from the parabolic approximation near the apex,
z(r) ~ z_max - r^2 / (2R) => R = r^2 / (2 * (z_max - z(r)))
half_angle: from the slope of the tip walls in the outer half,
tan(half_angle) = dz/dr => half_angle = arctan(slope)
"""
n = len(tip_profile)
r = np.linspace(0, (n - 1) * pixel_size, n)
# ── Tip radius from apex curvature ────────────────────────────────
# Use a few points near the apex for a parabolic fit: z = a - b*r^2
n_apex = max(3, n // 10)
r_apex = r[:n_apex]
z_apex = tip_profile[:n_apex]
if len(r_apex) >= 2 and r_apex[-1] > 0:
# Fit z = c0 + c1 * r^2
A = np.vstack([np.ones(n_apex), r_apex ** 2]).T
try:
coeffs = np.linalg.lstsq(A, z_apex, rcond=None)[0]
c1 = coeffs[1]
# z = z_max - r^2/(2R) => c1 = -1/(2R) => R = -1/(2*c1)
if c1 < 0:
tip_radius = -1.0 / (2.0 * c1)
else:
tip_radius = float('inf')
except np.linalg.LinAlgError:
tip_radius = float('inf')
else:
tip_radius = float('inf')
# ── Half-angle from outer wall slope ──────────────────────────────
# Use the outer 50% of the profile.
mid = n // 2
if mid < n - 1:
r_outer = r[mid:]
z_outer = tip_profile[mid:]
if len(r_outer) >= 2:
dr = r_outer[-1] - r_outer[0]
dz = z_outer[-1] - z_outer[0]
if dr > 0:
slope = abs(dz / dr)
half_angle = np.degrees(np.arctan(slope))
else:
half_angle = 0.0
else:
half_angle = 0.0
else:
half_angle = 0.0
return tip_radius, half_angle