get_frbl_maps.py

from typing import Tuple
import numpy as np


def norm_to_face(face_size: int, mat: np.ndarray) -> np.ndarray:
    """Convert from [-1,1] tp [0, face_size]"""
    return np.clip((mat+1)*0.5*face_size, 0, face_size)


def get_frbl_maps(out_width: int,
                  out_height: int,
                  face_size: int) -> Tuple[np.ndarray, np.ndarray]:
    """Generate xy-maps containing FRONT,RIGHT,BACK,LEFT info"""
    theta = sample_pixels(-np.pi/4, np.pi/4, out_width//4)
    phi = sample_pixels(-np.pi/2, np.pi/2, out_height)

    # Generate xy-maps for the FRONT face
    map_x = norm_to_face(face_size, np.tan(theta))
    map_y = norm_to_face(face_size, np.divide(np.tan(phi)[:, None],
                                              np.cos(theta)[None, :]))

    # Repeat it to get FRONT, RIGHT, BACK and LEFT
    map_x = np.tile(map_x, [4, 1])
    map_x += face_size*np.arange(4)[:, None]  # Add offsets
    map_x = np.tile(map_x.flatten(), [out_height, 1])

    map_y = np.tile(map_y, 4)

    # Roll
    map_x = np.roll(map_x, 3 * out_width // 8, axis=1)
    map_y = np.roll(map_y, 3 * out_width // 8, axis=1)

    return map_x, map_y