cube2panorama.py

from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
import subprocess


ROOT_DIR = './panorama'

FACES = ['f', 'b', 'l', 'r', 'u', 'd']

FACE_MAP = {
    'f': 'front',
    'b': 'back',
    'l': 'left',
    'r': 'right',
    'u': 'top',
    'd': 'bottom'
}


def stich_3d(images, *, output_path="out", image_format="PNG") -> None:
    """
    Generates an equirectangular map from six cube face images using cube2sphere CLI,
    with dynamic calculation of the output resolution based on the dimensions of the input images
    to accurately represent a full equirectangular projection of the cube map.

    :param front: Path to the source front cube face image.
    :param back: Path to the source back cube face image.
    :param right: Path to the source right cube face image.
    :param left: Path to the source left cube face image.
    :param top: Path to the source top cube face image.
    :param bottom: Path to the source bottom cube face image.
    :param output_path: Filename for the rendered map. Defaults to "out".
    :param image_format: Format to use when saving the map. Defaults to "PNG".
    """

#    swap front and back
    images['front'], images['back'] = images['back'], images['front']

    # Load one of the cube face images to determine its size
    with Image.open(images['front']) as img:
        width, height = img.size

    # Given the cube faces are square, calculate the appropriate equirectangular projection size
    resolution_width = 4 * width
    resolution_height = 2 * height

    resolution = f"{resolution_width} {resolution_height}"

    # Construct the command to execute
    cmd = [
        "cube2sphere",
        # Split resolution to pass as two arguments
        "-r", *resolution.split(' '),
        "-o", output_path,
        "-f", image_format,
        # "-t", "16",  # Set the number of threads to 16
        # front, back, right, left, top, bottom
        images['front'], images['back'], images['right'], images['left'], images['top'], images['bottom']
    ]

    # print(" ".join(cmd))

    try:
        # Execute the command
        subprocess.run(cmd, check=True)
        print(
            f"Equirectangular map generated successfully at {output_path} in {image_format} format with resolution {resolution}.")
    except subprocess.CalledProcessError as e:
        # If there's an error during execution, print it out
        print(f"An error occurred: {e}")


def get_panoramas(base_dir: str) -> list[str]:
    base_dir = Path(base_dir)
    return [item.name for item in base_dir.iterdir() if item.is_dir()]


def stitch_all_panoramas() -> None:
    with ThreadPoolExecutor(max_workers=5) as executor:
        panoramas = get_panoramas(ROOT_DIR)
        executor.map(stitch_panorama_from_name, panoramas)


def stitch_panorama_from_name(panorama: str) -> None:
    print(f'Stitching {panorama}...')

    base_dir = ROOT_DIR / Path(panorama)
    cube_path = base_dir / 'cube'

    # Get the local images
    images_files = [(cube_path / Path(face + '.jpg')) for face in FACES]

    if not all([image.exists() for image in images_files]):
        raise FileNotFoundError(
            f'Could not find all images for panorama {panorama}.')

    images_files = {
        FACE_MAP[face]: path.as_posix() for face, path in zip(FACES, images_files)
    }

    stich_3d(images_files,
             output_path=f'{base_dir}/panorama.', image_format='PNG')


if __name__ == '__main__':
    stitch_all_panoramas()