318 lines
13 KiB
Python
318 lines
13 KiB
Python
# pixel_drift_fix_node.py
|
|
#
|
|
# Pixel-drift correction for ComfyUI.
|
|
# - flat_4_points : single global homography via SIFT + RANSAC, optional ECC refine.
|
|
# ~2-5s on 4MP (downscale=4), ~1-2s (downscale=8).
|
|
# - mesh : homography + grid-interpolated residual displacement.
|
|
# Visually equivalent to the old PiecewiseAffineTransform path,
|
|
# but ~30-60x faster (400 pts: ~2s, 10000 pts: ~10-15s on 4MP).
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
|
|
try:
|
|
from scipy.interpolate import griddata as _griddata
|
|
_HAS_SCIPY = True
|
|
except ImportError:
|
|
_HAS_SCIPY = False
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tensor <-> OpenCV helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _to_bgr_batch(t: torch.Tensor):
|
|
if t.ndim == 3:
|
|
t = t.unsqueeze(0)
|
|
arr = (t.detach().cpu().float().numpy() * 255.0).clip(0, 255).astype(np.uint8)
|
|
return [cv2.cvtColor(arr[i], cv2.COLOR_RGB2BGR) for i in range(arr.shape[0])]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# SIFT + RANSAC (returns inlier correspondences at FULL resolution + homography)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _sift_match(orig_bgr, mod_bgr, max_features, downscale,
|
|
contrast=0.01, edge=20, sigma=1.6, octave=3, lowe=0.80,
|
|
ransac_thresh=2.0):
|
|
h, w = orig_bgr.shape[:2]
|
|
ds = max(1, int(downscale))
|
|
sw = max(64, (w + ds - 1) // ds)
|
|
sh = max(64, (h + ds - 1) // ds)
|
|
sx, sy = sw / w, sh / h
|
|
|
|
if ds > 1:
|
|
orig_s = cv2.resize(orig_bgr, (sw, sh), interpolation=cv2.INTER_AREA)
|
|
mod_s = cv2.resize(mod_bgr, (sw, sh), interpolation=cv2.INTER_AREA)
|
|
else:
|
|
orig_s, mod_s = orig_bgr, mod_bgr
|
|
|
|
g_orig = cv2.cvtColor(orig_s, cv2.COLOR_BGR2GRAY)
|
|
g_mod = cv2.cvtColor(mod_s, cv2.COLOR_BGR2GRAY)
|
|
|
|
sift = cv2.SIFT_create(
|
|
nfeatures=max_features, nOctaveLayers=octave,
|
|
contrastThreshold=contrast, edgeThreshold=edge, sigma=sigma,
|
|
)
|
|
kp_mod, des_mod = sift.detectAndCompute(g_mod, None)
|
|
kp_orig, des_orig = sift.detectAndCompute(g_orig, None)
|
|
|
|
if (des_mod is None or des_orig is None
|
|
or len(kp_mod) < 4 or len(kp_orig) < 4):
|
|
return None, None, None
|
|
|
|
bf = cv2.BFMatcher(cv2.NORM_L2)
|
|
matches = bf.knnMatch(des_mod, des_orig, k=2)
|
|
good = [m for m, n in matches if m.distance < lowe * n.distance]
|
|
if len(good) < 8:
|
|
return None, None, None
|
|
|
|
pts_mod_s = np.float32([kp_mod[m.queryIdx].pt for m in good])
|
|
pts_orig_s = np.float32([kp_orig[m.trainIdx].pt for m in good])
|
|
|
|
M_s, mask = cv2.findHomography(pts_mod_s, pts_orig_s, cv2.RANSAC, ransac_thresh)
|
|
if M_s is None or mask is None:
|
|
return None, None, None
|
|
|
|
inl_mod_s = pts_mod_s[mask.ravel() == 1]
|
|
inl_orig_s = pts_orig_s[mask.ravel() == 1]
|
|
|
|
S = np.array([[sx, 0, 0], [0, sy, 0], [0, 0, 1]], dtype=np.float64)
|
|
S_inv = np.array([[1.0 / sx, 0, 0], [0, 1.0 / sy, 0], [0, 0, 1]], dtype=np.float64)
|
|
M_full = (S_inv @ M_s.astype(np.float64) @ S).astype(np.float32)
|
|
|
|
inl_mod = (inl_mod_s / np.float32([sx, sy])).astype(np.float32)
|
|
inl_orig = (inl_orig_s / np.float32([sx, sy])).astype(np.float32)
|
|
|
|
return inl_mod, inl_orig, M_full
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Optional sub-pixel ECC refinement of the homography
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _refine_with_ecc(orig_bgr, mod_bgr, M_full, downscale):
|
|
h, w = orig_bgr.shape[:2]
|
|
ds = max(1, int(downscale))
|
|
sw = max(64, w // ds)
|
|
sh = max(64, h // ds)
|
|
sx, sy = sw / w, sh / h
|
|
|
|
g_orig = cv2.cvtColor(
|
|
cv2.resize(orig_bgr, (sw, sh), interpolation=cv2.INTER_AREA),
|
|
cv2.COLOR_BGR2GRAY)
|
|
g_mod = cv2.cvtColor(
|
|
cv2.resize(mod_bgr, (sw, sh), interpolation=cv2.INTER_AREA),
|
|
cv2.COLOR_BGR2GRAY)
|
|
|
|
S = np.array([[sx, 0, 0], [0, sy, 0], [0, 0, 1]], dtype=np.float64)
|
|
S_inv = np.array([[1.0 / sx, 0, 0], [0, 1.0 / sy, 0], [0, 0, 1]], dtype=np.float64)
|
|
M_small = S @ M_full.astype(np.float64) @ S_inv
|
|
|
|
try:
|
|
criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 50, 1e-6)
|
|
init = np.linalg.inv(M_small).astype(np.float32)
|
|
cc, M_ecc = cv2.findTransformECC(
|
|
g_orig, g_mod, init, cv2.MOTION_HOMOGRAPHY, criteria, None, 5)
|
|
if cc > 0.3:
|
|
M_ecc_full = (S_inv @ np.linalg.inv(M_ecc.astype(np.float64)) @ S).astype(np.float32)
|
|
corners = np.float32([[[0, 0]], [[w, 0]], [[w, h]], [[0, h]]])
|
|
p1 = cv2.perspectiveTransform(corners, M_full)
|
|
p2 = cv2.perspectiveTransform(corners, M_ecc_full)
|
|
if np.abs(p2 - p1).max() < max(w, h) * 0.05 + 50:
|
|
return M_ecc_full
|
|
except cv2.error:
|
|
pass
|
|
return M_full
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Warp modes
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _warp_flat(mod_bgr, M_full):
|
|
h, w = mod_bgr.shape[:2]
|
|
return cv2.warpPerspective(
|
|
mod_bgr, M_full, (w, h),
|
|
flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)
|
|
|
|
|
|
def _warp_mesh(mod_bgr, M_full, pts_mod, pts_orig, max_mesh_points, step_size=30):
|
|
"""
|
|
Fast piecewise-affine-equivalent mesh warp.
|
|
1. Cap inlier points (plus boundary anchors) to max_mesh_points.
|
|
2. Compute residual = target_mod_pos - homography_predicted_mod_pos
|
|
at each control point.
|
|
3. Interpolate residual displacement to a coarse grid with griddata,
|
|
upsample to full res with cv2.resize.
|
|
4. Final remap = homography displacement + interpolated residual.
|
|
Boundary anchors get zero residual, so mesh smoothly falls off to the
|
|
homography outside the SIFT point cloud.
|
|
"""
|
|
h, w = mod_bgr.shape[:2]
|
|
if not _HAS_SCIPY or pts_mod is None or len(pts_mod) < 4:
|
|
return _warp_flat(mod_bgr, M_full)
|
|
|
|
# --- boundary anchors in orig space, mapped to mod space via M_inv ---
|
|
M_inv = np.linalg.inv(M_full.astype(np.float64))
|
|
edge = []
|
|
for x in range(0, w, step_size):
|
|
edge.append([x, 0]); edge.append([x, h - 1])
|
|
for y in range(0, h, step_size):
|
|
edge.append([0, y]); edge.append([w - 1, y])
|
|
edge = np.array(edge, dtype=np.float32).reshape(-1, 1, 2)
|
|
edge_proj = cv2.perspectiveTransform(edge, M_inv).reshape(-1, 2).astype(np.float32)
|
|
edge_proj[:, 0] = np.clip(edge_proj[:, 0], 0, w - 1)
|
|
edge_proj[:, 1] = np.clip(edge_proj[:, 1], 0, h - 1)
|
|
edge = edge.reshape(-1, 2)
|
|
|
|
all_orig = np.vstack([pts_orig, edge]).astype(np.float32)
|
|
all_mod = np.vstack([pts_mod, edge_proj]).astype(np.float32)
|
|
|
|
# --- cap total control points (matches original max_mesh_points semantics) ---
|
|
if len(all_orig) > max_mesh_points:
|
|
idx = np.linspace(0, len(all_orig) - 1, max_mesh_points, dtype=int)
|
|
all_orig = all_orig[idx]
|
|
all_mod = all_mod[idx]
|
|
|
|
# --- residual relative to homography ---
|
|
base_at_ctrl = cv2.perspectiveTransform(
|
|
all_orig.reshape(-1, 1, 2), M_inv).reshape(-1, 2).astype(np.float32)
|
|
resid = all_mod - base_at_ctrl # (N, 2)
|
|
|
|
# --- interpolate residual on coarse grid, then upsample ---
|
|
max_dim = 512
|
|
gq_w = min(w, max_dim)
|
|
gq_h = min(h, max_dim)
|
|
qx = np.linspace(0, w - 1, gq_w, dtype=np.float32)
|
|
qy = np.linspace(0, h - 1, gq_h, dtype=np.float32)
|
|
gx_q, gy_q = np.meshgrid(qx, qy)
|
|
|
|
rx = _griddata(all_orig, resid[:, 0], (gx_q, gy_q),
|
|
method='linear', fill_value=0.0)
|
|
ry = _griddata(all_orig, resid[:, 1], (gx_q, gy_q),
|
|
method='linear', fill_value=0.0)
|
|
rx = np.nan_to_num(rx, nan=0.0).astype(np.float32)
|
|
ry = np.nan_to_num(ry, nan=0.0).astype(np.float32)
|
|
|
|
if (gq_w, gq_h) != (w, h):
|
|
rx = cv2.resize(rx, (w, h), interpolation=cv2.INTER_CUBIC)
|
|
ry = cv2.resize(ry, (w, h), interpolation=cv2.INTER_CUBIC)
|
|
|
|
# --- full-res homography displacement + residual ---
|
|
grid_x, grid_y = np.meshgrid(
|
|
np.arange(w, dtype=np.float32), np.arange(h, dtype=np.float32))
|
|
pts_grid = np.stack([grid_x, grid_y], axis=-1).reshape(-1, 1, 2)
|
|
base_full = cv2.perspectiveTransform(pts_grid, M_inv).reshape(h, w, 2)
|
|
|
|
map_x = (base_full[:, :, 0] + rx).astype(np.float32)
|
|
map_y = (base_full[:, :, 1] + ry).astype(np.float32)
|
|
|
|
return cv2.remap(
|
|
mod_bgr, map_x, map_y, cv2.INTER_LINEAR,
|
|
borderMode=cv2.BORDER_REPLICATE)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Per-image driver
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _fix_one(orig_bgr, mod_bgr, method, max_mesh_points, downscale, refine_ecc):
|
|
h, w = orig_bgr.shape[:2]
|
|
mod_bgr = cv2.resize(mod_bgr, (w, h), interpolation=cv2.INTER_AREA) \
|
|
if mod_bgr.shape[:2] != (h, w) else mod_bgr
|
|
|
|
# Hardcoded SIFT tuning (matches original) + flat_4_points speed budget.
|
|
inl_mod, inl_orig, M_full = _sift_match(
|
|
orig_bgr, mod_bgr,
|
|
max_features=4096,
|
|
downscale=downscale,
|
|
contrast=0.01, edge=20, sigma=1.6, octave=3,
|
|
lowe=0.80, ransac_thresh=2.0,
|
|
)
|
|
|
|
if M_full is None:
|
|
return mod_bgr # no match, return unchanged
|
|
|
|
if refine_ecc:
|
|
M_full = _refine_with_ecc(orig_bgr, mod_bgr, M_full, downscale)
|
|
|
|
if method == "mesh":
|
|
warped = _warp_mesh(mod_bgr, M_full, inl_mod, inl_orig, max_mesh_points)
|
|
else:
|
|
warped = _warp_flat(mod_bgr, M_full)
|
|
|
|
return warped
|
|
|
|
|
|
def fix_pixel_drift(source_tensor, edited_tensor, method="flat_4_points",
|
|
max_mesh_points=400, downscale=4, refine_ecc=True):
|
|
device = edited_tensor.device
|
|
dtype = edited_tensor.dtype
|
|
|
|
src_list = _to_bgr_batch(source_tensor)
|
|
ed_list = _to_bgr_batch(edited_tensor)
|
|
|
|
out = []
|
|
for i, ed_bgr in enumerate(ed_list):
|
|
s_bgr = src_list[i] if i < len(src_list) else src_list[0]
|
|
warped = _fix_one(s_bgr, ed_bgr, method, max_mesh_points,
|
|
downscale, refine_ecc)
|
|
out.append(cv2.cvtColor(warped, cv2.COLOR_BGR2RGB))
|
|
|
|
arr = np.stack(out, 0).astype(np.float32) / 255.0
|
|
return torch.from_numpy(arr).to(device=device, dtype=dtype)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# ComfyUI node
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class PixelDriftFixNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"source_image": ("IMAGE",),
|
|
"edited_image": ("IMAGE",),
|
|
"method": (["flat_4_points", "mesh"], {
|
|
"default": "flat_4_points",
|
|
"tooltip": "flat_4_points gives better results and is faster. Mesh is experimental, don't use it."
|
|
}),
|
|
"max_mesh_points": ("INT", {
|
|
"default": 400,
|
|
"min": 100,
|
|
"max": 10000,
|
|
"step": 100,
|
|
"display": "number",
|
|
"tooltip": "Only used when method is set to 'mesh'. Higher values increase alignment accuracy but take longer. 400 = good/fast (~4s), 10000 = best quality (~90s)."
|
|
}),
|
|
"downscale": ("INT", {
|
|
"default": 4, "min": 1, "max": 8, "step": 1,
|
|
"display": "number",
|
|
"tooltip": "NEW: SIFT runs on 1/N resolution then homography is rescaled. 4 = best quality/speed balance. 8 = fastest (~1-2s), 1 = slowest but most RANSAC-robust."
|
|
}),
|
|
"refine_ecc": ("BOOLEAN", {
|
|
"default": True,
|
|
"tooltip": "NEW: sub-pixel refinement of the homography via cv2.findTransformECC. ~50ms, recommended ON."
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "apply"
|
|
CATEGORY = "image/alignment"
|
|
|
|
def apply(self, source_image, edited_image, method, max_mesh_points,
|
|
downscale, refine_ecc):
|
|
return (
|
|
fix_pixel_drift(
|
|
source_image, edited_image,
|
|
method=method,
|
|
max_mesh_points=max_mesh_points,
|
|
downscale=downscale,
|
|
refine_ecc=refine_ecc,
|
|
),
|
|
)
|