File size: 3,826 Bytes
f4d2177 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
import numpy as np
from scipy.ndimage import map_coordinates
import cv2
# Based on https://github.com/sunset1995/py360convert
class Equirec2Cube:
def __init__(self, equ_h, equ_w, face_w):
'''
equ_h: int, height of the equirectangular image
equ_w: int, width of the equirectangular image
face_w: int, the length of each face of the cubemap
'''
self.equ_h = equ_h
self.equ_w = equ_w
self.face_w = face_w
self._xyzcube()
self._xyz2coor()
# For convert R-distance to Z-depth for CubeMaps
cosmap = 1 / np.sqrt((2 * self.grid[..., 0]) ** 2 + (2 * self.grid[..., 1]) ** 2 + 1)
self.cosmaps = np.concatenate(6 * [cosmap], axis=1)[..., np.newaxis]
def _xyzcube(self):
'''
Compute the xyz cordinates of the unit cube in [F R B L U D] format.
'''
self.xyz = np.zeros((self.face_w, self.face_w * 6, 3), np.float32)
rng = np.linspace(-0.5, 0.5, num=self.face_w, dtype=np.float32)
self.grid = np.stack(np.meshgrid(rng, -rng), -1)
# Front face (z = 0.5)
self.xyz[:, 0 * self.face_w:1 * self.face_w, [0, 1]] = self.grid
self.xyz[:, 0 * self.face_w:1 * self.face_w, 2] = 0.5
# Right face (x = 0.5)
self.xyz[:, 1 * self.face_w:2 * self.face_w, [2, 1]] = self.grid[:, ::-1]
self.xyz[:, 1 * self.face_w:2 * self.face_w, 0] = 0.5
# Back face (z = -0.5)
self.xyz[:, 2 * self.face_w:3 * self.face_w, [0, 1]] = self.grid[:, ::-1]
self.xyz[:, 2 * self.face_w:3 * self.face_w, 2] = -0.5
# Left face (x = -0.5)
self.xyz[:, 3 * self.face_w:4 * self.face_w, [2, 1]] = self.grid
self.xyz[:, 3 * self.face_w:4 * self.face_w, 0] = -0.5
# Up face (y = 0.5)
self.xyz[:, 4 * self.face_w:5 * self.face_w, [0, 2]] = self.grid[::-1, :]
self.xyz[:, 4 * self.face_w:5 * self.face_w, 1] = 0.5
# Down face (y = -0.5)
self.xyz[:, 5 * self.face_w:6 * self.face_w, [0, 2]] = self.grid
self.xyz[:, 5 * self.face_w:6 * self.face_w, 1] = -0.5
def _xyz2coor(self):
# x, y, z to longitude and latitude
x, y, z = np.split(self.xyz, 3, axis=-1)
lon = np.arctan2(x, z)
c = np.sqrt(x ** 2 + z ** 2)
lat = np.arctan2(y, c)
# longitude and latitude to equirectangular coordinate
self.coor_x = (lon / (2 * np.pi) + 0.5) * self.equ_w - 0.5
self.coor_y = (-lat / np.pi + 0.5) * self.equ_h - 0.5
def sample_equirec(self, e_img, order=0):
pad_u = np.roll(e_img[[0]], self.equ_w // 2, 1)
pad_d = np.roll(e_img[[-1]], self.equ_w // 2, 1)
e_img = np.concatenate([e_img, pad_d, pad_u], 0)
# pad_l = e_img[:, [0]]
# pad_r = e_img[:, [-1]]
# e_img = np.concatenate([e_img, pad_l, pad_r], 1)
return map_coordinates(e_img, [self.coor_y, self.coor_x],
order=order, mode='wrap')[..., 0]
def run(self, equ_img, equ_dep=None):
h, w = equ_img.shape[:2]
if h != self.equ_h or w != self.equ_w:
equ_img = cv2.resize(equ_img, (self.equ_w, self.equ_h))
if equ_dep is not None:
equ_dep = cv2.resize(equ_dep, (self.equ_w, self.equ_h), interpolation=cv2.INTER_NEAREST)
cube_img = np.stack([self.sample_equirec(equ_img[..., i], order=1)
for i in range(equ_img.shape[2])], axis=-1)
if equ_dep is not None:
cube_dep = np.stack([self.sample_equirec(equ_dep[..., i], order=0)
for i in range(equ_dep.shape[2])], axis=-1)
cube_dep = cube_dep * self.cosmaps
if equ_dep is not None:
return cube_img, cube_dep
else:
return cube_img |