HDR全景图切割算法（numpy版本）

原理参考：javascript - Convert 2:1 equirectangular panorama to cube map - Stack Overflow

这里修改成numpy和多线程优化，从用blender的sphere2cube（它本身也是遍历像素操作，而且无法保留亮度信息）几分钟到10s内。

import numpy as np
import sys, math, os
import time
import imageio
import threading
import shutil

imageio.plugins.freeimage.download()

FACES = ['pz','nx','nz','px','ny','py']
# FACES = ['front','left','back','right','bottom','top']

def sphere_to_cube(file_path, resolution=1024, format="hdr", output="output"):
    im = imageio.imread(file_path)
    im = np.array(im)
    hsize = resolution / 2
    pos_array = np.arange(0, resolution * resolution, 1)
    axA_array = np.floor_divide(pos_array, resolution)
    axB_array = np.fmod(pos_array, resolution)
    output_cubes = []
    tasks = []
    for i in range(0, 6):
        output_cube = os.path.join(output, "%s%s.%s" % ("sp_", FACES[i], "hdr"))
        output_cubes.append(output_cube)
        task = threading.Thread(target=sphere_to_cube_process,
                                args=(im, i, axA_array, axB_array, resolution, hsize, format, output_cube),
                                name="sphere_to_cube_" + str(i))
        task.start()  # 启动进程
        tasks.append(task)
    for task in tasks:
        task.join()
    return output_cubes

def sphere_to_cube_process(im, face_id, axA_array, axB_array, size, hsize, format, output_cube):
    # nz
    if FACES[face_id] == 'nz':
        x_array = np.full(size * size, hsize)
        y_array = - axB_array + np.full(size * size, hsize)
        z_array = - axA_array + np.full(size * size, hsize)
    # pz
    elif FACES[face_id] == 'pz':
        x_array = np.full(size * size, -hsize)
        y_array = axB_array + np.full(size * size, -hsize)
        z_array = - axA_array + np.full(size * size, hsize)
    # px
    elif FACES[face_id] == 'px':
        x_array = axB_array + np.full(size * size, -hsize)
        y_array = np.full(size * size, hsize)
        z_array = - axA_array + np.full(size * size, hsize)
    # nx
    elif FACES[face_id] == 'nx':
        x_array = - axB_array + np.full(size * size, hsize)
        y_array = np.full(size * size, -hsize)
        z_array = - axA_array + np.full(size * size, hsize)
    # py
    elif FACES[face_id] == 'py':
        x_array = axB_array + np.full(size * size, -hsize)
        y_array = axA_array + np.full(size * size, -hsize)
        z_array = np.full(size * size, hsize)
    # ny
    elif FACES[face_id] == 'ny':
        x_array = axB_array + np.full(size * size, -hsize)
        y_array = -axA_array + np.full(size * size, hsize)
        z_array = np.full(size * size, -hsize)

    r_array = np.sqrt(x_array * x_array + y_array * y_array + z_array * z_array)
    theta_array = np.arccos(z_array / r_array)
    phi_array = -np.arctan2(y_array, x_array)
    ix_array = np.floor_divide((im.shape[1] - 1) * phi_array, (2 * math.pi))
    iy_array = np.floor_divide((im.shape[0] - 1) * (theta_array), math.pi)
    ix_array = np.where(ix_array >= 0, ix_array, im.shape[1] + ix_array)
    iy_array = np.where(iy_array >= 0, iy_array, im.shape[0] + iy_array)
    index_array = iy_array * im.shape[1] + ix_array
    reshape_array = im.reshape((im.shape[0] * im.shape[1], 3))
    color_side = reshape_array[index_array.astype(int)]
    color_side = color_side.reshape((size, size, 3))

    if face_id == 5:
        color_side = np.rot90(color_side, 1)
    if face_id == 4:
        color_side = np.rot90(color_side, -1)
    imageio.imwrite(output_cube, color_side, format=format)