412 lines
18 KiB
Python
412 lines
18 KiB
Python
# File: ComfyUI/custom_nodes/ComfyUI-PanoTools/CameraCalibration.py
|
|
"""
|
|
Camera Calibration Node for ComfyUI-PanoTools
|
|
|
|
Estimates camera parameters (roll, pitch, vFoV) from single images using
|
|
various calibration methods based on scene geometry.
|
|
|
|
Methods supported:
|
|
- GeoCalib: Deep learning-based calibration (requires optional dependency)
|
|
- Horizon Detection: For outdoor scenes with visible horizon
|
|
- Vertical Lines: For architectural scenes with vertical structures
|
|
- Manual: User-specified parameters
|
|
"""
|
|
|
|
import numpy as np
|
|
import torch
|
|
from typing import Tuple, Optional
|
|
|
|
class CameraCalibration:
|
|
"""
|
|
Estimates camera calibration parameters for perspective to panorama conversion.
|
|
|
|
Returns focal_length_mm as an additional output for compatibility with other
|
|
custom nodes that may need 35mm equivalent focal length values.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self._geocalib_available = None
|
|
self._geocalib_model = None
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE", {"tooltip": "Single image or first frame to analyze for camera calibration parameters"}),
|
|
"calibration_method": ([
|
|
"manual",
|
|
"horizon_detection",
|
|
"vertical_lines",
|
|
"geocalib_auto"
|
|
], {"tooltip": "Calibration method: 'geocalib_auto' (AI-based, works for all scenes - recommended), 'horizon_detection' (for outdoor/ocean scenes), 'vertical_lines' (for architecture/buildings), or 'manual' (specify values manually)"}),
|
|
},
|
|
"optional": {
|
|
# Manual override parameters
|
|
"manual_roll": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.01, "tooltip": "Camera roll (rotation around viewing axis) in degrees. Positive = clockwise tilt. Only used when method is 'manual'."}),
|
|
"manual_pitch": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.01, "tooltip": "Camera pitch (up/down angle) in degrees. Positive = tilted up, negative = tilted down. Only used when method is 'manual'."}),
|
|
"manual_vfov": ("FLOAT", {"default": 60.0, "min": 10.0, "max": 170.0, "step": 0.1, "tooltip": "Vertical field of view in degrees. Typical values: 40-60° (standard lens), 70-90° (wide angle), 10-30° (telephoto). Used when method is 'manual' or as fallback for geometric methods."}),
|
|
|
|
# Horizon detection parameters
|
|
"horizon_roi_top": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Top boundary of search region for horizon (0.0 = image top, 1.0 = image bottom). Adjust to focus on area where horizon is expected. Only for 'horizon_detection' method."}),
|
|
"horizon_roi_bottom": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Bottom boundary of search region for horizon (0.0 = image top, 1.0 = image bottom). Should be below the horizon line. Only for 'horizon_detection' method."}),
|
|
|
|
# Vertical line detection parameters
|
|
"min_line_length": ("INT", {"default": 100, "min": 20, "max": 500, "step": 10, "tooltip": "Minimum length (in pixels) for lines to be detected. Increase for images with noise or small details. Only for 'vertical_lines' method."}),
|
|
"angle_threshold": ("FLOAT", {"default": 5.0, "min": 0.1, "max": 15.0, "step": 0.1, "tooltip": "Maximum angle deviation from perfectly vertical (in degrees). Lower = stricter vertical alignment. Increase if camera has significant roll. Only for 'vertical_lines' method."}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("FLOAT", "FLOAT", "FLOAT", "FLOAT", "FLOAT", "STRING")
|
|
RETURN_NAMES = ("roll", "pitch", "vfov", "focal_length", "focal_length_mm", "calibration_info")
|
|
OUTPUT_TOOLTIPS = ("Camera roll rotation in degrees",
|
|
"Camera pitch rotation in degrees",
|
|
"Vertical field of view in degrees",
|
|
"Focal length in pixels (used for projection calculations)",
|
|
"35mm equivalent focal length (e.g., 24mm wide, 50mm standard, 85mm portrait). Useful for comparing with real camera lenses.",
|
|
"Diagnostic information about the calibration method used and detected parameters")
|
|
FUNCTION = "calibrate_camera"
|
|
CATEGORY = "PanoTools"
|
|
|
|
def calibrate_camera(
|
|
self,
|
|
image,
|
|
calibration_method,
|
|
manual_roll=0.0,
|
|
manual_pitch=0.0,
|
|
manual_vfov=60.0,
|
|
horizon_roi_top=0.3,
|
|
horizon_roi_bottom=0.7,
|
|
min_line_length=100,
|
|
angle_threshold=5.0
|
|
):
|
|
"""
|
|
Calibrate camera parameters based on selected method.
|
|
"""
|
|
|
|
# Convert to numpy for processing
|
|
img_np = (image[0].cpu().numpy() * 255).astype(np.uint8)
|
|
height, width = img_np.shape[:2]
|
|
|
|
if calibration_method == "manual":
|
|
roll, pitch, vfov = manual_roll, manual_pitch, manual_vfov
|
|
focal_length = height / (2.0 * np.tan(np.radians(vfov) / 2.0))
|
|
focal_length_mm = 12.0 / np.tan(np.radians(vfov) / 2.0)
|
|
info = f"Manual calibration: roll={roll:.2f}°, pitch={pitch:.2f}°, vFoV={vfov:.2f}°"
|
|
|
|
elif calibration_method == "horizon_detection":
|
|
roll, pitch, vfov, focal_length, focal_length_mm, info = self._calibrate_from_horizon(
|
|
img_np, horizon_roi_top, horizon_roi_bottom, manual_vfov
|
|
)
|
|
|
|
elif calibration_method == "vertical_lines":
|
|
roll, pitch, vfov, focal_length, focal_length_mm, info = self._calibrate_from_vertical_lines(
|
|
img_np, min_line_length, angle_threshold, manual_vfov
|
|
)
|
|
|
|
elif calibration_method == "geocalib_auto":
|
|
roll, pitch, vfov, focal_length, focal_length_mm, info = self._calibrate_with_geocalib(img_np)
|
|
|
|
else:
|
|
roll, pitch, vfov = 0.0, 0.0, 60.0
|
|
focal_length = height / (2.0 * np.tan(np.radians(vfov) / 2.0))
|
|
focal_length_mm = 12.0 / np.tan(np.radians(vfov) / 2.0)
|
|
info = f"Unknown method '{calibration_method}', using defaults"
|
|
|
|
return (roll, pitch, vfov, focal_length, focal_length_mm, info)
|
|
|
|
def _calibrate_from_horizon(
|
|
self,
|
|
img: np.ndarray,
|
|
roi_top: float,
|
|
roi_bottom: float,
|
|
default_vfov: float
|
|
) -> Tuple[float, float, float, float, float, str]:
|
|
"""
|
|
Calibrate camera from horizon line detection.
|
|
|
|
Method:
|
|
1. Detect horizon line using edge detection in ROI
|
|
2. Calculate roll from horizon slope
|
|
3. Calculate pitch from horizon vertical position
|
|
4. Use default or estimate vFoV
|
|
"""
|
|
try:
|
|
import cv2
|
|
except ImportError:
|
|
focal_length = img.shape[0] / (2.0 * np.tan(np.radians(default_vfov) / 2.0))
|
|
focal_length_mm = 12.0 / np.tan(np.radians(default_vfov) / 2.0)
|
|
return 0.0, 0.0, default_vfov, focal_length, focal_length_mm, "ERROR: OpenCV not available. Install with: pip install opencv-python"
|
|
|
|
height, width = img.shape[:2]
|
|
|
|
# Define ROI for horizon detection
|
|
y1 = int(height * roi_top)
|
|
y2 = int(height * roi_bottom)
|
|
roi = img[y1:y2, :]
|
|
|
|
# Convert to grayscale
|
|
gray = cv2.cvtColor(roi, cv2.COLOR_RGB2GRAY)
|
|
|
|
# Edge detection
|
|
edges = cv2.Canny(gray, 50, 150, apertureSize=3)
|
|
|
|
# Hough line detection
|
|
lines = cv2.HoughLinesP(
|
|
edges,
|
|
rho=1,
|
|
theta=np.pi/180,
|
|
threshold=100,
|
|
minLineLength=width//4,
|
|
maxLineGap=50
|
|
)
|
|
|
|
if lines is None or len(lines) == 0:
|
|
focal_length = height / (2.0 * np.tan(np.radians(default_vfov) / 2.0))
|
|
focal_length_mm = 12.0 / np.tan(np.radians(default_vfov) / 2.0)
|
|
return 0.0, 0.0, default_vfov, focal_length, focal_length_mm, "WARNING: No horizon detected, using defaults"
|
|
|
|
# Find most horizontal line (smallest absolute angle)
|
|
best_angle = 90.0
|
|
best_y_pos = height / 2
|
|
|
|
for line in lines:
|
|
x1, y1_line, x2, y2_line = line[0]
|
|
|
|
# Calculate angle from horizontal
|
|
if x2 != x1:
|
|
angle = np.degrees(np.arctan2(y2_line - y1_line, x2 - x1))
|
|
|
|
# Only consider nearly horizontal lines
|
|
if abs(angle) < 15:
|
|
if abs(angle) < abs(best_angle):
|
|
best_angle = angle
|
|
best_y_pos = (y1_line + y2_line) / 2 + y1 # Adjust for ROI offset
|
|
|
|
# Calculate roll (negative because image y increases downward)
|
|
roll = -best_angle
|
|
|
|
# Calculate pitch from horizon position
|
|
# If horizon is above center, camera is pitched up (positive pitch)
|
|
horizon_offset = (height / 2 - best_y_pos) / height
|
|
pitch = horizon_offset * default_vfov # Approximate
|
|
|
|
focal_length = height / (2.0 * np.tan(np.radians(default_vfov) / 2.0))
|
|
focal_length_mm = 12.0 / np.tan(np.radians(default_vfov) / 2.0)
|
|
info = f"Horizon detected: roll={roll:.2f}°, pitch={pitch:.2f}°, vFoV={default_vfov:.2f}° (manual)"
|
|
|
|
return roll, pitch, default_vfov, focal_length, focal_length_mm, info
|
|
|
|
def _calibrate_from_vertical_lines(
|
|
self,
|
|
img: np.ndarray,
|
|
min_line_length: int,
|
|
angle_threshold: float,
|
|
default_vfov: float
|
|
) -> Tuple[float, float, float, float, float, str]:
|
|
"""
|
|
Calibrate camera from vertical line detection (building edges, poles, etc).
|
|
|
|
Method:
|
|
1. Detect long straight lines
|
|
2. Find lines close to vertical
|
|
3. Calculate roll from average vertical line angle
|
|
4. Estimate pitch from vanishing point if possible
|
|
"""
|
|
try:
|
|
import cv2
|
|
except ImportError:
|
|
height = img.shape[0]
|
|
focal_length = height / (2.0 * np.tan(np.radians(default_vfov) / 2.0))
|
|
focal_length_mm = 12.0 / np.tan(np.radians(default_vfov) / 2.0)
|
|
return 0.0, 0.0, default_vfov, focal_length, focal_length_mm, "ERROR: OpenCV not available. Install with: pip install opencv-python"
|
|
|
|
height, width = img.shape[:2]
|
|
|
|
# Convert to grayscale
|
|
gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
|
|
|
|
# Edge detection
|
|
edges = cv2.Canny(gray, 50, 150, apertureSize=3)
|
|
|
|
# Hough line detection
|
|
lines = cv2.HoughLinesP(
|
|
edges,
|
|
rho=1,
|
|
theta=np.pi/180,
|
|
threshold=80,
|
|
minLineLength=min_line_length,
|
|
maxLineGap=20
|
|
)
|
|
|
|
if lines is None or len(lines) == 0:
|
|
focal_length = height / (2.0 * np.tan(np.radians(default_vfov) / 2.0))
|
|
focal_length_mm = 12.0 / np.tan(np.radians(default_vfov) / 2.0)
|
|
return 0.0, 0.0, default_vfov, focal_length, focal_length_mm, "WARNING: No vertical lines detected, using defaults"
|
|
|
|
# Find vertical lines (angles close to ±90°)
|
|
vertical_angles = []
|
|
|
|
for line in lines:
|
|
x1, y1, x2, y2 = line[0]
|
|
|
|
# Calculate angle from horizontal
|
|
if x2 != x1:
|
|
angle = np.degrees(np.arctan2(y2 - y1, x2 - x1))
|
|
else:
|
|
angle = 90.0
|
|
|
|
# Normalize to [-90, 90]
|
|
while angle > 90:
|
|
angle -= 180
|
|
while angle < -90:
|
|
angle += 180
|
|
|
|
# Check if close to vertical (90° or -90°)
|
|
deviation_from_vertical = min(abs(90 - abs(angle)), abs(angle))
|
|
|
|
if deviation_from_vertical < angle_threshold:
|
|
vertical_angles.append(angle)
|
|
|
|
if len(vertical_angles) == 0:
|
|
focal_length = height / (2.0 * np.tan(np.radians(default_vfov) / 2.0))
|
|
focal_length_mm = 12.0 / np.tan(np.radians(default_vfov) / 2.0)
|
|
return 0.0, 0.0, default_vfov, focal_length, focal_length_mm, "WARNING: No sufficiently vertical lines found, using defaults"
|
|
|
|
# Calculate roll as deviation from perfect vertical
|
|
# If lines lean right, roll is positive
|
|
avg_angle = np.median(vertical_angles)
|
|
|
|
# Roll correction: if vertical lines are at 85°, camera is rolled -5°
|
|
if avg_angle > 0:
|
|
roll = avg_angle - 90.0
|
|
else:
|
|
roll = avg_angle + 90.0
|
|
|
|
# Pitch estimation is harder from vertical lines alone
|
|
# Use a simple heuristic or default
|
|
pitch = 0.0
|
|
|
|
focal_length = height / (2.0 * np.tan(np.radians(default_vfov) / 2.0))
|
|
focal_length_mm = 12.0 / np.tan(np.radians(default_vfov) / 2.0)
|
|
info = f"Vertical lines: {len(vertical_angles)} detected, roll={roll:.2f}°, pitch={pitch:.2f}° (estimated), vFoV={default_vfov:.2f}° (manual)"
|
|
|
|
return roll, pitch, default_vfov, focal_length, focal_length_mm, info
|
|
|
|
def _calibrate_with_geocalib(
|
|
self,
|
|
img: np.ndarray
|
|
) -> Tuple[float, float, float, float, float, str]:
|
|
"""
|
|
Calibrate camera using GeoCalib deep learning model.
|
|
|
|
GeoCalib is vendored (included) in this package.
|
|
"""
|
|
|
|
# Check if GeoCalib is available (cache result)
|
|
if self._geocalib_available is None:
|
|
try:
|
|
from .geocalib import GeoCalib
|
|
self._geocalib_available = True
|
|
print("[CameraCalibration] GeoCalib loaded (vendored)")
|
|
except ImportError as e:
|
|
self._geocalib_available = False
|
|
print(f"[CameraCalibration] GeoCalib import failed: {e}")
|
|
|
|
if not self._geocalib_available:
|
|
default_vfov = 60.0
|
|
focal_length_mm = 12.0 / np.tan(np.radians(default_vfov) / 2.0)
|
|
return 0.0, 0.0, default_vfov, 0.0, focal_length_mm, (
|
|
"ERROR: GeoCalib failed to load. Check that all dependencies are installed: "
|
|
"torch, torchvision, opencv-python, kornia, matplotlib"
|
|
)
|
|
|
|
try:
|
|
from .geocalib import GeoCalib
|
|
|
|
# Initialize model (lazy loading)
|
|
if self._geocalib_model is None:
|
|
print("[CameraCalibration] Loading GeoCalib model...")
|
|
self._geocalib_model = GeoCalib()
|
|
self._geocalib_model.eval()
|
|
print("[CameraCalibration] GeoCalib model loaded")
|
|
|
|
# Get image dimensions
|
|
height, width = img.shape[:2]
|
|
|
|
# Convert numpy (H,W,C) uint8 RGB -> torch (C,H,W) float [0,1]
|
|
img_tensor = torch.from_numpy(img).permute(2, 0, 1).float() / 255.0
|
|
|
|
# Run calibration using the correct method
|
|
with torch.no_grad():
|
|
result = self._geocalib_model.calibrate(img_tensor)
|
|
|
|
# Extract camera and gravity objects
|
|
camera = result['camera']
|
|
gravity = result['gravity']
|
|
|
|
# Get focal length from camera object
|
|
# camera.f is the focal length tensor with shape (..., 2) for [fx, fy]
|
|
f_tensor = camera.f
|
|
|
|
# Extract fy (focal length in y direction) which corresponds to vertical FoV
|
|
if f_tensor.dim() == 2:
|
|
# Batch dimension: shape (batch, 2)
|
|
focal_length = f_tensor[0, 1].item()
|
|
else:
|
|
# No batch: shape (2,)
|
|
focal_length = f_tensor[1].item()
|
|
|
|
# Calculate vFoV from focal length
|
|
if focal_length > 0:
|
|
vfov = 2 * np.degrees(np.arctan(height / (2 * focal_length)))
|
|
else:
|
|
vfov = 60.0
|
|
|
|
# Use the Gravity object's built-in roll and pitch properties
|
|
# These handle coordinate system transformations correctly for any image orientation
|
|
if hasattr(gravity, 'roll') and hasattr(gravity, 'pitch'):
|
|
# Gravity object provides roll/pitch that account for orientation
|
|
roll_rad = gravity.roll
|
|
pitch_rad = gravity.pitch
|
|
|
|
# Convert from radians to degrees, handling tensor extraction
|
|
if hasattr(roll_rad, 'item'):
|
|
roll = np.degrees(roll_rad.item() if roll_rad.numel() == 1 else roll_rad[0].item())
|
|
pitch = np.degrees(pitch_rad.item() if pitch_rad.numel() == 1 else pitch_rad[0].item())
|
|
else:
|
|
roll = np.degrees(float(roll_rad))
|
|
pitch = np.degrees(float(pitch_rad))
|
|
else:
|
|
# Fallback: extract gravity vector manually (less robust for non-landscape)
|
|
if hasattr(gravity, 'data'):
|
|
gravity_data = gravity.data.cpu().numpy()
|
|
else:
|
|
gravity_data = gravity.cpu().numpy()
|
|
|
|
# Handle batch dimension if present
|
|
if len(gravity_data.shape) > 1:
|
|
gx, gy, gz = gravity_data[0, :3]
|
|
else:
|
|
gx, gy, gz = gravity_data[:3]
|
|
|
|
# Calculate roll and pitch from gravity direction
|
|
roll = np.degrees(np.arctan2(gx, gy))
|
|
pitch = np.degrees(np.arctan2(-gz, np.sqrt(gx**2 + gy**2)))
|
|
|
|
focal_length_mm = 12.0 / np.tan(np.radians(vfov) / 2.0)
|
|
info = f"GeoCalib: roll={roll:.2f}°, pitch={pitch:.2f}°, vFoV={vfov:.2f}° (fl={focal_length:.1f}px)"
|
|
|
|
return roll, pitch, vfov, focal_length, focal_length_mm, info
|
|
|
|
except Exception as e:
|
|
import traceback
|
|
traceback.print_exc()
|
|
default_vfov = 60.0
|
|
focal_length_mm = 12.0 / np.tan(np.radians(default_vfov) / 2.0)
|
|
return 0.0, 0.0, default_vfov, 0.0, focal_length_mm, f"ERROR: GeoCalib failed: {str(e)}"
|
|
|
|
|
|
# --- Node Registration ---
|
|
NODE_CLASS_MAPPINGS = {"CameraCalibration": CameraCalibration}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"CameraCalibration": "Camera Calibration"}
|