fix perspective correction
This commit is contained in:
@@ -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]],
|
||||
|
||||
Reference in New Issue
Block a user