基于GPU实现立方体贴图到等矩形全景图的转换

需求背景 #

现有6张宽高比1:1、分辨率1000×1000、视场角（FOV）为90°的立方体贴图，需转换为宽高比2:1、分辨率4000×2000的等矩形全景图。当前CPU实现逻辑如下：

import numpy as np

def uv_pano_to_uv_cubes(u_pano, v_pano, cubemap_shape, panorama_shape):
    # 从2D图像坐标计算球坐标
    phi = (u_pano / panorama_shape[0] - 0.5) * np.pi
    theta = (v_pano / panorama_shape[1] + 0.25) * 2 * np.pi

    # 从球坐标计算笛卡尔坐标
    x = np.cos(phi) * np.cos(theta)
    y = np.sin(phi)
    z = np.cos(phi) * np.sin(theta)

    # 判断像素所属的立方体贴图面
    if abs(x) >= abs(y) and abs(x) >= abs(z):
        face = 'E' if x > 0 else 'W'
        a = z if x > 0 else -z
        b = y
        ma = abs(x)
    elif abs(y) >= abs(x) and abs(y) >= abs(z):
        face = 'D' if y > 0 else 'U'
        a = x
        b = z if y > 0 else -z
        ma = abs(y)
    else:
        face = 'S' if z > 0 else 'N'
        a = -x if z > 0 else x
        b = y
        ma = abs(z)

    # 计算对应立方体贴图上的像素坐标
    u_cube = int(((a / ma) + 1.0) * 0.5 * (cubemap_shape[0] - 1))
    v_cube = int(((b / ma) + 1.0) * 0.5 * (cubemap_shape[1] - 1))

    return face, u_cube, v_cube

该CPU实现转换耗时数秒，CPU多线程优化效果不佳，OpenGL着色器实现未成功，现提供以下Python环境下基于GPU的可行实现方案：

方案1：使用PyTorch（CUDA加速） #

适合已有深度学习环境的用户，利用PyTorch的CUDA张量批量并行计算，支持双线性插值采样，画质更平滑。

实现步骤 #

1. 将6张立方体贴图加载为PyTorch CUDA张量，按面存储为字典；
2. 生成全景图所有像素的坐标网格，转为CUDA张量；
3. 批量计算所有像素的球坐标、笛卡尔坐标，判断所属立方体面并计算对应UV；
4. 利用grid_sample函数完成像素采样，拼接得到全景图。

import torch
import numpy as np

def cubemap_to_panorama(cubemaps, panorama_size=(4000, 2000)):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    H_pano, W_pano = panorama_size
    cubemap_size = 1000

    # 生成全景图像素坐标网格
    u_pano = torch.linspace(0, W_pano-1, W_pano, device=device)
    v_pano = torch.linspace(0, H_pano-1, H_pano, device=device)
    u_grid, v_grid = torch.meshgrid(u_pano, v_pano, indexing='xy')

    # 计算球坐标
    phi = (u_grid / W_pano - 0.5) * np.pi
    theta = (v_grid / H_pano + 0.25) * 2 * np.pi

    # 计算笛卡尔坐标
    x = torch.cos(phi) * torch.cos(theta)
    y = torch.sin(phi)
    z = torch.cos(phi) * torch.sin(theta)

    # 计算各轴绝对值，判断像素所属立方体面
    abs_x = torch.abs(x)
    abs_y = torch.abs(y)
    abs_z = torch.abs(z)

    # 初始化UV张量和面索引
    u_cube = torch.zeros_like(u_grid)
    v_cube = torch.zeros_like(v_grid)
    face_indices = torch.zeros_like(u_grid, dtype=torch.int32)

    # 处理X轴最大的E/W面
    mask_x = (abs_x >= abs_y) & (abs_x >= abs_z)
    mask_e = mask_x & (x > 0)
    u_cube[mask_e] = ((z[mask_e]/x[mask_e]) + 1.0) * 0.5 * (cubemap_size-1)
    v_cube[mask_e] = ((y[mask_e]/x[mask_e]) + 1.0) * 0.5 * (cubemap_size-1)
    face_indices[mask_e] = 0

    mask_w = mask_x & (x <= 0)
    u_cube[mask_w] = ((-z[mask_w]/x[mask_w]) + 1.0) * 0.5 * (cubemap_size-1)
    v_cube[mask_w] = ((y[mask_w]/x[mask_w]) + 1.0) * 0.5 * (cubemap_size-1)
    face_indices[mask_w] = 1

    # 处理Y轴最大的D/U面
    mask_y = (abs_y >= abs_x) & (abs_y >= abs_z)
    mask_d = mask_y & (y > 0)
    u_cube[mask_d] = ((x[mask_d]/y[mask_d]) + 1.0) * 0.5 * (cubemap_size-1)
    v_cube[mask_d] = ((z[mask_d]/y[mask_d]) + 1.0) * 0.5 * (cubemap_size-1)
    face_indices[mask_d] = 2

    mask_u = mask_y & (y <= 0)
    u_cube[mask_u] = ((x[mask_u]/y[mask_u]) + 1.0) * 0.5 * (cubemap_size-1)
    v_cube[mask_u] = ((-z[mask_u]/y[mask_u]) + 1.0) * 0.5 * (cubemap_size-1)
    face_indices[mask_u] = 3

    # 处理Z轴最大的S/N面
    mask_z = (abs_z >= abs_x) & (abs_z >= abs_y)
    mask_s = mask_z & (z > 0)
    u_cube[mask_s] = ((-x[mask_s]/z[mask_s]) + 1.0) * 0.5 * (cubemap_size-1)
    v_cube[mask_s] = ((y[mask_s]/z[mask_s]) + 1.0) * 0.5 * (cubemap_size-1)
    face_indices[mask_s] = 4

    mask_n = mask_z & (z <= 0)
    u_cube[mask_n] = ((x[mask_n]/z[mask_n]) + 1.0) * 0.5 * (cubemap_size-1)
    v_cube[mask_n] = ((y[mask_n]/z[mask_n]) + 1.0) * 0.5 * (cubemap_size-1)
    face_indices[mask_n] = 5

    # 转换为grid_sample所需的[-1,1]归一化坐标
    u_cube_norm = (u_cube/(cubemap_size-1))*2 -1
    v_cube_norm = (v_cube/(cubemap_size-1))*2 -1

    # 初始化全景图张量并采样各面像素
    panorama = torch.zeros((3, H_pano, W_pano), device=device)
    face_keys = ['E', 'W', 'D', 'U', 'S', 'N']
    for idx, face in enumerate(face_keys):
        mask = face_indices == idx
        if mask.any():
            grid = torch.stack([u_cube_norm[mask], v_cube_norm[mask]], dim=-1).view(1, mask.sum(), 1, 2)
            sampled = torch.nn.functional.grid_sample(
                cubemaps[face].unsqueeze(0), grid, mode='bilinear', padding_mode='border', align_corners=True
            )
            panorama[:, mask] = sampled.squeeze().view(3, -1)

    return panorama.permute(1,2,0).cpu().numpy()

