Spaces:
Running
Running
| import numpy as np | |
| from scipy.ndimage import map_coordinates | |
| def xyzcube(face_w): | |
| ''' | |
| Return the xyz cordinates of the unit cube in [F R B L U D] format. | |
| ''' | |
| out = np.zeros((face_w, face_w * 6, 3), np.float32) | |
| rng = np.linspace(-0.5, 0.5, num=face_w, dtype=np.float32) | |
| grid = np.stack(np.meshgrid(rng, -rng), -1) | |
| # 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 | |
| # Right face (x = 0.5) | |
| out[:, 1*face_w:2*face_w, [2, 1]] = grid | |
| out[:, 1*face_w:2*face_w, 0] = 0.5 | |
| # Back face (z = -0.5) | |
| out[:, 2*face_w:3*face_w, [0, 1]] = grid | |
| out[:, 2*face_w:3*face_w, 2] = -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 | |
| # Up face (y = 0.5) | |
| out[:, 4*face_w:5*face_w, [0, 2]] = grid | |
| out[:, 4*face_w:5*face_w, 1] = 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 | |
| return out | |
| def equirect_uvgrid(h, w): | |
| u = np.linspace(-np.pi, np.pi, num=w, dtype=np.float32) | |
| v = np.linspace(np.pi, -np.pi, num=h, dtype=np.float32) / 2 | |
| return np.stack(np.meshgrid(u, v), axis=-1) | |
| def equirect_facetype(h, w): | |
| ''' | |
| 0F 1R 2B 3L 4U 5D | |
| ''' | |
| tp = np.roll(np.arange(4).repeat(w // 4)[None, :].repeat(h, 0), 3 * w // 8, 1) | |
| # Prepare ceil mask | |
| mask = np.zeros((h, w // 4), np.bool) | |
| idx = np.linspace(-np.pi, np.pi, w // 4) / 4 | |
| idx = h // 2 - np.round(np.arctan(np.cos(idx)) * h / np.pi).astype(int) | |
| for i, j in enumerate(idx): | |
| mask[:j, i] = 1 | |
| mask = np.roll(np.concatenate([mask] * 4, 1), 3 * w // 8, 1) | |
| tp[mask] = 4 | |
| tp[np.flip(mask, 0)] = 5 | |
| return tp.astype(np.int32) | |
| def xyzpers(h_fov, v_fov, u, v, out_hw, in_rot): | |
| out = np.ones((*out_hw, 3), np.float32) | |
| x_max = np.tan(h_fov / 2) | |
| y_max = np.tan(v_fov / 2) | |
| 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]) | |
| Ri = rotation_matrix(in_rot, np.array([0, 0, 1.0]).dot(Rx).dot(Ry)) | |
| return out.dot(Rx).dot(Ry).dot(Ri) | |
| def xyz2uv(xyz): | |
| ''' | |
| xyz: ndarray in shape of [..., 3] | |
| ''' | |
| x, y, z = np.split(xyz, 3, axis=-1) | |
| u = np.arctan2(x, z) | |
| c = np.sqrt(x**2 + z**2) | |
| v = np.arctan2(y, c) | |
| return np.concatenate([u, v], axis=-1) | |
| def uv2unitxyz(uv): | |
| u, v = np.split(uv, 2, axis=-1) | |
| y = np.sin(v) | |
| c = np.cos(v) | |
| x = c * np.sin(u) | |
| z = c * np.cos(u) | |
| return np.concatenate([x, y, z], axis=-1) | |
| def uv2coor(uv, h, w): | |
| ''' | |
| uv: ndarray in shape of [..., 2] | |
| h: int, height of the equirectangular image | |
| w: int, width of the equirectangular image | |
| ''' | |
| u, v = np.split(uv, 2, axis=-1) | |
| coor_x = (u / (2 * np.pi) + 0.5) * w - 0.5 | |
| coor_y = (-v / np.pi + 0.5) * h - 0.5 | |
| return np.concatenate([coor_x, coor_y], axis=-1) | |
| def coor2uv(coorxy, h, w): | |
| coor_x, coor_y = np.split(coorxy, 2, axis=-1) | |
| u = ((coor_x + 0.5) / w - 0.5) * 2 * np.pi | |
| v = -((coor_y + 0.5) / h - 0.5) * np.pi | |
| return np.concatenate([u, v], axis=-1) | |
| def sample_equirec(e_img, coor_xy, order): | |
| w = e_img.shape[1] | |
| coor_x, coor_y = np.split(coor_xy, 2, axis=-1) | |
| pad_u = np.roll(e_img[[0]], w // 2, 1) | |
| pad_d = np.roll(e_img[[-1]], w // 2, 1) | |
| e_img = np.concatenate([e_img, pad_d, pad_u], 0) | |
| return map_coordinates(e_img, [coor_y, coor_x], | |
| order=order, mode='wrap')[..., 0] | |
| def sample_cubefaces(cube_faces, tp, coor_y, coor_x, order): | |
| 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])) | |
| 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) | |
| # Pad left right | |
| pad_lr = np.zeros((6, cube_faces.shape[1], 2)) | |
| 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) | |
| return map_coordinates(cube_faces, [tp, coor_y, coor_x], order=order, mode='wrap') | |
| def cube_h2list(cube_h): | |
| assert cube_h.shape[0] * 6 == cube_h.shape[1] | |
| return np.split(cube_h, 6, axis=1) | |
| def cube_list2h(cube_list): | |
| assert len(cube_list) == 6 | |
| assert sum(face.shape == cube_list[0].shape for face in cube_list) == 6 | |
| return np.concatenate(cube_list, axis=1) | |
| def cube_h2dict(cube_h): | |
| cube_list = cube_h2list(cube_h) | |
| return dict([(k, cube_list[i]) | |
| for i, k in enumerate(['F', 'R', 'B', 'L', 'U', 'D'])]) | |
| def cube_dict2h(cube_dict, face_k=['F', 'R', 'B', 'L', 'U', 'D']): | |
| assert len(face_k) == 6 | |
| return cube_list2h([cube_dict[k] for k in face_k]) | |
| def cube_h2dice(cube_h): | |
| assert cube_h.shape[0] * 6 == cube_h.shape[1] | |
| w = cube_h.shape[0] | |
| 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 | |
| sxy = [(1, 1), (2, 1), (3, 1), (0, 1), (1, 0), (1, 2)] | |
| for i, (sx, sy) in enumerate(sxy): | |
| face = cube_list[i] | |
| if i in [1, 2]: | |
| face = np.flip(face, axis=1) | |
| if i == 4: | |
| face = np.flip(face, axis=0) | |
| cube_dice[sy*w:(sy+1)*w, sx*w:(sx+1)*w] = face | |
| return cube_dice | |
| def cube_dice2h(cube_dice): | |
| w = cube_dice.shape[0] // 3 | |
| assert cube_dice.shape[0] == w * 3 and cube_dice.shape[1] == w * 4 | |
| cube_h = np.zeros((w, w * 6, cube_dice.shape[2]), dtype=cube_dice.dtype) | |
| # Order: F R B L U 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] | |
| if i in [1, 2]: | |
| face = np.flip(face, axis=1) | |
| if i == 4: | |
| face = np.flip(face, axis=0) | |
| cube_h[:, i*w:(i+1)*w] = face | |
| return cube_h | |
| def rotation_matrix(rad, ax): | |
| ax = np.array(ax) | |
| assert len(ax.shape) == 1 and ax.shape[0] == 3 | |
| 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]]) | |
| return R |