From d35cdd69712d07af12e352366ee02cd72dffae05 Mon Sep 17 00:00:00 2001 From: matei jordache Date: Thu, 16 Apr 2026 22:41:56 -0700 Subject: [PATCH] fix perspective correction --- backend/nodes/perspective_correction.py | 98 ++++++++++---- docs/nodes/Perspective Correction.md | 23 ++-- frontend/src/CustomNode.tsx | 12 ++ frontend/src/PerspectiveOverlay.tsx | 150 +++++++++++++++++++++ frontend/src/styles.css | 67 +++++++++ frontend/src/types.ts | 2 + frontend/src/workflowCapture.ts | 1 + tests/node_tests/perspective_correction.py | 54 ++++++++ 8 files changed, 365 insertions(+), 42 deletions(-) create mode 100644 frontend/src/PerspectiveOverlay.tsx diff --git a/backend/nodes/perspective_correction.py b/backend/nodes/perspective_correction.py index f9e8eae..cef6edf 100644 --- a/backend/nodes/perspective_correction.py +++ b/backend/nodes/perspective_correction.py @@ -6,25 +6,34 @@ import numpy as np from scipy.ndimage import map_coordinates from backend.node_registry import register_node -from backend.data_types import DataField +from backend.data_types import DataField, datafield_to_uint8, encode_preview +from backend.execution_context import emit_overlay @register_node(display_name="Perspective Correction") class PerspectiveCorrection: + _CUSTOM_PREVIEW = True + @classmethod def INPUT_TYPES(cls): return { "required": { "field": ("DATA_FIELD",), - "top_left_x": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01}), - "top_left_y": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01}), - "top_right_x": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01}), - "top_right_y": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01}), - "bottom_left_x": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01}), - "bottom_left_y": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01}), - "bottom_right_x": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01}), - "bottom_right_y": ("FLOAT", {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01}), - } + "top_left_x": ("FLOAT", {"default": 0.1, "min": -1.0, "max": 1.0, "step": 0.01, "hidden": True}), + "top_left_y": ("FLOAT", {"default": 0.1, "min": -1.0, "max": 1.0, "step": 0.01, "hidden": True}), + "top_right_x": ("FLOAT", {"default": -0.1, "min": -1.0, "max": 1.0, "step": 0.01, "hidden": True}), + "top_right_y": ("FLOAT", {"default": 0.1, "min": -1.0, "max": 1.0, "step": 0.01, "hidden": True}), + "bottom_left_x": ("FLOAT", {"default": 0.1, "min": -1.0, "max": 1.0, "step": 0.01, "hidden": True}), + "bottom_left_y": ("FLOAT", {"default": -0.1, "min": -1.0, "max": 1.0, "step": 0.01, "hidden": True}), + "bottom_right_x": ("FLOAT", {"default": -0.1, "min": -1.0, "max": 1.0, "step": 0.01, "hidden": True}), + "bottom_right_y": ("FLOAT", {"default": -0.1, "min": -1.0, "max": 1.0, "step": 0.01, "hidden": True}), + }, + "optional": { + "top_left": ("COORD",), + "top_right": ("COORD",), + "bottom_left": ("COORD",), + "bottom_right": ("COORD",), + }, } OUTPUTS = ( @@ -33,9 +42,8 @@ class PerspectiveCorrection: FUNCTION = "process" DESCRIPTION = ( - "Fix perspective distortion by specifying corner offsets. Each corner " - "can be shifted by a fractional amount (relative to image size) to " - "define the distorted quadrilateral. The image is then warped back to " + "Fix perspective distortion by dragging corner handles. Each corner " + "offset defines a distorted quadrilateral that is warped back to " "a rectangle." ) @@ -45,11 +53,23 @@ class PerspectiveCorrection: top_left_x: float, top_left_y: float, top_right_x: float, top_right_y: float, bottom_left_x: float, bottom_left_y: float, - bottom_right_x: float, bottom_right_y: float) -> tuple: + bottom_right_x: float, bottom_right_y: float, + top_left: tuple[float, float] | None = None, + top_right: tuple[float, float] | None = None, + bottom_left: tuple[float, float] | None = None, + bottom_right: tuple[float, float] | None = None) -> tuple: + if top_left is not None: + top_left_x, top_left_y = float(top_left[0]), float(top_left[1]) + if top_right is not None: + top_right_x, top_right_y = float(top_right[0]), float(top_right[1]) + if bottom_left is not None: + bottom_left_x, bottom_left_y = float(bottom_left[0]), float(bottom_left[1]) + if bottom_right is not None: + bottom_right_x, bottom_right_y = float(bottom_right[0]), float(bottom_right[1]) + data = np.asarray(field.data, dtype=np.float64) yres, xres = data.shape - # Source corners (distorted) as fractional offsets from ideal corners src = np.array([ [top_left_y * yres, top_left_x * xres], [top_right_y * yres, top_right_x * xres + (xres - 1)], @@ -57,7 +77,6 @@ class PerspectiveCorrection: [(1 + bottom_right_y) * yres - 1, bottom_right_x * xres + (xres - 1)], ], dtype=np.float64) - # Destination corners (ideal rectangle) dst = np.array([ [0, 0], [0, xres - 1], @@ -65,33 +84,54 @@ class PerspectiveCorrection: [yres - 1, xres - 1], ], dtype=np.float64) - # Solve for perspective transform matrix (3x3) H = _solve_perspective(src, dst) - # Apply inverse warp yy, xx = np.mgrid[:yres, :xres] - coords = np.stack([yy.ravel(), xx.ravel(), np.ones(yres * xres)]) + coords = np.stack([xx.ravel(), yy.ravel(), np.ones(yres * xres)]) src_coords = H @ coords src_coords /= src_coords[2:3, :] - sy = src_coords[0].reshape(yres, xres) - sx = src_coords[1].reshape(yres, xres) + sx = src_coords[0].reshape(yres, xres) + sy = src_coords[1].reshape(yres, xres) result = map_coordinates(data, [sy, sx], order=1, mode='nearest') - return (field.replace(data=result),) + corrected = field.replace(data=result) + + source_rgb = datafield_to_uint8(field, field.colormap) + corrected_rgb = datafield_to_uint8(corrected, corrected.colormap) + + corners = [ + {"x": float(top_left_x), "y": float(top_left_y)}, + {"x": float(top_right_x), "y": float(top_right_y)}, + {"x": float(bottom_left_x), "y": float(bottom_left_y)}, + {"x": float(bottom_right_x), "y": float(bottom_right_y)}, + ] + + emit_overlay({ + "kind": "perspective", + "section_title": "Perspective", + "image": encode_preview(source_rgb), + "corrected_image": encode_preview(corrected_rgb), + "corners": corners, + }) + + return (corrected,) def _solve_perspective(src: np.ndarray, dst: np.ndarray) -> np.ndarray: - """Solve for 3x3 perspective matrix mapping dst -> src (for inverse warp).""" + """Solve for 3x3 perspective matrix mapping dst -> src (for inverse warp). + + Coordinates are (col, row) — the matrix is applied to [col, row, 1] vectors. + """ n = len(src) A = np.zeros((2 * n, 8)) b = np.zeros(2 * n) for i in range(n): - dy, dx = dst[i] - sy, sx = src[i] - A[2 * i] = [dx, dy, 1, 0, 0, 0, -sx * dx, -sx * dy] - A[2 * i + 1] = [0, 0, 0, dx, dy, 1, -sy * dx, -sy * dy] - b[2 * i] = sx - b[2 * i + 1] = sy + dr, dc = dst[i] # dest row, col + sr, sc = src[i] # src row, col + A[2 * i] = [dc, dr, 1, 0, 0, 0, -sc * dc, -sc * dr] + A[2 * i + 1] = [0, 0, 0, dc, dr, 1, -sr * dc, -sr * dr] + b[2 * i] = sc + b[2 * i + 1] = sr h, _, _, _ = np.linalg.lstsq(A, b, rcond=None) H = np.array([[h[0], h[1], h[2]], [h[3], h[4], h[5]], diff --git a/docs/nodes/Perspective Correction.md b/docs/nodes/Perspective Correction.md index b734e47..f397b1e 100644 --- a/docs/nodes/Perspective Correction.md +++ b/docs/nodes/Perspective Correction.md @@ -1,12 +1,16 @@ # Perspective Correction -Fix perspective distortion in a DATA_FIELD via a projective (homography) transform. Each corner can be shifted by a fractional offset to map a distorted quadrilateral back to a rectangle. Equivalent to Gwyddion's `correct_perspective.c` module. +Fix perspective distortion in a DATA_FIELD via a projective (homography) transform. Each corner can be shifted by a fractional offset to map a distorted quadrilateral back to a rectangle. ## Inputs | Name | Type | Required | Description | |------|------|----------|-------------| | field | DATA_FIELD | Yes | Input field with perspective distortion | +| top_left | COORD | No | Override top-left corner offset (x, y) | +| top_right | COORD | No | Override top-right corner offset (x, y) | +| bottom_left | COORD | No | Override bottom-left corner offset (x, y) | +| bottom_right | COORD | No | Override bottom-right corner offset (x, y) | ## Outputs @@ -14,22 +18,15 @@ Fix perspective distortion in a DATA_FIELD via a projective (homography) transfo |------|------|-------------| | corrected | DATA_FIELD | Perspective-corrected field | -## Controls +## Interactive preview -| Name | Type | Default | Description | -|------|------|---------|-------------| -| top_left_x | FLOAT | 0.0 | Horizontal offset of the top-left corner as a fraction of image width (-0.5-0.5) | -| top_left_y | FLOAT | 0.0 | Vertical offset of the top-left corner as a fraction of image height (-0.5-0.5) | -| top_right_x | FLOAT | 0.0 | Horizontal offset of the top-right corner as a fraction of image width (-0.5-0.5) | -| top_right_y | FLOAT | 0.0 | Vertical offset of the top-right corner as a fraction of image height (-0.5-0.5) | -| bottom_left_x | FLOAT | 0.0 | Horizontal offset of the bottom-left corner as a fraction of image width (-0.5-0.5) | -| bottom_left_y | FLOAT | 0.0 | Vertical offset of the bottom-left corner as a fraction of image height (-0.5-0.5) | -| bottom_right_x | FLOAT | 0.0 | Horizontal offset of the bottom-right corner as a fraction of image width (-0.5-0.5) | -| bottom_right_y | FLOAT | 0.0 | Vertical offset of the bottom-right corner as a fraction of image height (-0.5-0.5) | +The preview shows the source image with a draggable quadrilateral overlay. Drag any corner handle to adjust the perspective correction. Use the Source/Corrected tabs to switch between the input image (with handles) and the corrected result. + +Corner positions can also be set by connecting Coordinate nodes to the optional COORD inputs, which override the handle-driven values. ## Notes - All offsets are given as fractions of the image dimensions (0.0 = no shift, 0.1 = 10% shift). Positive x shifts right, positive y shifts down. - The transform uses bilinear interpolation to resample pixel values at non-integer locations. - For trapezoidal distortions (common in tilted AFM scans), typically only two corners need adjustment. -- Set all offsets to 0.0 to pass the field through unchanged. +- When all offsets are zero (default), the field passes through unchanged. diff --git a/frontend/src/CustomNode.tsx b/frontend/src/CustomNode.tsx index 2f39b4b..80587d4 100644 --- a/frontend/src/CustomNode.tsx +++ b/frontend/src/CustomNode.tsx @@ -15,6 +15,7 @@ const ThresholdHistogram = lazy(() => import('./ThresholdHistogram')); const RadialProfileOverlay = lazy(() => import('./RadialProfileOverlay')); const StraightenPathOverlay = lazy(() => import('./StraightenPathOverlay')); const MultiProfileOverlay = lazy(() => import('./MultiProfileOverlay')); +const PerspectiveOverlay = lazy(() => import('./PerspectiveOverlay')); import TextNoteNode from './TextNoteNode'; @@ -1202,6 +1203,7 @@ function CustomNode({ id, data }: { id: string; data: NodeData }) { || data.overlay.kind === 'radial_profile' || data.overlay.kind === 'straighten_path' || data.overlay.kind === 'multi_profile' + || data.overlay.kind === 'perspective' ); const hidePreviewForInteractiveMask = data.overlay?.kind === 'mask_paint' || data.overlay?.kind === 'markup'; const overlayTitle = data.overlay?.section_title @@ -1223,6 +1225,8 @@ function CustomNode({ id, data }: { id: string; data: NodeData }) { ? 'Path' : data.overlay?.kind === 'multi_profile' ? 'Preview' + : data.overlay?.kind === 'perspective' + ? 'Perspective' : 'Cross Section'); const headerMeta = (() => { if (data.className === 'Folder') { @@ -1597,6 +1601,14 @@ function CustomNode({ id, data }: { id: string; data: NodeData }) { nodeId={id} onWidgetChange={ctx!.onWidgetChange} /> + ) : data.overlay!.kind === 'perspective' ? ( + } + nodeId={id} + onWidgetChange={ctx!.onWidgetChange} + /> ) : data.overlay!.kind === 'angle_measure' ? ( = { + top_left: { ax: 0, ay: 0 }, + top_right: { ax: 1, ay: 0 }, + bottom_left: { ax: 0, ay: 1 }, + bottom_right: { ax: 1, ay: 1 }, +}; + +interface Corner { x: number; y: number } + +interface Props { + image: string; + correctedImage: string; + corners: Corner[]; + nodeId: string; + onWidgetChange: (nodeId: string, name: string, value: unknown) => void; +} + +function cornerToPercent(corner: Corner, name: CornerName) { + const anchor = CORNER_ANCHORS[name]; + return { + left: (anchor.ax + corner.x) * 100, + top: (anchor.ay + corner.y) * 100, + }; +} + +function cornersKey(c: Corner[]): string { + return c.map((p) => `${p.x},${p.y}`).join(';'); +} + +export default function PerspectiveOverlay({ + image, correctedImage, corners, nodeId, onWidgetChange, +}: Props) { + const containerRef = useRef(null); + const draggingRef = useRef(null); + const [draft, setDraft] = useState(null); + const pendingCommitRef = useRef(null); + const [showCorrected, setShowCorrected] = useState(false); + + useEffect(() => { + if (pendingCommitRef.current && cornersKey(corners) === pendingCommitRef.current) { + pendingCommitRef.current = null; + setDraft(null); + } + }, [corners]); + + const liveCorners = draft ?? corners; + + const onPointerDown = useCallback((corner: CornerName) => (e: React.PointerEvent) => { + e.stopPropagation(); + e.preventDefault(); + (e.target as HTMLElement).setPointerCapture(e.pointerId); + draggingRef.current = corner; + setDraft([...liveCorners]); + }, [liveCorners]); + + const onPointerMove = useCallback((e: React.PointerEvent) => { + const name = draggingRef.current; + if (!name || !containerRef.current) return; + const { fx, fy } = pointerToFraction(e, containerRef.current); + const anchor = CORNER_ANCHORS[name]; + const cx = Math.max(-1, Math.min(1, parseFloat((fx - anchor.ax).toFixed(3)))); + const cy = Math.max(-1, Math.min(1, parseFloat((fy - anchor.ay).toFixed(3)))); + const idx = CORNER_NAMES.indexOf(name); + setDraft((prev) => { + if (!prev) return prev; + const next = [...prev]; + next[idx] = { x: cx, y: cy }; + return next; + }); + }, []); + + const onPointerUp = useCallback(() => { + const name = draggingRef.current; + if (!name || !draft) { + draggingRef.current = null; + return; + } + draggingRef.current = null; + pendingCommitRef.current = cornersKey(draft); + for (let i = 0; i < CORNER_NAMES.length; i++) { + onWidgetChange(nodeId, `${CORNER_NAMES[i]}_x`, draft[i].x); + onWidgetChange(nodeId, `${CORNER_NAMES[i]}_y`, draft[i].y); + } + }, [draft, nodeId, onWidgetChange]); + + const positions = CORNER_NAMES.map((name, i) => cornerToPercent(liveCorners[i] || { x: 0, y: 0 }, name)); + const quadPoints = `${positions[0].left},${positions[0].top} ${positions[1].left},${positions[1].top} ${positions[3].left},${positions[3].top} ${positions[2].left},${positions[2].top}`; + + return ( +
+
+ + +
+ + {showCorrected ? ( +
+ corrected +
+ ) : ( +
+ source + + + + + + {CORNER_NAMES.map((name, i) => ( +
+ ))} +
+ )} +
+ ); +} diff --git a/frontend/src/styles.css b/frontend/src/styles.css index eafd791..4ad69ae 100644 --- a/frontend/src/styles.css +++ b/frontend/src/styles.css @@ -1801,6 +1801,73 @@ html, body, #root { pointer-events: none; } +/* ── Perspective correction overlay ──────────────────────────────────── */ +.perspective-overlay-wrap { + display: flex; + flex-direction: column; +} +.perspective-tab-bar { + display: flex; + gap: 1px; + background: var(--border-default); + border-bottom: 1px solid var(--border-default); +} +.perspective-tab { + flex: 1; + padding: 4px 0; + font-size: 11px; + font-weight: 500; + border: none; + background: var(--bg-deep); + color: var(--text-secondary); + cursor: pointer; + transition: background 0.1s, color 0.1s; +} +.perspective-tab:hover { + background: var(--bg-panel); +} +.perspective-tab.active { + background: var(--bg-panel); + color: var(--text-primary); +} +.perspective-overlay { + position: relative; + user-select: none; + touch-action: none; + overflow: hidden; +} +.perspective-overlay img { + width: 100%; + display: block; +} +.perspective-quad { + position: absolute; + inset: 0; + width: 100%; + height: 100%; + pointer-events: none; +} +.perspective-handle { + position: absolute; + width: 12px; + height: 12px; + border-radius: 50%; + background: var(--selection, #3b82f6); + border: 2px solid #fff; + box-shadow: 0 0 4px rgba(0, 0, 0, 0.5); + transform: translate(-50%, -50%); + cursor: grab; + z-index: 2; +} +.perspective-handle:hover, +.perspective-handle.dragging { + cursor: grabbing; + transform: translate(-50%, -50%) scale(1.2); +} +.is-panning .perspective-overlay { + pointer-events: none; +} + .angle-overlay { --angle-line-color: #ff9800; --angle-arc-color: rgb(255, 166, 77); diff --git a/frontend/src/types.ts b/frontend/src/types.ts index a6bed64..4d2b496 100644 --- a/frontend/src/types.ts +++ b/frontend/src/types.ts @@ -82,6 +82,8 @@ export interface OverlayData { row?: number; direction?: 'horizontal' | 'vertical'; max_index?: number; + corrected_image?: string; + corners?: Array<{ x: number; y: number }>; section_title?: string; line?: number[]; shape?: string; diff --git a/frontend/src/workflowCapture.ts b/frontend/src/workflowCapture.ts index 9eeaafa..046b6e4 100644 --- a/frontend/src/workflowCapture.ts +++ b/frontend/src/workflowCapture.ts @@ -14,6 +14,7 @@ export const OVERLAY_CAPTURE_SELECTORS = [ '.radial-overlay', // RadialProfileOverlay '.straighten-overlay', // StraightenPathOverlay '.multiprofile-overlay', // MultiProfileOverlay + '.perspective-overlay', // PerspectiveOverlay ]; function encodeBase64(bytes: Uint8Array) { diff --git a/tests/node_tests/perspective_correction.py b/tests/node_tests/perspective_correction.py index 3a1ac3c..06be04a 100644 --- a/tests/node_tests/perspective_correction.py +++ b/tests/node_tests/perspective_correction.py @@ -51,3 +51,57 @@ def test_output_shape(): bottom_right_x=0.0, bottom_right_y=0.0, ) assert result.data.shape == (48, 96) + + +def test_emits_perspective_overlay(): + from backend.execution_context import active_node, execution_callbacks + from backend.nodes.perspective_correction import PerspectiveCorrection + + node = PerspectiveCorrection() + field = make_field(shape=(64, 64)) + + overlays = [] + with execution_callbacks(overlay=lambda nid, d: overlays.append(d)), active_node("test"): + node.process( + field, + top_left_x=0.05, top_left_y=0.05, + top_right_x=-0.05, top_right_y=0.05, + bottom_left_x=0.05, bottom_left_y=-0.05, + bottom_right_x=-0.05, bottom_right_y=-0.05, + ) + + assert len(overlays) == 1 + ov = overlays[0] + assert ov["kind"] == "perspective" + assert ov["section_title"] == "Perspective" + assert ov["image"].startswith("data:image/png;base64,") + assert ov["corrected_image"].startswith("data:image/png;base64,") + assert len(ov["corners"]) == 4 + assert ov["corners"][0] == {"x": 0.05, "y": 0.05} + assert ov["corners"][3] == {"x": -0.05, "y": -0.05} + + +def test_coord_input_overrides_floats(): + from backend.nodes.perspective_correction import PerspectiveCorrection + + node = PerspectiveCorrection() + field = make_field(shape=(64, 64)) + + result_floats, = node.process( + field, + top_left_x=0.1, top_left_y=0.1, + top_right_x=0.0, top_right_y=0.0, + bottom_left_x=0.0, bottom_left_y=0.0, + bottom_right_x=0.0, bottom_right_y=0.0, + ) + + result_coord, = node.process( + field, + top_left_x=0.0, top_left_y=0.0, + top_right_x=0.0, top_right_y=0.0, + bottom_left_x=0.0, bottom_left_y=0.0, + bottom_right_x=0.0, bottom_right_y=0.0, + top_left=(0.1, 0.1), + ) + + assert np.allclose(result_floats.data, result_coord.data)