天天看點

HDR全景圖切割算法(numpy版本)

原理參考:https://stackoverflow.com/questions/29678510/convert-21-equirectangular-panorama-to-cube-map

這裡修改成numpy和多線程優化,從用blender的sphere2cube(它本身也是周遊像素操作,而且無法保留亮度資訊)幾分鐘到10s内。

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)