103 lines
3.7 KiB
Python
103 lines
3.7 KiB
Python
from __future__ import annotations
|
|
import numpy as np
|
|
from backend.node_registry import register_node
|
|
from backend.data_types import DataField, LineData, datafield_to_uint8, encode_preview
|
|
from backend.nodes.helpers import _extend_to_edges
|
|
|
|
|
|
@register_node(display_name="Cross Section")
|
|
class CrossSection:
|
|
"""Extract a 1-D height profile along an arbitrary line across the image."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"field": ("DATA_FIELD",),
|
|
"x1": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01, "hidden": True}),
|
|
"y1": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, "hidden": True}),
|
|
"x2": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.01, "hidden": True}),
|
|
"y2": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, "hidden": True}),
|
|
"extend": (["none", "to_edges"],),
|
|
"n_samples": ("INT", {"default": 0, "min": 0, "max": 4096, "step": 1}),
|
|
},
|
|
"optional": {
|
|
"marker_pair": ("COORDPAIR", {"label": "marker pair"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("LINE", "COORDPAIR",)
|
|
RETURN_NAMES = ("profile", "marker pair",)
|
|
FUNCTION = "process"
|
|
|
|
DESCRIPTION = (
|
|
"Extract a cross-section profile along a line between two points. "
|
|
"Drag the markers on the image to set the line endpoints. "
|
|
"Equivalent to gwy_data_field_get_profile."
|
|
)
|
|
|
|
_broadcast_overlay_fn = None
|
|
_current_node_id: str = ""
|
|
|
|
def process(
|
|
self, field: DataField,
|
|
x1: float, y1: float, x2: float, y2: float,
|
|
extend: str, n_samples: int,
|
|
marker_pair=None,
|
|
) -> tuple:
|
|
from scipy.ndimage import map_coordinates
|
|
|
|
if marker_pair is not None:
|
|
(x1, y1), (x2, y2) = marker_pair
|
|
|
|
marker_x1, marker_y1 = float(x1), float(y1)
|
|
marker_x2, marker_y2 = float(x2), float(y2)
|
|
|
|
xres, yres = field.xres, field.yres
|
|
|
|
if extend == "to_edges":
|
|
x1, y1, x2, y2 = _extend_to_edges(
|
|
float(x1), float(y1), float(x2), float(y2),
|
|
)
|
|
|
|
px1, py1 = float(x1) * (xres - 1), float(y1) * (yres - 1)
|
|
px2, py2 = float(x2) * (xres - 1), float(y2) * (yres - 1)
|
|
|
|
line_len_px = np.hypot(px2 - px1, py2 - py1)
|
|
if n_samples <= 0:
|
|
n_samples = max(2, int(np.ceil(line_len_px)))
|
|
|
|
t = np.linspace(0, 1, n_samples)
|
|
coords_y = py1 + t * (py2 - py1)
|
|
coords_x = px1 + t * (px2 - px1)
|
|
|
|
profile = map_coordinates(field.data, [coords_y, coords_x], order=3, mode="nearest")
|
|
|
|
if CrossSection._broadcast_overlay_fn is not None:
|
|
image_uri = encode_preview(datafield_to_uint8(field, field.colormap))
|
|
|
|
CrossSection._broadcast_overlay_fn(
|
|
CrossSection._current_node_id,
|
|
{
|
|
"image": image_uri,
|
|
"x1": marker_x1, "y1": marker_y1,
|
|
"x2": marker_x2, "y2": marker_y2,
|
|
"a_locked": marker_pair is not None,
|
|
"b_locked": marker_pair is not None,
|
|
},
|
|
)
|
|
|
|
dx_real = (x2 - x1) * field.xreal
|
|
dy_real = (y2 - y1) * field.yreal
|
|
distance_axis = np.linspace(0.0, float(np.hypot(dx_real, dy_real)), n_samples, dtype=np.float64)
|
|
|
|
return (
|
|
LineData(
|
|
data=profile.astype(np.float64),
|
|
x_axis=distance_axis,
|
|
x_unit=field.si_unit_xy,
|
|
y_unit=field.si_unit_z,
|
|
),
|
|
((marker_x1, marker_y1), (marker_x2, marker_y2)),
|
|
)
|