139 lines
4.8 KiB
Python
139 lines
4.8 KiB
Python
import numpy as np
|
|
import cv2
|
|
import math
|
|
|
|
def equirect_to_perspective(img, yaw, pitch, fov=90, size=512):
|
|
h, w = img.shape[:2]
|
|
|
|
fov_rad = math.radians(fov)
|
|
half = math.tan(fov_rad / 2)
|
|
|
|
xs, ys = np.meshgrid(
|
|
np.linspace(-half, half, size),
|
|
np.linspace(half, -half, size) # Reversed: sky at top
|
|
)
|
|
|
|
zs = np.ones_like(xs)
|
|
|
|
dirs = np.stack([xs, -ys, zs], axis=-1)
|
|
norm = np.linalg.norm(dirs, axis=-1, keepdims=True)
|
|
dirs /= norm
|
|
|
|
cy, sy = math.cos(yaw), math.sin(yaw)
|
|
cp, sp = math.cos(pitch), math.sin(pitch)
|
|
|
|
rot_y = np.array([[ cy, 0, sy],
|
|
[ 0, 1, 0],
|
|
[-sy, 0, cy]])
|
|
|
|
rot_x = np.array([[1, 0, 0],
|
|
[0, cp, -sp],
|
|
[0, sp, cp]])
|
|
|
|
dirs = dirs @ rot_y.T @ rot_x.T
|
|
|
|
lon = np.arctan2(dirs[...,0], dirs[...,2])
|
|
lat = np.arcsin(dirs[...,1])
|
|
|
|
u = (lon / (2*np.pi) + 0.5) * w
|
|
v = (0.5 - lat / np.pi) * h
|
|
|
|
map_x = u.astype(np.float32)
|
|
map_y = v.astype(np.float32)
|
|
|
|
return cv2.remap(img, map_x, map_y, cv2.INTER_LINEAR, borderMode=cv2.BORDER_WRAP)
|
|
|
|
|
|
def strip_to_equirect(strip_img, width, height, face_size, order=None, flip_sides=False):
|
|
"""Convert a 6-face strip back to equirectangular projection.
|
|
|
|
Args:
|
|
order: List of indices indicating which face is at each strip position.
|
|
e.g., [0,1,2,3,4,5] means strip position 0 has front, 1 has left, etc.
|
|
Face indices: 0=front, 1=left, 2=back, 3=right, 4=top, 5=bottom
|
|
flip_sides: If True, flip the first 4 faces (horizontal views) vertically before processing
|
|
"""
|
|
if order is None:
|
|
order = [0, 1, 2, 3, 4, 5]
|
|
|
|
# Create reverse mapping: for each face type, which strip position has it?
|
|
face_to_strip = {face_idx: strip_pos for strip_pos, face_idx in enumerate(order)}
|
|
|
|
faces = [
|
|
strip_img[:, i*face_size:(i+1)*face_size]
|
|
for i in range(len(order))
|
|
]
|
|
|
|
# Flip horizontal faces if requested
|
|
if flip_sides:
|
|
for i, face_idx in enumerate(order):
|
|
if face_idx < 4: # front, left, back, right
|
|
faces[i] = cv2.flip(faces[i], 0)
|
|
|
|
# Vectorized approach for speed
|
|
y_coords = np.arange(height)
|
|
x_coords = np.arange(width)
|
|
xx, yy = np.meshgrid(x_coords, y_coords)
|
|
|
|
lat = (0.5 - yy / height) * np.pi
|
|
lon = (xx / width - 0.5) * 2 * np.pi
|
|
|
|
# Direction vector
|
|
dx = np.cos(lat) * np.sin(lon)
|
|
dy = np.sin(lat)
|
|
dz = np.cos(lat) * np.cos(lon)
|
|
|
|
adx, ady, adz = np.abs(dx), np.abs(dy), np.abs(dz)
|
|
|
|
out = np.zeros((height, width, 3), dtype=np.uint8)
|
|
|
|
# Front face (z+) - face index 0
|
|
mask = (adz >= adx) & (adz >= ady) & (dz > 0)
|
|
u = dx[mask] / dz[mask]
|
|
v = dy[mask] / dz[mask]
|
|
px = np.clip(((u * 0.5 + 0.5) * (face_size - 1)).astype(int), 0, face_size - 1)
|
|
py = np.clip(((v * 0.5 + 0.5) * (face_size - 1)).astype(int), 0, face_size - 1)
|
|
out[mask] = faces[face_to_strip[0]][py, px]
|
|
|
|
# Back face (z-) - face index 2
|
|
mask = (adz >= adx) & (adz >= ady) & (dz <= 0)
|
|
u = dx[mask] / dz[mask]
|
|
v = -dy[mask] / dz[mask]
|
|
px = np.clip(((u * 0.5 + 0.5) * (face_size - 1)).astype(int), 0, face_size - 1)
|
|
py = np.clip(((v * 0.5 + 0.5) * (face_size - 1)).astype(int), 0, face_size - 1)
|
|
out[mask] = faces[face_to_strip[2]][py, px]
|
|
|
|
# Right face (x+) - face index 3
|
|
mask = (adx > adz) & (adx >= ady) & (dx > 0)
|
|
u = -dz[mask] / dx[mask]
|
|
v = dy[mask] / dx[mask]
|
|
px = np.clip(((u * 0.5 + 0.5) * (face_size - 1)).astype(int), 0, face_size - 1)
|
|
py = np.clip(((v * 0.5 + 0.5) * (face_size - 1)).astype(int), 0, face_size - 1)
|
|
out[mask] = faces[face_to_strip[3]][py, px]
|
|
|
|
# Left face (x-) - face index 1
|
|
mask = (adx > adz) & (adx >= ady) & (dx <= 0)
|
|
u = -dz[mask] / dx[mask]
|
|
v = -dy[mask] / dx[mask]
|
|
px = np.clip(((u * 0.5 + 0.5) * (face_size - 1)).astype(int), 0, face_size - 1)
|
|
py = np.clip(((v * 0.5 + 0.5) * (face_size - 1)).astype(int), 0, face_size - 1)
|
|
out[mask] = faces[face_to_strip[1]][py, px]
|
|
|
|
# Top face (y+) - face index 4
|
|
mask = (ady > adx) & (ady > adz) & (dy > 0)
|
|
u = dx[mask] / dy[mask]
|
|
v = -dz[mask] / dy[mask] # Flipped to match equirect_to_perspective
|
|
px = np.clip(((u * 0.5 + 0.5) * (face_size - 1)).astype(int), 0, face_size - 1)
|
|
py = np.clip(((v * 0.5 + 0.5) * (face_size - 1)).astype(int), 0, face_size - 1)
|
|
out[mask] = faces[face_to_strip[4]][py, px]
|
|
|
|
# Bottom face (y-) - face index 5
|
|
mask = (ady > adx) & (ady > adz) & (dy <= 0)
|
|
u = -dx[mask] / dy[mask] # Flipped horizontal
|
|
v = -dz[mask] / dy[mask]
|
|
px = np.clip(((u * 0.5 + 0.5) * (face_size - 1)).astype(int), 0, face_size - 1)
|
|
py = np.clip(((v * 0.5 + 0.5) * (face_size - 1)).astype(int), 0, face_size - 1)
|
|
out[mask] = faces[face_to_strip[5]][py, px]
|
|
|
|
return out
|