176 lines
8.3 KiB
Python
176 lines
8.3 KiB
Python
import cv2
|
|
import numpy as np
|
|
import torch
|
|
from skimage.transform import PiecewiseAffineTransform, warp
|
|
|
|
class PixelDriftFixNode:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@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."}),
|
|
"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)."
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("fixed_image",)
|
|
FUNCTION = "fix_pixel_drift"
|
|
CATEGORY = "image/processing"
|
|
|
|
def fix_pixel_drift(self, source_image, edited_image, method, max_mesh_points):
|
|
# ComfyUI image tensors are shapes: [B, H, W, C] and range 0.0 to 1.0
|
|
b1, h1, w1, c1 = source_image.shape
|
|
b2, h2, w2, c2 = edited_image.shape
|
|
|
|
batch_size = min(b1, b2)
|
|
output_tensors = []
|
|
|
|
for i in range(batch_size):
|
|
# 1. Convert PyTorch tensor to uint8 BGR for OpenCV processing
|
|
img_orig_rgb = (source_image[i].cpu().numpy() * 255).astype(np.uint8)
|
|
img_mod_rgb = (edited_image[i].cpu().numpy() * 255).astype(np.uint8)
|
|
|
|
img_orig = cv2.cvtColor(img_orig_rgb, cv2.COLOR_RGB2BGR)
|
|
img_mod = cv2.cvtColor(img_mod_rgb, cv2.COLOR_RGB2BGR)
|
|
|
|
height, width = img_orig.shape[:2]
|
|
|
|
# Grayscale conversion
|
|
gray_orig = cv2.cvtColor(img_orig, cv2.COLOR_BGR2GRAY)
|
|
gray_mod = cv2.cvtColor(img_mod, cv2.COLOR_BGR2GRAY)
|
|
|
|
# Ultra-Sensitive SIFT
|
|
sift = cv2.SIFT_create(
|
|
nfeatures=0,
|
|
nOctaveLayers=3,
|
|
contrastThreshold=0.01,
|
|
edgeThreshold=20,
|
|
sigma=1.6
|
|
)
|
|
|
|
kp_orig, des_orig = sift.detectAndCompute(gray_orig, None)
|
|
kp_mod, des_mod = sift.detectAndCompute(gray_mod, None)
|
|
|
|
# Fallback check if features are completely missing
|
|
if des_orig is None or des_mod is None or len(kp_orig) < 10 or len(kp_mod) < 10:
|
|
print(f"[PixelDriftFix] Warning: Not enough features found in batch {i}. Passing edited image through.")
|
|
output_tensors.append(edited_image[i])
|
|
continue
|
|
|
|
# Match features
|
|
bf = cv2.BFMatcher()
|
|
matches = bf.knnMatch(des_mod, des_orig, k=2)
|
|
|
|
# Lowe's ratio test
|
|
good_matches = []
|
|
for match_pair in matches:
|
|
if len(match_pair) == 2:
|
|
m, n = match_pair
|
|
if m.distance < 0.80 * n.distance:
|
|
good_matches.append(m)
|
|
|
|
if len(good_matches) < 10:
|
|
print(f"[PixelDriftFix] Warning: Too few good matches ({len(good_matches)}) in batch {i}. Passing edited image through.")
|
|
output_tensors.append(edited_image[i])
|
|
continue
|
|
|
|
# Extract coordinates
|
|
pts_mod = np.float32([kp_mod[m.queryIdx].pt for m in good_matches]).reshape(-1, 2)
|
|
pts_orig = np.float32([kp_orig[m.trainIdx].pt for m in good_matches]).reshape(-1, 2)
|
|
|
|
# RANSAC filter to get Homography matrix
|
|
M, mask = cv2.findHomography(pts_mod, pts_orig, cv2.RANSAC, 5.0)
|
|
|
|
if M is None or mask is None:
|
|
print(f"[PixelDriftFix] Warning: Homography matrix computation failed in batch {i}.")
|
|
output_tensors.append(edited_image[i])
|
|
continue
|
|
|
|
inliers_mod = pts_mod[mask.ravel() == 1]
|
|
inliers_orig = pts_orig[mask.ravel() == 1]
|
|
|
|
if len(inliers_mod) < 10:
|
|
print(f"[PixelDriftFix] Warning: Too few static features ({len(inliers_mod)}) found after RANSAC filtering. Images should be similar.")
|
|
output_tensors.append(edited_image[i])
|
|
continue
|
|
|
|
# Calculate Global Layer (always needed either as final output or as mesh fallback boundary)
|
|
global_warped = cv2.warpPerspective(
|
|
img_mod, M, (width, height), borderMode=cv2.BORDER_REPLICATE
|
|
)
|
|
|
|
# Route calculation based on selected UI dropdown method
|
|
if method == "mesh":
|
|
try:
|
|
# Construct rigid outer boundary anchors
|
|
edge_points = []
|
|
step_size = 30
|
|
for x in range(0, width, step_size):
|
|
edge_points.append([x, 0])
|
|
edge_points.append([x, height - 1])
|
|
for y in range(0, height, step_size):
|
|
edge_points.append([0, y])
|
|
edge_points.append([width - 1, y])
|
|
|
|
edge_points.extend([[0, 0], [width - 1, 0], [0, height - 1], [width - 1, height - 1]])
|
|
corners = np.array(edge_points, dtype=np.float32).reshape(-1, 2)
|
|
corners = np.unique(corners, axis=0)
|
|
|
|
corners_proj = cv2.perspectiveTransform(corners.reshape(-1, 1, 2), M).reshape(-1, 2)
|
|
corners_proj[:, 0] = np.clip(corners_proj[:, 0], 0, width - 1)
|
|
corners_proj[:, 1] = np.clip(corners_proj[:, 1], 0, height - 1)
|
|
|
|
final_mod = np.vstack([inliers_mod, corners])
|
|
final_orig = np.vstack([inliers_orig, corners_proj])
|
|
|
|
# Cap points based on user's slider choice
|
|
if len(final_mod) > max_mesh_points:
|
|
idx = np.linspace(0, len(final_mod) - 1, max_mesh_points, dtype=int)
|
|
final_mod = final_mod[idx]
|
|
final_orig = final_orig[idx]
|
|
print(f"PixelDriftFix: using mesh with {len(final_mod)} points")
|
|
|
|
# Adaptive API initialization for scikit-image cross-version compatibility
|
|
if hasattr(PiecewiseAffineTransform, 'from_estimate'):
|
|
tform = PiecewiseAffineTransform.from_estimate(final_orig, final_mod)
|
|
else:
|
|
tform = PiecewiseAffineTransform()
|
|
tform.estimate(final_orig, final_mod)
|
|
|
|
local_warped_raw = warp(img_mod, tform, output_shape=(height, width), order=1, mode='edge')
|
|
local_warped = (local_warped_raw * 255).astype(np.uint8)
|
|
|
|
coverage_input = np.ones((height, width), dtype=np.float32)
|
|
coverage_warped = warp(coverage_input, tform, output_shape=(height, width), order=0, cval=0)
|
|
valid_mesh_mask = coverage_warped > 0.5
|
|
|
|
final_img = np.where(valid_mesh_mask[:, :, None], local_warped, global_warped)
|
|
except Exception as warp_error:
|
|
print(f"[PixelDriftFix] Error during dense warping: {warp_error}. Falling back to linear matrix transform.")
|
|
final_img = global_warped
|
|
else:
|
|
# Default "flat_4_points" path: clean global warp matrix execution
|
|
print("PixelDriftFix: using flat_4_points")
|
|
final_img = global_warped
|
|
|
|
# 2. Convert final BGR image back to RGB and then normalize PyTorch Tensor [H, W, C]
|
|
final_rgb = cv2.cvtColor(final_img, cv2.COLOR_BGR2RGB)
|
|
out_tensor = torch.from_numpy(final_rgb).float() / 255.0
|
|
output_tensors.append(out_tensor)
|
|
|
|
# Stack separate batch images back into uniform [B, H, W, C] format
|
|
fixed_image_batch = torch.stack(output_tensors, dim=0)
|
|
return (fixed_image_batch,) |