280 lines
9.1 KiB
Python
280 lines
9.1 KiB
Python
"""SamplePanorama node for ComfyUI-Sharp.
|
||
|
||
Samples perspective cutouts from an equirectangular panorama for use with SHARP.
|
||
"""
|
||
|
||
import math
|
||
|
||
import numpy as np
|
||
import torch
|
||
import torch.nn.functional as F
|
||
|
||
|
||
def create_rotation_matrix(yaw: float, pitch: float) -> torch.Tensor:
|
||
"""Create a rotation matrix from yaw and pitch angles.
|
||
|
||
Args:
|
||
yaw: Horizontal angle in radians (0 = forward, positive = right)
|
||
pitch: Vertical angle in radians (0 = horizon, positive = up)
|
||
|
||
Returns:
|
||
3x3 rotation matrix (camera-to-world)
|
||
"""
|
||
# Rotation around Y axis (yaw - looking left/right)
|
||
cy, sy = math.cos(yaw), math.sin(yaw)
|
||
R_yaw = torch.tensor([
|
||
[cy, 0, sy],
|
||
[0, 1, 0],
|
||
[-sy, 0, cy]
|
||
], dtype=torch.float32)
|
||
|
||
# Rotation around X axis (pitch - looking up/down)
|
||
cp, sp = math.cos(pitch), math.sin(pitch)
|
||
R_pitch = torch.tensor([
|
||
[1, 0, 0],
|
||
[0, cp, -sp],
|
||
[0, sp, cp]
|
||
], dtype=torch.float32)
|
||
|
||
# Combined: first pitch, then yaw
|
||
R = R_yaw @ R_pitch
|
||
return R
|
||
|
||
|
||
def sample_perspective_from_equirectangular(
|
||
panorama: torch.Tensor,
|
||
yaw: float,
|
||
pitch: float,
|
||
fov_radians: float,
|
||
output_size: int,
|
||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||
"""Sample a perspective view from an equirectangular panorama.
|
||
|
||
Args:
|
||
panorama: Equirectangular image [H, W, 3] or [B, H, W, 3]
|
||
yaw: Horizontal angle in radians (0 = center, positive = right)
|
||
pitch: Vertical angle in radians (0 = horizon, positive = up)
|
||
fov_radians: Field of view in radians
|
||
output_size: Output image size (square)
|
||
|
||
Returns:
|
||
Tuple of (perspective_image, extrinsics_4x4, intrinsics_4x4)
|
||
"""
|
||
if panorama.dim() == 4:
|
||
panorama = panorama[0] # Take first if batched
|
||
|
||
H, W, C = panorama.shape
|
||
device = panorama.device
|
||
|
||
# Compute focal length in pixels
|
||
f_px = (output_size / 2) / math.tan(fov_radians / 2)
|
||
|
||
# Principal point at center
|
||
cx = (output_size - 1) / 2
|
||
cy = (output_size - 1) / 2
|
||
|
||
# Create pixel grid for output image
|
||
u = torch.arange(output_size, dtype=torch.float32, device=device)
|
||
v = torch.arange(output_size, dtype=torch.float32, device=device)
|
||
uu, vv = torch.meshgrid(u, v, indexing='xy') # [H, W]
|
||
|
||
# Convert to camera-space ray directions
|
||
# d_cam = [(u - cx) / f, (v - cy) / f, 1]
|
||
dx = (uu - cx) / f_px
|
||
dy = (vv - cy) / f_px
|
||
dz = torch.ones_like(dx)
|
||
|
||
# Stack and normalize
|
||
rays_cam = torch.stack([dx, dy, dz], dim=-1) # [H, W, 3]
|
||
rays_cam = F.normalize(rays_cam, dim=-1)
|
||
|
||
# Rotate rays to world space
|
||
R = create_rotation_matrix(yaw, pitch).to(device) # [3, 3]
|
||
rays_world = torch.einsum('ij,hwj->hwi', R, rays_cam) # [H, W, 3]
|
||
|
||
# Convert world rays to spherical coordinates
|
||
# x = right, y = up, z = forward
|
||
rx, ry, rz = rays_world[..., 0], rays_world[..., 1], rays_world[..., 2]
|
||
|
||
# Yaw: angle around Y axis (horizontal)
|
||
ray_yaw = torch.atan2(rx, rz) # [-pi, pi]
|
||
|
||
# Pitch: angle from horizon (vertical)
|
||
ray_pitch = torch.asin(torch.clamp(ry, -1, 1)) # [-pi/2, pi/2]
|
||
|
||
# Map to equirectangular pixel coordinates
|
||
# Yaw: -pi..pi -> 0..W
|
||
# Pitch: -pi/2..pi/2 -> H..0 (note: inverted, top of image is up)
|
||
eq_x = (ray_yaw / math.pi + 1) * (W - 1) / 2
|
||
eq_y = (0.5 - ray_pitch / math.pi) * (H - 1)
|
||
|
||
# Normalize to [-1, 1] for grid_sample
|
||
grid_x = eq_x / (W - 1) * 2 - 1
|
||
grid_y = eq_y / (H - 1) * 2 - 1
|
||
grid = torch.stack([grid_x, grid_y], dim=-1).unsqueeze(0) # [1, H, W, 2]
|
||
|
||
# Sample from panorama
|
||
panorama_nchw = panorama.permute(2, 0, 1).unsqueeze(0) # [1, C, H, W]
|
||
sampled = F.grid_sample(
|
||
panorama_nchw,
|
||
grid,
|
||
mode='bilinear',
|
||
padding_mode='border',
|
||
align_corners=True
|
||
)
|
||
perspective = sampled[0].permute(1, 2, 0) # [H, W, C]
|
||
|
||
# Build intrinsics matrix (4x4)
|
||
intrinsics = torch.tensor([
|
||
[f_px, 0, cx, 0],
|
||
[0, f_px, cy, 0],
|
||
[0, 0, 1, 0],
|
||
[0, 0, 0, 1],
|
||
], dtype=torch.float32, device=device)
|
||
|
||
# Build extrinsics matrix (4x4, world-to-camera)
|
||
# World-to-camera is the inverse (transpose for rotation)
|
||
R_w2c = R.T
|
||
extrinsics = torch.eye(4, dtype=torch.float32, device=device)
|
||
extrinsics[:3, :3] = R_w2c
|
||
# No translation (camera at origin)
|
||
|
||
return perspective, extrinsics, intrinsics
|
||
|
||
|
||
class SamplePanorama:
|
||
"""Sample perspective cutouts from an equirectangular panorama.
|
||
|
||
Automatically calculates the number of samples needed to cover the full
|
||
360° x 180° panorama based on FOV and overlap settings.
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"panorama": ("IMAGE",),
|
||
"fov_degrees": ("FLOAT", {
|
||
"default": 65.0,
|
||
"min": 30.0,
|
||
"max": 120.0,
|
||
"step": 1.0,
|
||
"tooltip": "Field of view in degrees for each perspective cutout"
|
||
}),
|
||
"overlap_percent": ("FLOAT", {
|
||
"default": 10.0,
|
||
"min": 0.0,
|
||
"max": 50.0,
|
||
"step": 1.0,
|
||
"tooltip": "Overlap between adjacent samples as percentage of FOV"
|
||
}),
|
||
"output_size": ("INT", {
|
||
"default": 1536,
|
||
"min": 256,
|
||
"max": 2048,
|
||
"step": 64,
|
||
"tooltip": "Output resolution for each perspective image (square). 1536 is SHARP's native resolution."
|
||
}),
|
||
},
|
||
"optional": {
|
||
"skip_poles": ("BOOLEAN", {
|
||
"default": True,
|
||
"tooltip": "Skip samples pointing straight up/down (often low quality in panoramas)"
|
||
}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "EXTRINSICS", "INTRINSICS", "INT", "INT")
|
||
RETURN_NAMES = ("images", "extrinsics", "intrinsics", "num_horizontal", "num_vertical")
|
||
FUNCTION = "sample"
|
||
CATEGORY = "SHARP"
|
||
DESCRIPTION = "Sample perspective views from a 360° equirectangular panorama for SHARP processing."
|
||
|
||
def sample(
|
||
self,
|
||
panorama: torch.Tensor,
|
||
fov_degrees: float = 65.0,
|
||
overlap_percent: float = 10.0,
|
||
output_size: int = 1536,
|
||
skip_poles: bool = True,
|
||
):
|
||
"""Sample perspective cutouts covering the full panorama."""
|
||
|
||
fov_radians = math.radians(fov_degrees)
|
||
|
||
# Calculate step angle based on overlap
|
||
step_degrees = fov_degrees * (1 - overlap_percent / 100)
|
||
step_radians = math.radians(step_degrees)
|
||
|
||
# Calculate number of samples needed
|
||
num_horizontal = math.ceil(360 / step_degrees)
|
||
|
||
# For vertical, we cover -90° to +90° (180° total)
|
||
# But if skip_poles, we limit to avoid extreme angles
|
||
if skip_poles:
|
||
# Skip top and bottom 15°
|
||
vertical_range = 150 # degrees
|
||
vertical_start = -75 # degrees
|
||
else:
|
||
vertical_range = 180
|
||
vertical_start = -90
|
||
|
||
num_vertical = max(1, math.ceil(vertical_range / step_degrees))
|
||
|
||
print(f"[SamplePanorama] FOV: {fov_degrees}°, Overlap: {overlap_percent}%")
|
||
print(f"[SamplePanorama] Step angle: {step_degrees:.1f}°")
|
||
print(f"[SamplePanorama] Samples: {num_horizontal} horizontal × {num_vertical} vertical = {num_horizontal * num_vertical} total")
|
||
|
||
# Handle batch dimension
|
||
if panorama.dim() == 3:
|
||
panorama = panorama.unsqueeze(0)
|
||
|
||
# Take first image if batch
|
||
pano = panorama[0] # [H, W, C]
|
||
|
||
all_images = []
|
||
all_extrinsics = []
|
||
intrinsics = None
|
||
|
||
for v_idx in range(num_vertical):
|
||
# Calculate pitch angle
|
||
pitch_degrees = vertical_start + (v_idx + 0.5) * step_degrees
|
||
pitch_radians = math.radians(pitch_degrees)
|
||
|
||
for h_idx in range(num_horizontal):
|
||
# Calculate yaw angle
|
||
yaw_degrees = -180 + (h_idx + 0.5) * step_degrees
|
||
yaw_radians = math.radians(yaw_degrees)
|
||
|
||
# Sample perspective view
|
||
perspective, extrinsics, intr = sample_perspective_from_equirectangular(
|
||
pano,
|
||
yaw_radians,
|
||
pitch_radians,
|
||
fov_radians,
|
||
output_size,
|
||
)
|
||
|
||
all_images.append(perspective)
|
||
all_extrinsics.append(extrinsics)
|
||
|
||
if intrinsics is None:
|
||
intrinsics = intr
|
||
|
||
# Stack results
|
||
images_batch = torch.stack(all_images, dim=0) # [N, H, W, C]
|
||
extrinsics_batch = torch.stack(all_extrinsics, dim=0) # [N, 4, 4]
|
||
|
||
print(f"[SamplePanorama] Output shape: {images_batch.shape}")
|
||
|
||
return (images_batch, extrinsics_batch, intrinsics, num_horizontal, num_vertical)
|
||
|
||
|
||
NODE_CLASS_MAPPINGS = {
|
||
"SamplePanorama": SamplePanorama,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"SamplePanorama": "Sample Panorama (Equirect → Perspective)",
|
||
}
|