223 lines
7.7 KiB
Python
223 lines
7.7 KiB
Python
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
|
|
from backend.data_types import (
|
|
DataField,
|
|
ImageData,
|
|
MeasureTable,
|
|
_apply_angle_measure_overlay,
|
|
encode_preview,
|
|
image_metadata,
|
|
image_to_uint8,
|
|
render_datafield_preview,
|
|
)
|
|
from backend.execution_context import emit_overlay, emit_table
|
|
from backend.node_registry import register_node
|
|
from backend.nodes.helpers import _normalize_markup_color
|
|
|
|
ANGLE_DEFAULT_COLOR = "#ff9800"
|
|
|
|
|
|
def _clamp01(value: float) -> float:
|
|
return float(np.clip(value, 0.0, 1.0))
|
|
|
|
|
|
def _sanitize_stroke_width(value: float | None, default: float = 1.35) -> float:
|
|
try:
|
|
numeric = float(value)
|
|
except (TypeError, ValueError):
|
|
numeric = default
|
|
if not np.isfinite(numeric):
|
|
numeric = default
|
|
return float(np.clip(numeric, 0.35, 6.0))
|
|
|
|
|
|
def _angle_metrics(ax: float, ay: float, vx: float, vy: float, bx: float, by: float) -> tuple[float, float, float]:
|
|
va = np.array([ax - vx, ay - vy], dtype=np.float64)
|
|
vb = np.array([bx - vx, by - vy], dtype=np.float64)
|
|
len_a = float(np.hypot(va[0], va[1]))
|
|
len_b = float(np.hypot(vb[0], vb[1]))
|
|
|
|
if len_a <= 1e-12 or len_b <= 1e-12:
|
|
return (0.0, len_a, len_b)
|
|
|
|
cos_theta = float(np.dot(va, vb) / (len_a * len_b))
|
|
cos_theta = float(np.clip(cos_theta, -1.0, 1.0))
|
|
angle_deg = float(np.degrees(np.arccos(cos_theta)))
|
|
return (angle_deg, len_a, len_b)
|
|
|
|
|
|
def _angle_overlay_spec(
|
|
x1: float,
|
|
y1: float,
|
|
xm: float,
|
|
ym: float,
|
|
x2: float,
|
|
y2: float,
|
|
angle_deg: float,
|
|
label_dx: float,
|
|
label_dy: float,
|
|
stroke_width: float,
|
|
color: str,
|
|
) -> dict[str, float | str]:
|
|
return {
|
|
"kind": "angle_measure",
|
|
"x1": x1,
|
|
"y1": y1,
|
|
"xm": xm,
|
|
"ym": ym,
|
|
"x2": x2,
|
|
"y2": y2,
|
|
"angle_deg": float(angle_deg),
|
|
"label_dx": float(label_dx),
|
|
"label_dy": float(label_dy),
|
|
"stroke_width": float(stroke_width),
|
|
"color": str(color),
|
|
}
|
|
|
|
|
|
@register_node(display_name="Angle Measure")
|
|
class AngleMeasure:
|
|
_CUSTOM_PREVIEW = True
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"input": ("ANNOTATION_SOURCE", {
|
|
"label": "Input",
|
|
"accepted_types": ["DATA_FIELD", "IMAGE"],
|
|
}),
|
|
"color": ("STRING", {"default": ANGLE_DEFAULT_COLOR, "color_picker": True}),
|
|
"stroke_width": ("FLOAT", {
|
|
"default": 1.35,
|
|
"min": 0.35,
|
|
"max": 6.0,
|
|
"step": 0.05,
|
|
"label": "stroke width",
|
|
}),
|
|
"x1": ("FLOAT", {"default": 0.22, "min": 0.0, "max": 1.0, "step": 0.01, "hidden": True}),
|
|
"y1": ("FLOAT", {"default": 0.72, "min": 0.0, "max": 1.0, "step": 0.01, "hidden": True}),
|
|
"xm": ("FLOAT", {"default": 0.50, "min": 0.0, "max": 1.0, "step": 0.01, "hidden": True}),
|
|
"ym": ("FLOAT", {"default": 0.50, "min": 0.0, "max": 1.0, "step": 0.01, "hidden": True}),
|
|
"x2": ("FLOAT", {"default": 0.78, "min": 0.0, "max": 1.0, "step": 0.01, "hidden": True}),
|
|
"y2": ("FLOAT", {"default": 0.28, "min": 0.0, "max": 1.0, "step": 0.01, "hidden": True}),
|
|
"label_dx": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01, "hidden": True}),
|
|
"label_dy": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01, "hidden": True}),
|
|
},
|
|
"optional": {
|
|
"line_thickness": ("FLOAT", {"hidden": True}),
|
|
"line_thickness_input": ("FLOAT", {"hidden": True}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("ANNOTATION_SOURCE", "MEASURE_TABLE")
|
|
RETURN_NAMES = ("output", "measurements")
|
|
FUNCTION = "process"
|
|
|
|
DESCRIPTION = (
|
|
"Measure the included angle between two draggable line segments over a DATA_FIELD or IMAGE. "
|
|
"Drag either endpoint to change that arm, drag the middle joint to move the whole widget, "
|
|
"drag the angle label to reposition it, choose the overlay color, and adjust stroke width "
|
|
"with the widget or its FLOAT socket."
|
|
)
|
|
|
|
def process(
|
|
self,
|
|
input,
|
|
color: str,
|
|
stroke_width: float,
|
|
x1: float,
|
|
y1: float,
|
|
xm: float,
|
|
ym: float,
|
|
x2: float,
|
|
y2: float,
|
|
label_dx: float,
|
|
label_dy: float,
|
|
line_thickness: float | None = None,
|
|
line_thickness_input: float | None = None,
|
|
) -> tuple:
|
|
x1 = _clamp01(x1)
|
|
y1 = _clamp01(y1)
|
|
xm = _clamp01(xm)
|
|
ym = _clamp01(ym)
|
|
x2 = _clamp01(x2)
|
|
y2 = _clamp01(y2)
|
|
label_dx = float(np.clip(label_dx, -1.0, 1.0))
|
|
label_dy = float(np.clip(label_dy, -1.0, 1.0))
|
|
resolved_color = _normalize_markup_color(color, default=ANGLE_DEFAULT_COLOR)
|
|
legacy_stroke_width = line_thickness_input if line_thickness_input is not None else line_thickness
|
|
resolved_stroke_width = _sanitize_stroke_width(
|
|
legacy_stroke_width if legacy_stroke_width is not None else stroke_width,
|
|
)
|
|
|
|
if isinstance(input, DataField):
|
|
preview_base = render_datafield_preview(input, input.colormap)
|
|
ax = float(input.xoff + x1 * input.xreal)
|
|
ay = float(input.yoff + y1 * input.yreal)
|
|
vx = float(input.xoff + xm * input.xreal)
|
|
vy = float(input.yoff + ym * input.yreal)
|
|
bx = float(input.xoff + x2 * input.xreal)
|
|
by = float(input.yoff + y2 * input.yreal)
|
|
length_unit = input.si_unit_xy
|
|
else:
|
|
preview_base = image_to_uint8(input)
|
|
height, width = preview_base.shape[:2]
|
|
ax = float(x1 * max(width - 1, 0))
|
|
ay = float(y1 * max(height - 1, 0))
|
|
vx = float(xm * max(width - 1, 0))
|
|
vy = float(ym * max(height - 1, 0))
|
|
bx = float(x2 * max(width - 1, 0))
|
|
by = float(y2 * max(height - 1, 0))
|
|
length_unit = "px"
|
|
|
|
angle_deg, length_a, length_b = _angle_metrics(ax, ay, vx, vy, bx, by)
|
|
angle_overlay = _angle_overlay_spec(
|
|
x1=x1,
|
|
y1=y1,
|
|
xm=xm,
|
|
ym=ym,
|
|
x2=x2,
|
|
y2=y2,
|
|
angle_deg=angle_deg,
|
|
label_dx=label_dx,
|
|
label_dy=label_dy,
|
|
stroke_width=resolved_stroke_width,
|
|
color=resolved_color,
|
|
)
|
|
table = MeasureTable([
|
|
{"quantity": "Angle", "value": angle_deg, "unit": "deg"},
|
|
{"quantity": "Arm A length", "value": length_a, "unit": length_unit},
|
|
{"quantity": "Arm B length", "value": length_b, "unit": length_unit},
|
|
{"quantity": "Arm A x", "value": ax, "unit": length_unit},
|
|
{"quantity": "Arm A y", "value": ay, "unit": length_unit},
|
|
{"quantity": "Vertex x", "value": vx, "unit": length_unit},
|
|
{"quantity": "Vertex y", "value": vy, "unit": length_unit},
|
|
{"quantity": "Arm B x", "value": bx, "unit": length_unit},
|
|
{"quantity": "Arm B y", "value": by, "unit": length_unit},
|
|
])
|
|
|
|
if isinstance(input, DataField):
|
|
output = input.replace(
|
|
overlays=[
|
|
*input.overlays,
|
|
angle_overlay,
|
|
],
|
|
)
|
|
else:
|
|
output = ImageData(
|
|
_apply_angle_measure_overlay(preview_base, None, angle_overlay),
|
|
metadata=image_metadata(input),
|
|
)
|
|
|
|
emit_overlay({
|
|
"section_title": "Angle",
|
|
"image": encode_preview(preview_base),
|
|
**angle_overlay,
|
|
})
|
|
emit_table(table)
|
|
|
|
return (output, table)
|