121 lines
4.4 KiB
Python
121 lines
4.4 KiB
Python
"""Straighten path — extract cross-section along a curved spline path."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
from scipy.interpolate import CubicSpline
|
|
from scipy.ndimage import map_coordinates
|
|
|
|
from backend.node_registry import register_node
|
|
from backend.data_types import DataField, LineData, datafield_to_uint8, encode_preview
|
|
from backend.execution_context import emit_overlay
|
|
|
|
|
|
@register_node(display_name="Straighten Path")
|
|
class StraightenPath:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"field": ("DATA_FIELD",),
|
|
"points_x": ("STRING", {"default": "0.25, 0.5, 0.75", "hidden": True}),
|
|
"points_y": ("STRING", {"default": "0.5, 0.3, 0.5", "hidden": True}),
|
|
"thickness": ("INT", {"default": 1, "min": 1, "max": 100, "step": 1}),
|
|
"n_samples": ("INT", {"default": 256, "min": 10, "max": 2048, "step": 1}),
|
|
}
|
|
}
|
|
|
|
OUTPUTS = (
|
|
('DATA_FIELD', 'straightened'),
|
|
('LINE', 'profile'),
|
|
)
|
|
FUNCTION = "process"
|
|
|
|
DESCRIPTION = (
|
|
"Extract a cross-section along an arbitrary curved path defined by "
|
|
"control points. The path is a natural cubic spline through the "
|
|
"points. Drag the points on the preview to reshape the path; the "
|
|
"shaded band shows the sampling thickness. "
|
|
)
|
|
|
|
KEYWORDS = ("unbend", "unroll", "spline", "curved profile", "extract path")
|
|
|
|
def process(self, field: DataField, points_x: str, points_y: str,
|
|
thickness: int, n_samples: int) -> tuple:
|
|
data = np.asarray(field.data, dtype=np.float64)
|
|
yres, xres = data.shape
|
|
|
|
fx = [float(v.strip()) for v in points_x.split(",") if v.strip()]
|
|
fy = [float(v.strip()) for v in points_y.split(",") if v.strip()]
|
|
n_pts = min(len(fx), len(fy))
|
|
fx, fy = fx[:n_pts], fy[:n_pts]
|
|
|
|
emit_overlay({
|
|
"kind": "straighten_path",
|
|
"section_title": "Path",
|
|
"image": encode_preview(datafield_to_uint8(field, field.colormap)),
|
|
"points": [{"x": float(fx[i]), "y": float(fy[i])} for i in range(n_pts)],
|
|
"thickness": int(thickness),
|
|
"xres": int(xres),
|
|
"yres": int(yres),
|
|
})
|
|
|
|
if n_pts < 2:
|
|
empty_line = LineData(
|
|
data=np.zeros(0, dtype=np.float64),
|
|
x_axis=np.zeros(0, dtype=np.float64),
|
|
x_unit=field.si_unit_xy,
|
|
y_unit=field.si_unit_z,
|
|
)
|
|
return (field, empty_line)
|
|
|
|
px = [f * (xres - 1) for f in fx]
|
|
py = [f * (yres - 1) for f in fy]
|
|
|
|
t_ctrl = np.linspace(0, 1, n_pts)
|
|
t_sample = np.linspace(0, 1, n_samples)
|
|
if n_pts >= 3:
|
|
cx = CubicSpline(t_ctrl, px, bc_type="natural")(t_sample)
|
|
cy = CubicSpline(t_ctrl, py, bc_type="natural")(t_sample)
|
|
else:
|
|
cx = np.interp(t_sample, t_ctrl, px)
|
|
cy = np.interp(t_sample, t_ctrl, py)
|
|
|
|
if thickness <= 1:
|
|
values = map_coordinates(data, [cy, cx], order=1, mode='nearest')
|
|
result = values.reshape(1, -1)
|
|
else:
|
|
dcx = np.gradient(cx)
|
|
dcy = np.gradient(cy)
|
|
length = np.sqrt(dcx**2 + dcy**2)
|
|
length = np.maximum(length, 1e-10)
|
|
nx = -dcy / length
|
|
ny = dcx / length
|
|
|
|
offsets = np.linspace(-(thickness - 1) / 2, (thickness - 1) / 2, thickness)
|
|
result = np.zeros((thickness, n_samples))
|
|
for i, off in enumerate(offsets):
|
|
sx = cx + off * nx
|
|
sy = cy + off * ny
|
|
result[i] = map_coordinates(data, [sy, sx], order=1, mode='nearest')
|
|
|
|
total_length = 0.0
|
|
for i in range(1, len(cx)):
|
|
dx_phys = (cx[i] - cx[i - 1]) * field.dx
|
|
dy_phys = (cy[i] - cy[i - 1]) * field.dy
|
|
total_length += np.sqrt(dx_phys**2 + dy_phys**2)
|
|
|
|
center_values = map_coordinates(data, [cy, cx], order=1, mode='nearest')
|
|
profile = LineData(
|
|
data=center_values,
|
|
x_axis=np.linspace(0.0, total_length, n_samples),
|
|
x_unit=field.si_unit_xy,
|
|
y_unit=field.si_unit_z,
|
|
)
|
|
|
|
straightened = field.replace(
|
|
data=result, xreal=total_length,
|
|
yreal=thickness * max(field.dx, field.dy),
|
|
)
|
|
return (straightened, profile)
|