# 使用示例
# cubemaps = {
#     'E': torch.from_numpy(im_e).permute(2,0,1).cuda(),
#     'W': torch.from_numpy(im_w).permute(2,0,1).cuda(),
#     'D': torch.from_numpy(im_d).permute(2,0,1).cuda(),
#     'U': torch.from_numpy(im_u).permute(2,0,1).cuda(),
#     'S': torch.from_numpy(im_s).permute(2,0,1).cuda(),
#     'N': torch.from_numpy(im_n).permute(2,0,1).cuda()
# }
# panorama = cubemap_to_panorama(cubemaps)

方案2：使用CuPy（CUDA加速的NumPy替代） #

完全兼容NumPy语法，无需学习新框架，适合熟悉NumPy的用户，代码逻辑与原CPU实现高度一致。

实现步骤 #

1. 安装CuPy并配置CUDA环境；
2. 将立方体贴图加载为CuPy数组；
3. 生成全景图坐标网格，批量计算所有像素的对应立方体贴图坐标；
4. 通过布尔掩码批量赋值，得到全景图。

import cupy as cp
import numpy as np

def cubemap_to_panorama_cupy(cubemaps, panorama_size=(4000, 2000)):
    H_pano, W_pano = panorama_size
    cubemap_size = 1000

    # 生成全景图坐标网格
    u_pano = cp.linspace(0, W_pano-1, W_pano)
    v_pano = cp.linspace(0, H_pano-1, H_pano)
    u_grid, v_grid = cp.meshgrid(u_pano, v_pano, indexing='xy')

    # 计算球坐标
    phi = (u_grid / W_pano - 0.5) * np.pi
    theta = (v_grid / H_pano + 0.25) * 2 * np.pi

    # 计算笛卡尔坐标
    x = cp.cos(phi) * cp.cos(theta)
    y = cp.sin(phi)
    z = cp.cos(phi) * cp.sin(theta)

    # 初始化全景图数组
    panorama = cp.zeros((H_pano, W_pano, 3), dtype=cp.uint8)

    # 处理X轴最大的E/W面
    mask_x = (cp.abs(x) >= cp.abs(y)) & (cp.abs(x) >= cp.abs(z))
    mask_e = mask_x & (x > 0)
    u_cube = cp.int32(((z[mask_e]/x[mask_e]) +1.0)*0.5*(cubemap_size-1))
    v_cube = cp.int32(((y[mask_e]/x[mask_e]) +1.0)*0.5*(cubemap_size-1))
    panorama[mask_e] = cubemaps['E'][v_cube, u_cube]

    mask_w = mask_x & (x <=0)
    u_cube = cp.int32(((-z[mask_w]/x[mask_w]) +1.0)*0.5*(cubemap_size-1))
    v_cube = cp.int32(((y[mask_w]/x[mask_w]) +1.0)*0.5*(cubemap_size-1))
    panorama[mask_w] = cubemaps['W'][v_cube, u_cube]

    # 处理Y轴最大的D/U面
    mask_y = (cp.abs(y) >= cp.abs(x)) & (cp.abs(y) >= cp.abs(z))
    mask_d = mask_y & (y >0)
    u_cube = cp.int32(((x[mask_d]/y[mask_d]) +1.0)*0.5*(cubemap_size-1))
    v_cube = cp.int32(((z[mask_d]/y[mask_d]) +1.0)*0.5*(cubemap_size-1))
    panorama[mask_d] = cubemaps['D'][v_cube, u_cube]

    mask_u = mask_y & (y <=0)
    u_cube = cp.int32(((x[mask_u]/y[mask_u]) +1.0)*0.5*(cubemap_size-1))
    v_cube = cp.int32(((-z[mask_u]/y[mask_u]) +1.0)*0.5*(cubemap_size-1))
    panorama[mask_u] = cubemaps['U'][v_cube, u_cube]

    # 处理Z轴最大的S/N面
    mask_z = (cp.abs(z) >= cp.abs(x)) & (cp.abs(z) >= cp.abs(y))
    mask_s = mask_z & (z >0)
    u_cube = cp.int32(((-x[mask_s]/z[mask_s]) +1.0)*0.5*(cubemap_size-1))
    v_cube = cp.int32(((y[mask_s]/z[mask_s]) +1.0)*0.5*(cubemap_size-1))
    panorama[mask_s] = cubemaps['S'][v_cube, u_cube]

    mask_n = mask_z & (z <=0)
    u_cube = cp.int32(((x[mask_n]/z[mask_n]) +1.0)*0.5*(cubemap_size-1))
    v_cube = cp.int32(((y[mask_n]/z[mask_n]) +1.0)*0.5*(cubemap_size-1))
    panorama[mask_n] = cubemaps['N'][v_cube, u_cube]

    return cp.asnumpy(panorama)

