Skip to content

Commit

Permalink
Merge pull request #36 from sunset1995/equirect-facetype-docs-and-rou…
Browse files Browse the repository at this point in the history 
…nding

Equirect facetype docs and rounding
  • Loading branch information
@BrianPugh
BrianPugh authored Dec 20, 2024
2 parents 499de10 + 028fd5d commit c1012f2
Show file tree
Hide file tree
Showing 2 changed files with 124 additions and 78 deletions.
10 changes: 3 additions & 7 deletions py360convert/c2e.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -156,12 +156,8 @@ def c2e(
coor_x_norm = (np.clip(coor_x, -0.5, 0.5) + 0.5) * face_w
coor_y_norm = (np.clip(coor_y, -0.5, 0.5) + 0.5) * face_w

equirec = np.stack(
[
sample_cubefaces(cube_faces[..., i], tp, coor_y_norm, coor_x_norm, order=order)
for i in range(cube_faces.shape[3])
],
axis=-1,
)
equirec = np.empty((h, w, cube_faces.shape[3]))
for i in range(cube_faces.shape[3]):
equirec[..., i] = sample_cubefaces(cube_faces[..., i], tp, coor_y_norm, coor_x_norm, order=order)

return equirec[..., 0] if squeeze else equirec
192 changes: 121 additions & 71 deletions py360convert/utils.py
View file
Original file line number Diff line number Diff line change
@@ -1,9 +1,11 @@
from collections.abc import Sequence
from enum import IntEnum
from typing import Any, Literal, Optional, TypeVar, Union

import numpy as np
from numpy.typing import NDArray
from scipy.ndimage import map_coordinates
from scipy.spatial.transform import Rotation

_mode_to_order = {
"nearest": 0,
Expand Down Expand Up @@ -38,6 +40,23 @@
DType = TypeVar("DType", bound=np.generic, covariant=True)


class Face(IntEnum):
"""Face type indexing for numpy vectorization."""

FRONT = 0
RIGHT = 1
BACK = 2
LEFT = 3
UP = 4
DOWN = 5


class Dim(IntEnum):
X = 0
Y = 1
Z = 2


def mode_to_order(mode: InterpolationMode) -> int:
"""Convert a human-friendly interpolation string to integer equivalent.
Expand All @@ -56,6 +75,11 @@ def mode_to_order(mode: InterpolationMode) -> int:
raise ValueError(f'Unknown mode "{mode}".') from None


def slice_chunk(index: int, width: int):
start = index * width
return slice(start, start + width)


def xyzcube(face_w: int) -> NDArray[np.float32]:
"""
Return the xyz coordinates of the unit cube in [F R B L U D] format.
Expand All @@ -78,30 +102,32 @@ def xyzcube(face_w: int) -> NDArray[np.float32]:
rng = np.linspace(-0.5, 0.5, num=face_w, dtype=np.float32)
grid = np.stack(np.meshgrid(rng, -rng), -1)

def face_slice(index):
return slice_chunk(index, face_w)

# Front face (z = 0.5)
out[:, 0 * face_w : 1 * face_w, [0, 1]] = grid
out[:, 0 * face_w : 1 * face_w, 2] = 0.5
out[:, face_slice(Face.FRONT), [Dim.X, Dim.Y]] = grid
out[:, face_slice(Face.FRONT), Dim.Z] = 0.5

# Right face (x = 0.5)
out[:, 1 * face_w : 2 * face_w, [2, 1]] = grid
out[:, 1 * face_w : 2 * face_w, [2, 1]] = np.flip(grid, axis=1)
out[:, 1 * face_w : 2 * face_w, 0] = 0.5
out[:, face_slice(Face.RIGHT), [Dim.Z, Dim.Y]] = np.flip(grid, axis=1)
out[:, face_slice(Face.RIGHT), Dim.X] = 0.5

# Back face (z = -0.5)
out[:, 2 * face_w : 3 * face_w, [0, 1]] = np.flip(grid, axis=1)
out[:, 2 * face_w : 3 * face_w, 2] = -0.5
out[:, face_slice(Face.BACK), [Dim.X, Dim.Y]] = np.flip(grid, axis=1)
out[:, face_slice(Face.BACK), Dim.Z] = -0.5

# Left face (x = -0.5)
out[:, 3 * face_w : 4 * face_w, [2, 1]] = grid
out[:, 3 * face_w : 4 * face_w, 0] = -0.5
out[:, face_slice(Face.LEFT), [Dim.Z, Dim.Y]] = grid
out[:, face_slice(Face.LEFT), Dim.X] = -0.5

# Up face (y = 0.5)
out[:, 4 * face_w : 5 * face_w, [0, 2]] = np.flip(grid, axis=0)
out[:, 4 * face_w : 5 * face_w, 1] = 0.5
out[:, face_slice(Face.UP), [Dim.X, Dim.Z]] = np.flip(grid, axis=0)
out[:, face_slice(Face.UP), Dim.Y] = 0.5

# Down face (y = -0.5)
out[:, 5 * face_w : 6 * face_w, [0, 2]] = grid
out[:, 5 * face_w : 6 * face_w, 1] = -0.5
out[:, face_slice(Face.DOWN), [Dim.X, Dim.Z]] = grid
out[:, face_slice(Face.DOWN), Dim.Y] = -0.5

return out

Expand All @@ -125,6 +151,8 @@ def equirect_facetype(h: int, w: int) -> NDArray[np.int32]:
* 4 - up
* 5 - down
See ``Face``.
Example:
>>> equirect_facetype(8, 12)
Expand Down Expand Up @@ -158,12 +186,14 @@ def equirect_facetype(h: int, w: int) -> NDArray[np.int32]:
mask = np.zeros((h, w // 4), np.bool_)
idx = np.linspace(-np.pi, np.pi, w // 4) / 4
idx = np.round(h / 2 - np.arctan(np.cos(idx)) * h / np.pi).astype(np.int32)
for i, j in enumerate(idx):
mask[:j, i] = True
mask = np.roll(np.concatenate([mask] * 4, 1), 3 * w // 8, 1)

tp[mask] = 4
tp[np.flip(mask, 0)] = 5
row_idx = np.arange(len(mask))[:, None]
mask[row_idx < idx[None, :]] = True

mask = np.roll(np.tile(mask, (1, 4)), 3 * w // 8, 1)

tp[mask] = Face.UP
tp[np.flip(mask, 0)] = Face.DOWN

return tp.astype(np.int32)

Expand All @@ -176,8 +206,8 @@ def xyzpers(h_fov: float, v_fov: float, u: float, v: float, out_hw: tuple[int, i
x_rng = np.linspace(-x_max, x_max, num=out_hw[1], dtype=np.float32)
y_rng = np.linspace(-y_max, y_max, num=out_hw[0], dtype=np.float32)
out[..., :2] = np.stack(np.meshgrid(x_rng, -y_rng), -1)
Rx = rotation_matrix(v, [1, 0, 0])
Ry = rotation_matrix(u, [0, 1, 0])
Rx = rotation_matrix(v, Dim.X)
Ry = rotation_matrix(u, Dim.Y)
Ri = rotation_matrix(in_rot, np.array([0, 0, 1.0]).dot(Rx).dot(Ry))

return out.dot(Rx).dot(Ry).dot(Ri)
Expand Down Expand Up @@ -277,44 +307,56 @@ def sample_equirec(e_img: NDArray[DType], coor_xy: NDArray, order: int) -> NDArr
def sample_cubefaces(
cube_faces: NDArray[DType], tp: NDArray, coor_y: NDArray, coor_x: NDArray, order: int
) -> NDArray[DType]:
cube_faces = cube_faces.copy()
# cube_faces[1] = np.flip(cube_faces[1], 1)
# cube_faces[2] = np.flip(cube_faces[2], 1)
# cube_faces[4] = np.flip(cube_faces[4], 0)

# Pad up down
pad_ud = np.zeros((6, 2, cube_faces.shape[2]), dtype=cube_faces.dtype)
pad_ud[0, 0] = cube_faces[5, 0, :]
pad_ud[0, 1] = cube_faces[4, -1, :]
pad_ud[1, 0] = cube_faces[5, :, -1]
pad_ud[1, 1] = cube_faces[4, ::-1, -1]
pad_ud[2, 0] = cube_faces[5, -1, ::-1]
pad_ud[2, 1] = cube_faces[4, 0, ::-1]
pad_ud[3, 0] = cube_faces[5, ::-1, 0]
pad_ud[3, 1] = cube_faces[4, :, 0]
pad_ud[4, 0] = cube_faces[0, 0, :]
pad_ud[4, 1] = cube_faces[2, 0, ::-1]
pad_ud[5, 0] = cube_faces[2, -1, ::-1]
pad_ud[5, 1] = cube_faces[0, -1, :]
cube_faces = np.concatenate([cube_faces, pad_ud], 1, dtype=cube_faces.dtype)

# Pad left right
pad_lr = np.zeros((6, cube_faces.shape[1], 2), dtype=cube_faces.dtype)
pad_lr[0, :, 0] = cube_faces[1, :, 0]
pad_lr[0, :, 1] = cube_faces[3, :, -1]
pad_lr[1, :, 0] = cube_faces[2, :, 0]
pad_lr[1, :, 1] = cube_faces[0, :, -1]
pad_lr[2, :, 0] = cube_faces[3, :, 0]
pad_lr[2, :, 1] = cube_faces[1, :, -1]
pad_lr[3, :, 0] = cube_faces[0, :, 0]
pad_lr[3, :, 1] = cube_faces[2, :, -1]
pad_lr[4, 1:-1, 0] = cube_faces[1, 0, ::-1]
pad_lr[4, 1:-1, 1] = cube_faces[3, 0, :]
pad_lr[5, 1:-1, 0] = cube_faces[1, -2, :]
pad_lr[5, 1:-1, 1] = cube_faces[3, -2, ::-1]
cube_faces = np.concatenate([cube_faces, pad_lr], 2, dtype=cube_faces.dtype)

return map_coordinates(cube_faces, [tp, coor_y, coor_x], order=order, mode="wrap") # pyright: ignore[reportReturnType]
"""Sample cube faces.
Parameters
----------
cube_faces: numpy.ndarray
(6, H, W) Cube faces.
tp: numpy.ndarray
(H, W) facetype image from ``equirect_facetype``.
coor_y: numpy.ndarray
(H, W) Y coordinates to sample.
coor_x: numpy.ndarray
(H, W) X coordinates to sample.
order: int
The order of the spline interpolation. See ``scipy.ndimage.map_coordinates``.
"""
ABOVE = (0, slice(None))
BELOW = (-1, slice(None))
LEFT = (slice(None), 0)
RIGHT = (slice(None), -1)
padded = np.pad(cube_faces, ((0, 0), (1, 1), (1, 1)), mode="constant")

# Pad above/below
padded[Face.FRONT][ABOVE] = padded[Face.UP, -2, :]
padded[Face.FRONT][BELOW] = padded[Face.DOWN, 1, :]
padded[Face.RIGHT][ABOVE] = padded[Face.UP, ::-1, -2]
padded[Face.RIGHT][BELOW] = padded[Face.DOWN, :, -2]
padded[Face.BACK][ABOVE] = padded[Face.UP, 1, ::-1]
padded[Face.BACK][BELOW] = padded[Face.DOWN, -2, ::-1]
padded[Face.LEFT][ABOVE] = padded[Face.UP, :, 1]
padded[Face.LEFT][BELOW] = padded[Face.DOWN, ::-1, 1]
padded[Face.UP][ABOVE] = padded[Face.BACK, 1, ::-1]
padded[Face.UP][BELOW] = padded[Face.FRONT, 1, :]
padded[Face.DOWN][ABOVE] = padded[Face.FRONT, -2, :]
padded[Face.DOWN][BELOW] = padded[Face.BACK, -2, ::-1]

# Pad left/right
padded[Face.FRONT][LEFT] = padded[Face.LEFT, :, -2]
padded[Face.FRONT][RIGHT] = padded[Face.RIGHT, :, 1]
padded[Face.RIGHT][LEFT] = padded[Face.FRONT, :, -2]
padded[Face.RIGHT][RIGHT] = padded[Face.BACK, :, 1]
padded[Face.BACK][LEFT] = padded[Face.RIGHT, :, -2]
padded[Face.BACK][RIGHT] = padded[Face.LEFT, :, 1]
padded[Face.LEFT][LEFT] = padded[Face.BACK, :, -2]
padded[Face.LEFT][RIGHT] = padded[Face.FRONT, :, 1]
padded[Face.UP][LEFT] = padded[Face.LEFT, 1, :]
padded[Face.UP][RIGHT] = padded[Face.RIGHT, 1, ::-1]
padded[Face.DOWN][LEFT] = padded[Face.LEFT, -2, ::-1]
padded[Face.DOWN][RIGHT] = padded[Face.RIGHT, -2, :]

return map_coordinates(padded, [tp, coor_y + 1, coor_x + 1], order=order) # pyright: ignore[reportReturnType]


def cube_h2list(cube_h: NDArray[DType]) -> list[NDArray[DType]]:
Expand Down Expand Up @@ -359,10 +401,16 @@ def cube_h2dice(cube_h: NDArray[DType]) -> NDArray[DType]:
cube_dice = np.zeros((w * 3, w * 4, cube_h.shape[2]), dtype=cube_h.dtype)
cube_list = cube_h2list(cube_h)
# Order: F R B L U D
# ┌────┐
# │ U │
# ┌────┼────┼────┬────┐
# │ L │ F │ R │ B │
# └────┼────┼────┴────┘
# │ D │
# └────┘
sxy = [(1, 1), (2, 1), (3, 1), (0, 1), (1, 0), (1, 2)]
for i, (sx, sy) in enumerate(sxy):
face = cube_list[i]
cube_dice[sy * w : (sy + 1) * w, sx * w : (sx + 1) * w] = face
for (sx, sy), face in zip(sxy, cube_list):
cube_dice[slice_chunk(sy, w), slice_chunk(sx, w)] = face
return cube_dice


Expand All @@ -374,22 +422,24 @@ def cube_dice2h(cube_dice: NDArray[DType]) -> NDArray[DType]:
raise ValueError(f'Dice width must be 4 "faces" (4x{w}={4*w}) wide.')
cube_h = np.zeros((w, w * 6, cube_dice.shape[2]), dtype=cube_dice.dtype)
# Order: F R B L U D
# ┌────┐
# │ U │
# ┌────┼────┼────┬────┐
# │ L │ F │ R │ B │
# └────┼────┼────┴────┘
# │ D │
# └────┘
sxy = [(1, 1), (2, 1), (3, 1), (0, 1), (1, 0), (1, 2)]
for i, (sx, sy) in enumerate(sxy):
face = cube_dice[sy * w : (sy + 1) * w, sx * w : (sx + 1) * w]
cube_h[:, i * w : (i + 1) * w] = face
cube_h[:, slice_chunk(i, w)] = cube_dice[slice_chunk(sy, w), slice_chunk(sx, w)]
return cube_h


def rotation_matrix(rad: float, ax: Union[NDArray, Sequence]):
def rotation_matrix(rad: float, ax: Union[int, NDArray, Sequence]):
if isinstance(ax, int):
ax = (np.arange(3) == ax).astype(float)
ax = np.array(ax)
if ax.shape != (3,):
raise ValueError(f"ax must be shape (3,); got {ax.shape}")
ax = ax / np.sqrt((ax**2).sum())
R = np.diag([np.cos(rad)] * 3)
R = R + np.outer(ax, ax) * (1.0 - np.cos(rad))

ax = ax * np.sin(rad)
R = R + np.array([[0, -ax[2], ax[1]], [ax[2], 0, -ax[0]], [-ax[1], ax[0], 0]])

R = Rotation.from_rotvec(rad * ax).as_matrix()
return R

0 comments on commit c1012f2

Please sign in to comment.