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"}), "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)