# 使用示例
# cubemaps = {
#     'E': cp.array(im_e),
#     'W': cp.array(im_w),
#     'D': cp.array(im_d),
#     'U': cp.array(im_u),
#     'S': cp.array(im_s),
#     'N': cp.array(im_n)
# }
# panorama = cubemap_to_panorama_cupy(cubemaps)

方案3：使用Numba CUDA #

允许直接用Python编写CUDA核函数，适合需要底层控制计算逻辑的场景，无需依赖大型框架。

实现步骤 #

1. 定义CUDA核函数，处理单个像素的坐标转换与采样；
2. 将立方体贴图转为Numba设备数组；
3. 配置线程块与网格大小，启动核函数并行处理所有像素。

from numba import cuda
import numpy as np

@cuda.jit
def cubemap_to_panorama_kernel(panorama, cubemap_E, cubemap_W, cubemap_D, cubemap_U, cubemap_S, cubemap_N):
    H_pano, W_pano = panorama.shape[:2]
    cubemap_size = cubemap_E.shape[0]
    u_pano, v_pano = cuda.grid(2)
    if u_pano >= W_pano or v_pano >= H_pano:
        return

    # 计算球坐标
    phi = (u_pano / W_pano - 0.5) * np.pi
    theta = (v_pano / H_pano + 0.25) * 2 * np.pi

    # 计算笛卡尔坐标
    x = np.cos(phi) * np.cos(theta)
    y = np.sin(phi)
    z = np.cos(phi) * np.sin(theta)

    abs_x = abs(x)
    abs_y = abs(y)
    abs_z = abs(z)

    # 判断所属面并采样
    if abs_x >= abs_y and abs_x >= abs_z:
        if x > 0:
            u_cube = int(((z/x)+1.0)*0.5*(cubemap_size-1))
            v_cube = int(((y/x)+1.0)*0.5*(cubemap_size-1))
            panorama[v_pano, u_pano] = cubemap_E[v_cube, u_cube]
        else:
            u_cube = int(((-z/x)+1.0)*0.5*(cubemap_size-1))
            v_cube = int(((y/x)+1.0)*0.5*(cubemap_size-1))
            panorama[v_pano, u_pano] = cubemap_W[v_cube, u_cube]
    elif abs_y >= abs_x and abs_y >= abs_z:
        if y > 0:
            u_cube = int(((x/y)+1.0)*0.5*(cubemap_size-1))
            v_cube = int(((z/y)+1.0)*0.5