305 lines
9.0 KiB
Python
305 lines
9.0 KiB
Python
"""
|
|
Fisheye Transform Nodes for ComfyUI
|
|
- FisheyeUndistort: Convert fisheye image to rectilinear (flat) image
|
|
- FisheyeRedistort: Convert rectilinear image back to fisheye
|
|
"""
|
|
|
|
import numpy as np
|
|
import torch
|
|
import cv2
|
|
|
|
|
|
class FisheyeUndistort:
|
|
"""
|
|
Undistort a fisheye image to rectilinear (flat perspective).
|
|
Use this BEFORE inpainting for correct perspective generation.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"strength": (
|
|
"FLOAT",
|
|
{
|
|
"default": 1.0,
|
|
"min": 0.0,
|
|
"max": 3.0,
|
|
"step": 0.05,
|
|
"display": "slider",
|
|
},
|
|
),
|
|
"zoom": (
|
|
"FLOAT",
|
|
{
|
|
"default": 1.0,
|
|
"min": 0.5,
|
|
"max": 2.0,
|
|
"step": 0.05,
|
|
"display": "slider",
|
|
},
|
|
),
|
|
},
|
|
"optional": {
|
|
"mask": ("MASK",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK", "FISHEYE_PARAMS")
|
|
RETURN_NAMES = ("image", "mask", "params")
|
|
FUNCTION = "undistort"
|
|
CATEGORY = "image/transform"
|
|
|
|
def undistort(self, image, strength, zoom, mask=None):
|
|
# Convert from torch tensor [B, H, W, C] to numpy
|
|
batch_size = image.shape[0]
|
|
results = []
|
|
mask_results = []
|
|
|
|
# Store params for re-distortion
|
|
params = {
|
|
"strength": strength,
|
|
"zoom": zoom,
|
|
"original_shape": (image.shape[1], image.shape[2]),
|
|
}
|
|
|
|
for i in range(batch_size):
|
|
img_np = (image[i].cpu().numpy() * 255).astype(np.uint8)
|
|
h, w = img_np.shape[:2]
|
|
|
|
# Create undistortion maps
|
|
map_x, map_y = self._create_undistort_map(w, h, strength, zoom)
|
|
|
|
# Apply undistortion
|
|
undistorted = cv2.remap(
|
|
img_np, map_x, map_y, cv2.INTER_LANCZOS4, borderMode=cv2.BORDER_REFLECT
|
|
)
|
|
results.append(undistorted)
|
|
|
|
# Process mask if provided
|
|
if mask is not None:
|
|
if len(mask.shape) == 3:
|
|
mask_np = (mask[i].cpu().numpy() * 255).astype(np.uint8)
|
|
else:
|
|
mask_np = (mask.cpu().numpy() * 255).astype(np.uint8)
|
|
mask_undist = cv2.remap(
|
|
mask_np,
|
|
map_x,
|
|
map_y,
|
|
cv2.INTER_NEAREST,
|
|
borderMode=cv2.BORDER_CONSTANT,
|
|
borderValue=0,
|
|
)
|
|
mask_results.append(mask_undist)
|
|
|
|
# Convert back to torch tensor
|
|
result_tensor = torch.from_numpy(np.stack(results).astype(np.float32) / 255.0)
|
|
|
|
if mask is not None:
|
|
mask_tensor = torch.from_numpy(
|
|
np.stack(mask_results).astype(np.float32) / 255.0
|
|
)
|
|
else:
|
|
# Return empty mask matching image dimensions
|
|
mask_tensor = torch.zeros(
|
|
(batch_size, result_tensor.shape[1], result_tensor.shape[2])
|
|
)
|
|
|
|
return (result_tensor, mask_tensor, params)
|
|
|
|
def _create_undistort_map(self, w, h, strength, zoom):
|
|
"""Create remapping arrays for fisheye undistortion."""
|
|
# Center of image
|
|
cx, cy = w / 2, h / 2
|
|
|
|
# Create coordinate grids
|
|
x = np.arange(w, dtype=np.float32)
|
|
y = np.arange(h, dtype=np.float32)
|
|
map_x, map_y = np.meshgrid(x, y)
|
|
|
|
# Normalize to [-1, 1]
|
|
nx = (map_x - cx) / cx
|
|
ny = (map_y - cy) / cy
|
|
|
|
# Distance from center
|
|
r = np.sqrt(nx**2 + ny**2)
|
|
|
|
# Avoid division by zero
|
|
r = np.where(r == 0, 1e-10, r)
|
|
|
|
# Fisheye undistortion: inverse barrel distortion
|
|
# For undistort, we apply the inverse: r_undist = r * (1 + k * r^2)
|
|
k = -strength * 0.3 # Negative for undistort (removes barrel)
|
|
r_undist = r * (1 + k * r**2)
|
|
|
|
# Apply zoom
|
|
r_undist = r_undist / zoom
|
|
|
|
# Scale factor
|
|
scale = np.where(r > 0, r_undist / r, 1)
|
|
|
|
# New coordinates
|
|
map_x = (nx * scale) * cx + cx
|
|
map_y = (ny * scale) * cy + cy
|
|
|
|
return map_x.astype(np.float32), map_y.astype(np.float32)
|
|
|
|
|
|
class FisheyeRedistort:
|
|
"""
|
|
Re-apply fisheye distortion to a rectilinear image.
|
|
Use this AFTER inpainting to restore fisheye perspective.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"params": ("FISHEYE_PARAMS",),
|
|
},
|
|
"optional": {
|
|
"mask": ("MASK",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK")
|
|
RETURN_NAMES = ("image", "mask")
|
|
FUNCTION = "redistort"
|
|
CATEGORY = "image/transform"
|
|
|
|
def redistort(self, image, params, mask=None):
|
|
strength = params["strength"]
|
|
zoom = params["zoom"]
|
|
|
|
batch_size = image.shape[0]
|
|
results = []
|
|
mask_results = []
|
|
|
|
for i in range(batch_size):
|
|
img_np = (image[i].cpu().numpy() * 255).astype(np.uint8)
|
|
h, w = img_np.shape[:2]
|
|
|
|
# Create redistortion maps (inverse of undistort)
|
|
map_x, map_y = self._create_redistort_map(w, h, strength, zoom)
|
|
|
|
# Apply redistortion
|
|
redistorted = cv2.remap(
|
|
img_np, map_x, map_y, cv2.INTER_LANCZOS4, borderMode=cv2.BORDER_REFLECT
|
|
)
|
|
results.append(redistorted)
|
|
|
|
# Process mask if provided
|
|
if mask is not None:
|
|
if len(mask.shape) == 3:
|
|
mask_np = (mask[i].cpu().numpy() * 255).astype(np.uint8)
|
|
else:
|
|
mask_np = (mask.cpu().numpy() * 255).astype(np.uint8)
|
|
mask_redist = cv2.remap(
|
|
mask_np,
|
|
map_x,
|
|
map_y,
|
|
cv2.INTER_NEAREST,
|
|
borderMode=cv2.BORDER_CONSTANT,
|
|
borderValue=0,
|
|
)
|
|
mask_results.append(mask_redist)
|
|
|
|
# Convert back to torch tensor
|
|
result_tensor = torch.from_numpy(np.stack(results).astype(np.float32) / 255.0)
|
|
|
|
if mask is not None:
|
|
mask_tensor = torch.from_numpy(
|
|
np.stack(mask_results).astype(np.float32) / 255.0
|
|
)
|
|
else:
|
|
mask_tensor = torch.zeros(
|
|
(batch_size, result_tensor.shape[1], result_tensor.shape[2])
|
|
)
|
|
|
|
return (result_tensor, mask_tensor)
|
|
|
|
def _create_redistort_map(self, w, h, strength, zoom):
|
|
"""Create remapping arrays for fisheye redistortion (barrel distortion)."""
|
|
cx, cy = w / 2, h / 2
|
|
|
|
x = np.arange(w, dtype=np.float32)
|
|
y = np.arange(h, dtype=np.float32)
|
|
map_x, map_y = np.meshgrid(x, y)
|
|
|
|
nx = (map_x - cx) / cx
|
|
ny = (map_y - cy) / cy
|
|
|
|
r = np.sqrt(nx**2 + ny**2)
|
|
r = np.where(r == 0, 1e-10, r)
|
|
|
|
# Fisheye redistortion: apply barrel distortion
|
|
k = strength * 0.3 # Positive for redistort (adds barrel)
|
|
|
|
# Apply zoom first (inverse of undistort zoom)
|
|
r_zoomed = r * zoom
|
|
|
|
# Then apply distortion
|
|
r_dist = r_zoomed * (1 + k * r_zoomed**2)
|
|
|
|
scale = np.where(r > 0, r_dist / r, 1)
|
|
|
|
map_x = (nx * scale) * cx + cx
|
|
map_y = (ny * scale) * cy + cy
|
|
|
|
return map_x.astype(np.float32), map_y.astype(np.float32)
|
|
|
|
|
|
class FisheyeStrengthEstimate:
|
|
"""
|
|
Estimate fisheye distortion strength from an image.
|
|
Uses edge detection to estimate barrel distortion amount.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("FLOAT",)
|
|
RETURN_NAMES = ("estimated_strength",)
|
|
FUNCTION = "estimate"
|
|
CATEGORY = "image/transform"
|
|
|
|
def estimate(self, image):
|
|
# Simple heuristic: typical fisheye cameras have strength around 0.8-1.2
|
|
# This could be improved with actual calibration
|
|
img_np = (image[0].cpu().numpy() * 255).astype(np.uint8)
|
|
h, w = img_np.shape[:2]
|
|
|
|
# Check aspect ratio - wider images typically have more distortion
|
|
aspect = w / h
|
|
|
|
# Heuristic estimation
|
|
if aspect > 1.5:
|
|
strength = 1.0
|
|
elif aspect > 1.2:
|
|
strength = 0.9
|
|
else:
|
|
strength = 0.8
|
|
|
|
return (strength,)
|
|
|
|
|
|
# Node registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"FisheyeUndistort": FisheyeUndistort,
|
|
"FisheyeRedistort": FisheyeRedistort,
|
|
"FisheyeStrengthEstimate": FisheyeStrengthEstimate,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"FisheyeUndistort": "🐟 Fisheye Undistort",
|
|
"FisheyeRedistort": "🐟 Fisheye Redistort",
|
|
"FisheyeStrengthEstimate": "🐟 Fisheye Strength Estimate",
|
|
}
|