From 2fa7b565fe5ed5f7f004b640409e765bebd1dba5 Mon Sep 17 00:00:00 2001 From: Dmitry Rude <3952800@gmail.com> Date: Sun, 29 Dec 2024 19:28:15 +0100 Subject: [PATCH 1/4] caching speeds up batch processing --- py360convert/utils.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/py360convert/utils.py b/py360convert/utils.py index 333a916..3acd37d 100644 --- a/py360convert/utils.py +++ b/py360convert/utils.py @@ -1,5 +1,6 @@ from collections.abc import Sequence from enum import IntEnum +from functools import cache from typing import Any, Literal, Optional, TypeVar, Union import numpy as np @@ -80,11 +81,13 @@ def mode_to_order(mode: InterpolationMode) -> int: raise ValueError(f'Unknown mode "{mode}".') from None +@cache def slice_chunk(index: int, width: int, offset=0): start = index * width + offset return slice(start, start + width) +@cache def xyzcube(face_w: int) -> NDArray[np.float32]: """ Return the xyz coordinates of the unit cube in [F R B L U D] format. @@ -148,12 +151,14 @@ def face_slice(index): return out +@cache def equirect_uvgrid(h: int, w: int) -> tuple[NDArray[np.float32], NDArray[np.float32]]: u = np.linspace(-np.pi, np.pi, num=w, dtype=np.float32) v = np.linspace(np.pi / 2, -np.pi / 2, num=h, dtype=np.float32) return np.meshgrid(u, v) # pyright: ignore[reportReturnType] +@cache def equirect_facetype(h: int, w: int) -> NDArray[np.int32]: """Generate a 2D equirectangular segmentation image for each facetype. From 38c81008b9c50557cb659ebcf48c1c2c3c01296d Mon Sep 17 00:00:00 2001 From: Brian Pugh Date: Thu, 2 Jan 2025 10:17:46 -0500 Subject: [PATCH 2/4] use lru_cache(8) instead of cache to bound memory usage. Make returned numpy arrays write=False. --- py360convert/utils.py | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/py360convert/utils.py b/py360convert/utils.py index 3acd37d..a95ae37 100644 --- a/py360convert/utils.py +++ b/py360convert/utils.py @@ -1,6 +1,6 @@ from collections.abc import Sequence from enum import IntEnum -from functools import cache +from functools import lru_cache from typing import Any, Literal, Optional, TypeVar, Union import numpy as np @@ -44,6 +44,7 @@ "quintic", ] DType = TypeVar("DType", bound=np.generic, covariant=True) +_CACHE_SIZE = 8 class Face(IntEnum): @@ -81,13 +82,12 @@ def mode_to_order(mode: InterpolationMode) -> int: raise ValueError(f'Unknown mode "{mode}".') from None -@cache def slice_chunk(index: int, width: int, offset=0): start = index * width + offset return slice(start, start + width) -@cache +@lru_cache(_CACHE_SIZE) def xyzcube(face_w: int) -> NDArray[np.float32]: """ Return the xyz coordinates of the unit cube in [F R B L U D] format. @@ -148,17 +148,23 @@ def face_slice(index): out[:, face_slice(Face.DOWN), Dim.Y] = -0.5 out[:, face_slice(Face.DOWN), Dim.Z] = y + # Since we are using lru_cache, we want the return value to be immutable. + out.setflags(write=False) return out -@cache +@lru_cache(_CACHE_SIZE) def equirect_uvgrid(h: int, w: int) -> tuple[NDArray[np.float32], NDArray[np.float32]]: u = np.linspace(-np.pi, np.pi, num=w, dtype=np.float32) v = np.linspace(np.pi / 2, -np.pi / 2, num=h, dtype=np.float32) - return np.meshgrid(u, v) # pyright: ignore[reportReturnType] + uu, vv = np.meshgrid(u, v) + # Since we are using lru_cache, we want the return value to be immutable. + uu.setflags(write=False) + vv.setflags(write=False) + return uu, vv # pyright: ignore[reportReturnType] -@cache +@lru_cache(_CACHE_SIZE) def equirect_facetype(h: int, w: int) -> NDArray[np.int32]: """Generate a 2D equirectangular segmentation image for each facetype. @@ -230,6 +236,9 @@ def equirect_facetype(h: int, w: int) -> NDArray[np.int32]: tp[:h3, s.stop :][mask[:, :remainder]] = Face.UP # pyright: ignore[reportPossiblyUnboundVariable] tp[-h3:, s.stop :][flip_mask[:, :remainder]] = Face.DOWN # pyright: ignore[reportPossiblyUnboundVariable] + # Since we are using lru_cache, we want the return value to be immutable. + tp.setflags(write=False) + return tp From 1382e987d33422176fe8d2477d33fec0ba4c3987 Mon Sep 17 00:00:00 2001 From: Brian Pugh Date: Thu, 2 Jan 2025 11:04:51 -0500 Subject: [PATCH 3/4] cache EquirecSampler --- py360convert/e2c.py | 7 +----- py360convert/e2p.py | 15 ++++-------- py360convert/utils.py | 56 +++++++++++++++++++++++++++++++++++++++++-- 3 files changed, 60 insertions(+), 18 deletions(-) diff --git a/py360convert/e2c.py b/py360convert/e2c.py index 9ef68c0..5aee826 100644 --- a/py360convert/e2c.py +++ b/py360convert/e2c.py @@ -79,12 +79,7 @@ def e2c( h, w = e_img.shape[:2] order = mode_to_order(mode) - - xyz = xyzcube(face_w) - u, v = xyz2uv(xyz) - coor_x, coor_y = uv2coor(u, v, h, w) - - sampler = EquirecSampler(coor_x, coor_y, order) + sampler = EquirecSampler.from_cubemap(face_w, h, w, order) cubemap = np.stack( [sampler(e_img[..., i]) for i in range(e_img.shape[2])], axis=-1, diff --git a/py360convert/e2p.py b/py360convert/e2p.py index fd7c852..a574739 100644 --- a/py360convert/e2p.py +++ b/py360convert/e2p.py @@ -59,21 +59,16 @@ def e2p( h, w = e_img.shape[:2] if isinstance(fov_deg, Real): - h_fov = v_fov = np.deg2rad(float(fov_deg)) + h_fov = v_fov = float(np.deg2rad(float(fov_deg))) else: h_fov, v_fov = map(np.deg2rad, fov_deg) - in_rot = np.deg2rad(in_rot_deg) - order = mode_to_order(mode) - u = -u_deg * np.pi / 180 - v = v_deg * np.pi / 180 - xyz = xyzpers(h_fov, v_fov, u, v, out_hw, in_rot) - u, v = xyz2uv(xyz) - coor_x, coor_y = uv2coor(u, v, h, w) - - sampler = EquirecSampler(coor_x, coor_y, order) + u = -float(np.deg2rad(u_deg)) + v = float(np.deg2rad(v_deg)) + in_rot = float(np.deg2rad(in_rot_deg)) + sampler = EquirecSampler.from_perspective(h_fov, v_fov, u, v, in_rot, h, w, order) pers_img = np.stack([sampler(e_img[..., i]) for i in range(e_img.shape[2])], axis=-1) return pers_img[..., 0] if squeeze else pers_img diff --git a/py360convert/utils.py b/py360convert/utils.py index a95ae37..acb8634 100644 --- a/py360convert/utils.py +++ b/py360convert/utils.py @@ -242,7 +242,9 @@ def equirect_facetype(h: int, w: int) -> NDArray[np.int32]: return tp -def xyzpers(h_fov: float, v_fov: float, u: float, v: float, out_hw: tuple[int, int], in_rot: float) -> NDArray: +def xyzpers( + h_fov: float, v_fov: float, u: float, v: float, out_hw: tuple[int, int], in_rot: float +) -> NDArray[np.float32]: out = np.ones((*out_hw, 3), np.float32) x_max = np.tan(h_fov / 2) @@ -254,7 +256,7 @@ def xyzpers(h_fov: float, v_fov: float, u: float, v: float, out_hw: tuple[int, i 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) + return out.dot(Rx).dot(Ry).dot(Ri).astype(np.float32) def xyz2uv(xyz: NDArray[DType]) -> tuple[NDArray[DType], NDArray[DType]]: @@ -394,6 +396,56 @@ def _pad(self, img: NDArray[DType]) -> NDArray[DType]: padded[-1, :] = np.roll(img[[-1]], w // 2, 1) return padded + @classmethod + @lru_cache(_CACHE_SIZE) + def from_cubemap(cls, face_w: int, h: int, w: int, order: int): + """Construct a EquirecSampler from cubemap specs. + + Parameters + ---------- + face_w: int + Length of each face of the output cubemap. + h: int + Height of input equirec image. + w: int + Width of input equirec image. + order: int + The order of the spline interpolation. See ``scipy.ndimage.map_coordinates``. + """ + xyz = xyzcube(face_w) + u, v = xyz2uv(xyz) + coor_x, coor_y = uv2coor(u, v, h, w) + return cls(coor_x, coor_y, order=order) + + @classmethod + @lru_cache(_CACHE_SIZE) + def from_perspective(cls, h_fov: float, v_fov: float, u, v, in_rot: float, h: int, w: int, order: int): + """Construct a EquirecSampler from perspective specs. + + Parameters + ---------- + h_fov: float + Horizontal field of view in radians. + v_fov: float + Horizontal field of view in radians. + u: float + Horizontal viewing angle in radians + v: float + Vertical viewing angle in radians + in_rot: float + Inplane rotation in radians. + h: int + Height of input equirec image. + w: int + Width of input equirec image. + order: int + The order of the spline interpolation. See ``scipy.ndimage.map_coordinates``. + """ + xyz = xyzpers(h_fov, v_fov, u, v, (h, w), in_rot) + u, v = xyz2uv(xyz) + coor_x, coor_y = uv2coor(u, v, h, w) + return cls(coor_x, coor_y, order=order) + class CubeFaceSampler: """Arranged as a class so coordinate computations can be re-used across multiple image interpolations.""" From d278d1baffa116cc2cdc4d1a827bee6b4b8bb4b8 Mon Sep 17 00:00:00 2001 From: Brian Pugh Date: Thu, 2 Jan 2025 11:27:05 -0500 Subject: [PATCH 4/4] cache CubemapSampler --- py360convert/c2e.py | 39 +------------------------------ py360convert/e2c.py | 3 --- py360convert/e2p.py | 3 --- py360convert/utils.py | 53 +++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 54 insertions(+), 44 deletions(-) diff --git a/py360convert/c2e.py b/py360convert/c2e.py index e13da3d..571cb6d 100644 --- a/py360convert/c2e.py +++ b/py360convert/c2e.py @@ -7,13 +7,10 @@ CubeFaceSampler, CubeFormat, DType, - Face, InterpolationMode, cube_dice2list, cube_dict2list, cube_h2list, - equirect_facetype, - equirect_uvgrid, mode_to_order, ) @@ -127,43 +124,9 @@ def c2e( raise ValueError("Cubemap faces must be square.") face_w = cube_faces.shape[2] - u, v = equirect_uvgrid(h, w) - - # Get face id to each pixel: 0F 1R 2B 3L 4U 5D - tp = equirect_facetype(h, w) - - coor_x = np.empty((h, w), dtype=np.float32) - coor_y = np.empty((h, w), dtype=np.float32) - face_w2 = face_w / 2 - - # Middle band (front/right/back/left) - mask = tp < Face.UP - angles = u[mask] - (np.pi / 2 * tp[mask]) - tan_angles = np.tan(angles) - cos_angles = np.cos(angles) - tan_v = np.tan(v[mask]) - - coor_x[mask] = face_w2 * tan_angles - coor_y[mask] = -face_w2 * tan_v / cos_angles - - mask = tp == Face.UP - c = face_w2 * np.tan(np.pi / 2 - v[mask]) - coor_x[mask] = c * np.sin(u[mask]) - coor_y[mask] = c * np.cos(u[mask]) - - mask = tp == Face.DOWN - c = face_w2 * np.tan(np.pi / 2 - np.abs(v[mask])) - coor_x[mask] = c * np.sin(u[mask]) - coor_y[mask] = -c * np.cos(u[mask]) - - # Final renormalize - coor_x += face_w2 - coor_y += face_w2 - coor_x.clip(0, face_w, out=coor_x) - coor_y.clip(0, face_w, out=coor_y) + sampler = CubeFaceSampler.from_equirec(face_w, h, w, order) equirec = np.empty((h, w, cube_faces.shape[3]), dtype=cube_faces[0].dtype) - sampler = CubeFaceSampler(tp, coor_x, coor_y, order, face_w, face_w) for i in range(cube_faces.shape[3]): equirec[..., i] = sampler(cube_faces[..., i]) diff --git a/py360convert/e2c.py b/py360convert/e2c.py index 5aee826..7fa2346 100644 --- a/py360convert/e2c.py +++ b/py360convert/e2c.py @@ -12,9 +12,6 @@ cube_h2dict, cube_h2list, mode_to_order, - uv2coor, - xyz2uv, - xyzcube, ) diff --git a/py360convert/e2p.py b/py360convert/e2p.py index a574739..1c93695 100644 --- a/py360convert/e2p.py +++ b/py360convert/e2p.py @@ -9,9 +9,6 @@ EquirecSampler, InterpolationMode, mode_to_order, - uv2coor, - xyz2uv, - xyzpers, ) diff --git a/py360convert/utils.py b/py360convert/utils.py index acb8634..bece5bc 100644 --- a/py360convert/utils.py +++ b/py360convert/utils.py @@ -577,6 +577,59 @@ def _pad(self, cube_faces: NDArray[DType]) -> NDArray[DType]: return padded + @classmethod + @lru_cache(_CACHE_SIZE) + def from_equirec(cls, face_w: int, h: int, w: int, order: int): + """Construct a CubemapSampler from equirectangular specs. + + Parameters + ---------- + face_w: int + Length of each face of the input cubemap. + h: int + Output equirectangular image height. + w: int + Output equirectangular image width. + order: int + The order of the spline interpolation. See ``scipy.ndimage.map_coordinates``. + """ + u, v = equirect_uvgrid(h, w) + + # Get face id to each pixel: 0F 1R 2B 3L 4U 5D + tp = equirect_facetype(h, w) + + coor_x = np.empty((h, w), dtype=np.float32) + coor_y = np.empty((h, w), dtype=np.float32) + face_w2 = face_w / 2 + + # Middle band (front/right/back/left) + mask = tp < Face.UP + angles = u[mask] - (np.pi / 2 * tp[mask]) + tan_angles = np.tan(angles) + cos_angles = np.cos(angles) + tan_v = np.tan(v[mask]) + + coor_x[mask] = face_w2 * tan_angles + coor_y[mask] = -face_w2 * tan_v / cos_angles + + mask = tp == Face.UP + c = face_w2 * np.tan(np.pi / 2 - v[mask]) + coor_x[mask] = c * np.sin(u[mask]) + coor_y[mask] = c * np.cos(u[mask]) + + mask = tp == Face.DOWN + c = face_w2 * np.tan(np.pi / 2 - np.abs(v[mask])) + coor_x[mask] = c * np.sin(u[mask]) + coor_y[mask] = -c * np.cos(u[mask]) + + # Final renormalize + coor_x += face_w2 + coor_y += face_w2 + coor_x.clip(0, face_w, out=coor_x) + coor_y.clip(0, face_w, out=coor_y) + + return cls(tp, coor_x, coor_y, order, face_w, face_w) + def cube_h2list(cube_h: NDArray[DType]) -> list[NDArray[DType]]: """Split an image into a list of 6 faces."""