stich-cubemap.py

from typing import Dict
from PIL import Image
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from stich_360 import create_panorama_gpu

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 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_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]

    # stich the images to create the panorama
    images = dict((face, Image.open(image)) for face, image in zip(FACES, images_files))

    # convert the keys before passing to the stitch_panorama function
    images = {FACE_MAP[face]: image for face, image in images.items()}

    panorama = stich_skybox(images)

    panorama.save(base_dir / 'panorama.png')

def stitch_all_panoramas() -> None:

    with ThreadPoolExecutor(max_workers=4) as executor:
        for panorama in get_panoramas(ROOT_DIR):
            executor.submit(create_panorama_gpu, panorama)


def stich_skybox(images:Dict[str, Image.Image]): 
    width, height = images['front'].size
    panorama = Image.new('RGBA', (width * 4, height * 3))

    '''
    stich the images to create the panorama
        u
    b l f r
        d
    '''
    faces = [
        ['', '', 'u', ''],
        ['b', 'l', 'f', 'r'],
        ['', '', 'd', ''],
    ]
    
    for y, row in enumerate(faces):
        for x, face in enumerate(row):
            if face:
                panorama.paste(images[face], (x * width, y * height))
    
    return panorama

if __name__ == '__main__':
    stitch_panorama_from_name('FF2335BF_DA2B_A7B6_41D4_543AA9F298C6')