lazy python script to convert a skysphere images to 6 skybox images. probably can't smooth pixels from small images very well

skysphere-to-skybox.py

import numpy as np
from PIL import Image


def raycast_box_imgs(sphere_img, new_img_size, out_dir):
    """renders and saves all 6 sides of the skybox"""
    rays = create_rays(new_img_size)
    rotations = {
        "ft": np.eye(3),
        "lt": rot_y_mat(-.5 * np.pi),
        "rt": rot_y_mat(.5 * np.pi),
        "bk": rot_y_mat(np.pi),
        "up": rot_x_mat(.5 * np.pi),
        "dn": rot_x_mat(-.5 * np.pi)
    }
    for direction, rotation in rotations.items():
        rotated_rays = np.apply_along_axis(lambda x: np.dot(rotation, x), 2, rays)
        box_img = create_raycast_image(sphere_img, rotated_rays, new_img_size)
        box_img.save(f"{out_dir}/skybox-{direction}.png")


def create_raycast_image(sphere_img, rays, new_img_size):
    """create a skybox image reading pixels from the skysphere img"""
    box_img = Image.new('RGB', (new_img_size, new_img_size))
    
    for y in range(new_img_size):
        for x in range(new_img_size):
            uv_x, uv_y = get_uv_coord(rays[x][y])
            box_img.putpixel((x, y), get_pixel(sphere_img, uv_x, uv_y))
    return box_img


def create_rays(img_size):
    """create ray directions for all pixels of a skybox image"""
    rays = np.zeros((img_size, img_size, 3))

    for x in range(img_size):
        for y in range(img_size):
            rays[x, y] = np.array([(x + .5) / img_size - .5, (y + .5) / img_size - .5, .5])
    return rays


def get_uv_coord(ray):
    """returns the uv coordinate for a ray from the middle of the sphere"""
    ray /= np.linalg.norm(ray)
    yaw = np.arctan2(ray[0], ray[2])
    pitch = np.arcsin(ray[1])
    uv_x = .5 + yaw / (2 * np.pi)
    uv_y = .5 + pitch / np.pi
    return uv_x, uv_y


def get_pixel(img, uv_x, uv_y):
    """returns pixel of image by uv coordinate"""
    return img.getpixel((uv_x * (img.width - 1), uv_y * (img.height - 1)))


def rot_x_mat(theta):
    """creates rotation matrix for rotation around x axis (only in 90 degree steps)"""
    sin_theta = np.round(np.sin(theta))
    cos_theta = np.round(np.cos(theta))
    return np.array(
        [[1, 0, 0],
        [0, cos_theta, -sin_theta],
        [0, sin_theta, cos_theta]])


def rot_y_mat(theta):
    sin_theta = np.round(np.sin(theta))
    cos_theta = np.round(np.cos(theta))
    return np.array(
        [[cos_theta, 0, sin_theta],
        [0, 1, 0],
        [-sin_theta, 0, cos_theta]])


if __name__ == "__main__":
    sphere_img = Image.open("./sphere-milkyway.png")
    box_imgs = raycast_box_imgs(sphere_img, sphere_img.height // 2, "./skybox")