Cedar ConnorandClaude
ea1aa67936
- Implement 3D rotation (yaw/pitch/roll) for equirectangular images - Add horizon line adjustment functionality - Support 180-degree longitude cropping with custom output dimensions - Include bilinear interpolation for smooth transformations - Provide three node types: rotation-only, crop-only, and combined processor - Add comprehensive mathematical coordinate transformations - Support batch processing with progress indicators 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
190 lines
7.5 KiB
Python
190 lines
7.5 KiB
Python
import numpy as np
|
|
import torch
|
|
import cv2
|
|
from typing import Tuple, Optional
|
|
from scipy.spatial.transform import Rotation
|
|
|
|
|
|
class EquirectangularProcessor:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@staticmethod
|
|
def equirectangular_to_spherical(x: np.ndarray, y: np.ndarray, width: int, height: int) -> Tuple[np.ndarray, np.ndarray]:
|
|
"""Convert equirectangular coordinates to spherical coordinates (lat, lon)"""
|
|
lon = (x / width) * 2 * np.pi - np.pi # -π to π
|
|
lat = np.pi/2 - (y / height) * np.pi # π/2 to -π/2
|
|
return lat, lon
|
|
|
|
@staticmethod
|
|
def spherical_to_cartesian(lat: np.ndarray, lon: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
|
|
"""Convert spherical coordinates to 3D cartesian coordinates"""
|
|
x = np.cos(lat) * np.cos(lon)
|
|
y = np.cos(lat) * np.sin(lon)
|
|
z = np.sin(lat)
|
|
return x, y, z
|
|
|
|
@staticmethod
|
|
def cartesian_to_spherical(x: np.ndarray, y: np.ndarray, z: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
|
|
"""Convert 3D cartesian coordinates to spherical coordinates"""
|
|
lat = np.arcsin(np.clip(z, -1, 1))
|
|
lon = np.arctan2(y, x)
|
|
return lat, lon
|
|
|
|
@staticmethod
|
|
def spherical_to_equirectangular(lat: np.ndarray, lon: np.ndarray, width: int, height: int) -> Tuple[np.ndarray, np.ndarray]:
|
|
"""Convert spherical coordinates to equirectangular coordinates"""
|
|
# Normalize longitude to [0, 2π] then to [0, width]
|
|
lon_norm = (lon + np.pi) / (2 * np.pi)
|
|
x = lon_norm * width
|
|
|
|
# Convert latitude from [π/2, -π/2] to [0, height]
|
|
lat_norm = (np.pi/2 - lat) / np.pi
|
|
y = lat_norm * height
|
|
|
|
return x, y
|
|
|
|
@staticmethod
|
|
def create_rotation_matrix(yaw: float, pitch: float, roll: float) -> np.ndarray:
|
|
"""Create 3D rotation matrix from yaw, pitch, roll angles in degrees"""
|
|
yaw_rad = np.radians(yaw)
|
|
pitch_rad = np.radians(pitch)
|
|
roll_rad = np.radians(roll)
|
|
|
|
# Using scipy for robust rotation matrix creation
|
|
r = Rotation.from_euler('zyx', [yaw_rad, pitch_rad, roll_rad])
|
|
return r.as_matrix()
|
|
|
|
@staticmethod
|
|
def apply_rotation(x: np.ndarray, y: np.ndarray, z: np.ndarray, rotation_matrix: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
|
|
"""Apply rotation matrix to 3D points"""
|
|
points = np.stack([x, y, z], axis=-1)
|
|
original_shape = points.shape[:-1]
|
|
points_flat = points.reshape(-1, 3)
|
|
|
|
rotated_points = points_flat @ rotation_matrix.T
|
|
rotated_points = rotated_points.reshape(*original_shape, 3)
|
|
|
|
return rotated_points[..., 0], rotated_points[..., 1], rotated_points[..., 2]
|
|
|
|
@staticmethod
|
|
def bilinear_interpolate(image: np.ndarray, x: np.ndarray, y: np.ndarray) -> np.ndarray:
|
|
"""Perform bilinear interpolation on image at given coordinates"""
|
|
height, width = image.shape[:2]
|
|
|
|
# Clamp coordinates to valid range
|
|
x = np.clip(x, 0, width - 1)
|
|
y = np.clip(y, 0, height - 1)
|
|
|
|
x0 = np.floor(x).astype(int)
|
|
x1 = np.minimum(x0 + 1, width - 1)
|
|
y0 = np.floor(y).astype(int)
|
|
y1 = np.minimum(y0 + 1, height - 1)
|
|
|
|
# Get fractional parts
|
|
dx = x - x0
|
|
dy = y - y0
|
|
|
|
# Get pixel values at corners
|
|
if len(image.shape) == 3: # Color image
|
|
I00 = image[y0, x0]
|
|
I01 = image[y1, x0]
|
|
I10 = image[y0, x1]
|
|
I11 = image[y1, x1]
|
|
|
|
# Expand dimensions for broadcasting
|
|
dx = dx[..., np.newaxis]
|
|
dy = dy[..., np.newaxis]
|
|
else: # Grayscale image
|
|
I00 = image[y0, x0]
|
|
I01 = image[y1, x0]
|
|
I10 = image[y0, x1]
|
|
I11 = image[y1, x1]
|
|
|
|
# Bilinear interpolation
|
|
I0 = I00 * (1 - dx) + I10 * dx
|
|
I1 = I01 * (1 - dx) + I11 * dx
|
|
interpolated = I0 * (1 - dy) + I1 * dy
|
|
|
|
return interpolated
|
|
|
|
@classmethod
|
|
def rotate_equirectangular(cls, image: np.ndarray, yaw: float = 0, pitch: float = 0, roll: float = 0, horizon_offset: float = 0) -> np.ndarray:
|
|
"""Rotate an equirectangular image with optional horizon adjustment"""
|
|
height, width = image.shape[:2]
|
|
|
|
# Create coordinate grids
|
|
x_coords, y_coords = np.meshgrid(np.arange(width), np.arange(height))
|
|
|
|
# Convert to spherical coordinates
|
|
lat, lon = cls.equirectangular_to_spherical(x_coords, y_coords, width, height)
|
|
|
|
# Apply horizon offset
|
|
lat += np.radians(horizon_offset)
|
|
lat = np.clip(lat, -np.pi/2, np.pi/2) # Clamp to valid latitude range
|
|
|
|
# Convert to 3D cartesian
|
|
x_3d, y_3d, z_3d = cls.spherical_to_cartesian(lat, lon)
|
|
|
|
# Create and apply rotation matrix
|
|
rotation_matrix = cls.create_rotation_matrix(yaw, pitch, roll)
|
|
x_rot, y_rot, z_rot = cls.apply_rotation(x_3d, y_3d, z_3d, rotation_matrix)
|
|
|
|
# Convert back to spherical
|
|
lat_rot, lon_rot = cls.cartesian_to_spherical(x_rot, y_rot, z_rot)
|
|
|
|
# Convert back to equirectangular coordinates
|
|
x_new, y_new = cls.spherical_to_equirectangular(lat_rot, lon_rot, width, height)
|
|
|
|
# Handle longitude wrapping
|
|
x_new = np.mod(x_new, width)
|
|
|
|
# Interpolate to get final image
|
|
rotated_image = cls.bilinear_interpolate(image, x_new, y_new)
|
|
|
|
return rotated_image
|
|
|
|
@staticmethod
|
|
def crop_to_180(image: np.ndarray, output_width: Optional[int] = None, output_height: Optional[int] = None) -> np.ndarray:
|
|
"""Crop equirectangular image to 180 degree field of view"""
|
|
height, width = image.shape[:2]
|
|
|
|
# Calculate crop boundaries (center 180 degrees out of 360)
|
|
start_x = width // 4 # Start at 90 degrees
|
|
end_x = start_x + width // 2 # End at 270 degrees (180 degree span)
|
|
|
|
# Crop the image
|
|
cropped = image[:, start_x:end_x]
|
|
|
|
# Resize if output dimensions specified
|
|
if output_width is not None and output_height is not None:
|
|
if len(image.shape) == 3:
|
|
cropped = cv2.resize(cropped, (output_width, output_height), interpolation=cv2.INTER_CUBIC)
|
|
else:
|
|
cropped = cv2.resize(cropped, (output_width, output_height), interpolation=cv2.INTER_CUBIC)
|
|
|
|
return cropped
|
|
|
|
@classmethod
|
|
def process_equirectangular(cls,
|
|
image: np.ndarray,
|
|
yaw: float = 0,
|
|
pitch: float = 0,
|
|
roll: float = 0,
|
|
horizon_offset: float = 0,
|
|
crop_to_180: bool = False,
|
|
output_width: Optional[int] = None,
|
|
output_height: Optional[int] = None) -> np.ndarray:
|
|
"""Main processing function that combines all operations"""
|
|
|
|
# Apply rotation and horizon adjustment
|
|
if yaw != 0 or pitch != 0 or roll != 0 or horizon_offset != 0:
|
|
processed_image = cls.rotate_equirectangular(image, yaw, pitch, roll, horizon_offset)
|
|
else:
|
|
processed_image = image.copy()
|
|
|
|
# Apply 180-degree cropping if requested
|
|
if crop_to_180:
|
|
processed_image = cls.crop_to_180(processed_image, output_width, output_height)
|
|
|
|
return processed_image |