GS-IR：3D 高斯喷溅用于逆向渲染

一、训练代码

Installation

create the basic environment

conda env create --file environment.yml
conda activate gsir

pip install kornia

install some extensions

cd gs-ir && python setup.py develop && cd ..

cd submodules
git clone https://github.com/NVlabs/nvdiffrast
pip install ./nvdiffrast

pip install ./simple-knn
pip install ./diff-gaussian-rasterization # or cd ./diff-gaussian-rasterization && python setup.py develop && cd ../..

TensoIR-Synthetic.Take the lego case as an example.

Stage1 (Initial Stage)

python train.py \
-m outputs/lego/ \
-s datasets/TensoIR/lego/ \
--iterations 30000 \
--eval

Baking

python baking.py \
-m outputs/lego/ \
--checkpoint outputs/lego/chkpnt30000.pth \
--bound 1.5 \
--occlu_res 128 \
--occlusion 0.25

Stage2 (Decomposition Stage)

python train.py \
-m outputs/lego/ \
-s datasets/TensoIR/lego/ \
--start_checkpoint outputs/lego/chkpnt30000.pth \
--iterations 35000 \
--eval \
--gamma \
--indirect

set --gamma to enable linear_to_sRGB will cause better relighting results but worse novel view synthesis results set --indirect to enable indirect illumination modelling

将 --gamma 设置为启用 linear_to_sRGB 会导致重光照结果更好，但新视角合成结果更差，将 --indirect 设置为启用间接光照建模

Evaluation (Novel View Synthesis)

python render.py \
-m outputs/lego \
-s datasets/TensoIR/lego/ \
--checkpoint outputs/lego/chkpnt35000.pth \
--eval \
--skip_train \
--pbr \
--gamma \
--indirect

Evaluation (Normal)

python normal_eval.py \
--gt_dir datasets/TensoIR/lego/ \
--output_dir outputs/lego/test/ours_None

Evaluation (Albedo)

python render.py \
-m outputs/lego \
-s datasets/TensoIR/lego/ \
--checkpoint outputs/lego/chkpnt35000.pth \
--eval \
--skip_train \
--brdf_eval

Relighting

python relight.py \
-m outputs/lego \
-s datasets/TensoIR/lego/ \
--checkpoint outputs/lego/chkpnt35000.pth \
--hdri datasets/TensoIR/Environment_Maps/high_res_envmaps_2k/bridge.hdr \
--eval \
--gamma

set --gamma to enable linear_to_sRGB will cause better relighting results but worse novel view synthesis results

Relighting Evaluation

python relight_eval.py \
--output_dir outputs/lego/test/ours_None/relight/ \
--gt_dir datasets/TensoIR/lego/

二、Stage1 (Initial Stage)

Stage1 (Initial Stage)
python train.py \
-m outputs/lego/ \
-s datasets/TensoIR/lego/ \
--iterations 30000 \
--eval

渲染结果

• "render": 最终渲染的RGB图像
• "viewspace_points": 屏幕空间点坐标，2D
• "visibility_filter": 可见性过滤器(radii>0)
• "radii": 高斯点在屏幕上的半径
• "opacity_map": 不透明度图
• "depth_map": 深度图
• "normal_map_from_depth": 从深度图梯度计算的法线图
• "normal_from_depth_mask": 从深度图梯度计算的法线图的掩码
• "normal_map": 原始法线图
• "normal_mask": 原始法线图的掩码
• "albedo_map": 反照率图
• "roughness_map": 粗糙度图
• "metallic_map": 金属度图

|---------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------|-------------|
| ​mask_from_depth | ​标识深度法线图中的无效区域​（需替换）； |----------------------------------------------------------| | 修复无效法线 ​： 将全零法线替换为背景法线 normal_bg，避免归一化错误（除零）或光照失真。 | | (normal_map_from_depth == 0.0).all(0, keepdim=True) 检测法线是否全零（RGB通道均为0） | [1, H, W] |
| ​normal_from_depth_mask | ​标识深度法线图中的有效区域​（全非零），哪些像素是有效的 | (normal_map_from_depth != 0).all(0) 检测法线是否全非零（RGB通道均非0） | [H, W] |

submodules/diff-gaussian-rasterization/diff_gaussian_rasterization/init.py中

GaussianRasterizer中forward输出下图，上图右部分是forward中间步骤

其他图计算

类似alpha混合计算颜色

每个高斯点对像素的贡献权重为 weight = alpha * T，所有属性值都通过权重进行加权累积

const float weight = alpha * T;
			// Eq. (3) from 3D Gaussian splatting paper.
			for (int ch = 0; ch < CHANNELS; ch++) {
				C[ch] += features[collected_id[j] * CHANNELS + ch] * weight;
				A[ch] += albedo[collected_id[j] * CHANNELS + ch] * weight;
                //if (NoV > 0.0f) // NOTE: the trick from GIR, do not make scene for scenes
					N[ch] += normals[collected_id[j] * CHANNELS + ch] * weight;
			}
			R += roughness[collected_id[j]] * weight;
			M += metallic[collected_id[j]] * weight;

			// softmax weight
			D += depth[collected_id[j]] * weight;
			O += weight;

			if (weight > max_weight) {
				except_depth = depth[collected_id[j]];
				max_weight = weight;
		}
//---------
if (inside)
	{
		final_T[pix_id] = T;
		n_contrib[pix_id] = last_contributor;
		for (int ch = 0; ch < CHANNELS; ch++) {
			out_color[ch * H * W + pix_id] = C[ch] + T * bg_color[ch];//alpha混合
			out_normal[ch * H * W + pix_id] = N[ch];
			out_albedo[ch * H * W + pix_id] = A[ch];
		}
		if (inference) {//训练时是false
			out_roughness[pix_id] = R + T;
		} else {
			out_roughness[pix_id] = R;
		}
		out_metallic[pix_id] = M;
		if (O > 1e-6) {
			out_depth[pix_id] = argmax_depth ? except_depth : D / O;
			//gaussian_renderer/__init__.py中 argmax_depth 为false
			//except_depth 使用最大权重对应的深度值
		} else {
			out_depth[pix_id] = 0.0f;
		}
		out_opacity[pix_id] = O;
	}

normal_from_depth计算

            normal_from_depth = _C.depth_to_normal(
                raster_settings.image_width,
                raster_settings.image_height,
                focal_x,
                focal_y,
                raster_settings.viewmatrix,
                depth_filter,
            )

在submodules/diff-gaussian-rasterization/ext.cpp中

在submodules/diff-gaussian-rasterization/rasterize_points.cu中

在submodules/diff-gaussian-rasterization/cuda_rasterizer/rasterizer_impl.cu中

在submodules/diff-gaussian-rasterization/cuda_rasterizer/forward.cu中

depthmapToNormalCUDA

1. 初始化与边界处理​

• ​块索引分配 ​：根据线程块索引确定当前线程块处理的像素区域范围（pix_min到pix_max）。

• ​边缘像素跳过 ​：若当前像素位于图像边缘（x=0、x=W-1、y=0、y=H-1），直接返回（无法计算完整邻域梯度）。

• ​深度有效性检查 ​：若当前像素深度值小于阈值（如0.01），视为无效点并跳过。

• ​无效邻域跳过 ​：若​±2像素范围内邻域内任一像素深度无效，放弃当前像素的法线计算。

​2. 邻域深度采样​

• 获取当前像素左、右、上、下 四个方向的邻域深度值（depth_left、depth_right、depth_up、depth_down）。

​3. 3D坐标转换​

• ​相机坐标系转换 ​：将当前像素及四个邻域像素的2D坐标转换为相机坐标系下的3D坐标

​4**. 梯度方向自适应选择**​

• ​X方向梯度​：比较左/右邻域深度差，选择深度变化更小的方向计算水平梯度向量

• ​Y方向梯度​：比较上/下邻域深度差，选择深度变化更小的方向计算垂直梯度向量（公式类同）。

​5. 法线计算​

• ​叉乘法线向量 ​：通过水平梯度（ddx）和垂直梯度（ddy）的叉乘得到相机坐标系下的法线，并归一化，使用视图矩阵（viewmatrix）的旋转部分（前3×3子矩阵）将法线从相机坐标系变换到世界坐标系，将计算的世界坐标系法线按通道（x, y, z）写入输出缓冲区，存储布局为[3, H, W]。

损失

像素级约束： L1 + SSIM确保渲染图像接近真实图像
几何约束 ：法线一致性损失确保几何结构正确，normal_map（直接从高斯点的法线属性渲染得到）和normal_map_from_depth（从深度图计算得到的法线）只在有效区域（mask）内计算l1损失
**平滑性约束：**TV损失（计算相邻像素间的差异（梯度），利用真实图像和法线图的垂直水平方向梯度加权），确保法线图平滑自然

• 平滑区域 （gt_image梯度小 → 权重 ≈ 1）：强烈惩罚预测图像的梯度（强制平滑）。
• 边缘区域 （gt_image梯度大 → 权重 ≈ 0）：弱化惩罚（允许保留边缘）。

三、baking

• 初始化设置：读取命令行参数，加载预训练的3D高斯模型。

• 定义立方体贴图：设置6个方向的视图矩阵（前、后、上、下、左、右），用于生成立方体贴图。

• 创建3D网格：在指定边界内生成一个3D网格，将点云坐标映射到网格索引空间

• 标记有效网格点 ：根据3D高斯点的位置，标记每个高斯点周围的8个相邻网格点为有效。

• 逐网格点计算环境光遮蔽 （AO）：

  • 对于每个有效网格点，从6个方向渲染立方体贴图（颜色、透明度、深度）。

  • 将6个不同视角的深度图合并成一个连续的环境贴图，尺寸变成了 [256, 512, 1]。其中

    # 立方体贴图渲染使用256×256分辨率
    # 环境贴图方向网格使用256×512分辨率

  • 根据深度信息计算遮挡掩码（occlusion mask）。

  • 使用球谐函数（Spherical Harmonics, SH）编码环境贴图的遮蔽信息：将遮挡掩码与立体角权重相乘，得到加权颜色值，利用加权颜色与球谐函数分量的乘积求和，计算遮挡球谐系数。

• 填充未计算网格点 ：对未计算的网格点（occlusion_ids=-1），使用邻近的有效值填充其遮蔽信息。

• 保存结果 ：将遮挡体积数据保存为文件。

  "occlusion_ids": occlusion_ids,#网格点ID，有效的分配连续的索引编号；；；其他无效的还是-1
   "occlusion_coefficients": occlusion_coefficients,#所有网格点的遮挡球谐系数
  "bound": args.bound,#边界
   "degree": occlu_sh_degree,#sh阶数，4
   "occlusion_threshold": occlusion_threshold,#遮挡阈值， #未被遮挡的像素（深度<阈值）

四、 Stage2 (Decomposition Stage)

python train.py \
-m outputs/lego/ \
-s datasets/TensoIR/lego/ \
--start_checkpoint outputs/lego/chkpnt30000.pth \
--iterations 35000 \
--eval \
--gamma \
--indirect

30000次迭代之前与第一次训练类似，

30000次到35000次迭代：高斯迭代，再使用adam优化irradiance_volumes的参数irradiance_coefficients每个网格点的球谐系数；cubemap中的self.base,原始的立方体贴图数据。

核心流程概述​

1.​条件判断​

当启用间接光照（indirect=True）且需要计算遮挡时（occlusion_flag=True）：

加载baking.py预计算的遮挡体积数据 ​（occlusion_volumes.pth），提取遮挡系数、球谐阶数、边界框（aabb）等信息

2.​重建遮挡项与辐照度项​

结合视图方向 view_dirs 和深度 depth_map 计算世界坐标系中的点坐标 points

通过 recon_occlusion() 函数重构每个点的遮挡系数 occlusion，输出形状 [H, W, 1]

调用 irradiance_volumes.query_irradiance() 获取辐照度 irradiance，输出形状 [H, W, 1]

​直接光照默认值​（不启用间接光照）：遮挡设为全1，辐照度设为全0

3.​PBR着色​

环境光 cubemap 构建 Mipmap

调用 pbr_shading() 函数，结合光照、材质、遮挡与辐照度数据，输出包含渲染图像 render_rgb

4.​后处理与损失计算​

无效法线区域（normal_mask=0）用背景色填充

​主损失 ​：pbr_render_loss = L1_loss(render_rgb, gt_image)

​BRDF正则化 ​：根据法线掩码状态计算总变差损失 brdf_tv_loss

计算 gt_image 的水平和垂直梯度权重 rgb_grad_h 和 rgb_grad_w；计算 prediction 的 TV 损失项 tv_h 和 tv_w；构造掩码项 mask_h 和 mask_w；加权求和得到最终的 TV loss。

​材质属性约束 ​：对粗糙度和金属度施加正则项 lamb_loss

​环境贴图平滑约束 ​：计算环境光的TV损失 env_tv_loss

​总损失组合​

total_loss = pbr_render_loss + 
             brdf_tv_loss * weight1 + 
             lamb_loss * weight2 + 
             env_tv_loss * weight3

遮挡体​重建

sparse_interpolate_coefficients

通过稀疏插值获取系数和对应的ID： coefficients, # [HW, d2, 1] ；coeff_ids, # [HW, 8]

1. ​线程初始化与边界检查​

   • 计算线程 ID (ray_id = blockIdx.x * blockDim.x + threadIdx.x)，若超出总光线数 num_rays 则直接返回。

2. ​坐标转换与体素索引计算​

   • 包围盒处理 ：读取场景 AABB 边界 aabb_min 和 aabb_max。
   • 点坐标映射 ：

     • 将世界空间点坐标 point 映射到体素空间：

       n_xyz = (point - aabb_min) / grid_size  // grid_size = (aabb_max - aabb_min) / (occlu_res - 1)

     • 约束坐标范围至 [0, occlu_res-1]，并取整数部分 quat = floor(n_xyz)。

3. ​方向向量生成与法线掩码​

   • 方向向量 ：计算当前点指向周围 8 个相邻体素中心的向量（dir000 至 dir111）。
   • 法线感知掩码 ：
     • 通过 dot(dir, normal) > 0 判断方向是否位于法线正面半球，生成二值掩码（mask000-mask111）。

4. ​三线性权重计算与归一化​

   • 权重计算 ：基于坐标小数部分 o_xyz = n_xyz - quat 计算初始三线性权重：

     weight000 = (1 - o_xyz.x) * (1 - o_xyz.y) * (1 - o_xyz.z) * mask000;
     // 其他7个权重类似

   • 归一化 ：权重总和 weight_sum 归一化，避免零除（EPS 为极小值）。

5. ​稀疏体素索引查询​

   • 索引转换 ：将三维体素坐标 (x, y, z) 转换为一维索引：

     index = x * (occlu_res²) + y * (occlu_res) + z;

   • 边界保护 ：使用 min(quat + 1, occlu_res - 1) 防止索引越界。

   • 输出索引 ：将 8 个相邻体素的 SH 系数索引写入 output_ids。

6. ​球谐系数插值​

   • 系数读取 ：根据索引从 coeffs_ptr 中获取 8 个体素的 SH 系数（每个体素 d2 = sh_degree² 个系数）。
   • 加权求和：对每个 SH 系数通道执行加权融合

SH_reconstruction

|-----------------|----------|---------------------|---------------------------|
| ​输入​ | | | |
| coeffs_ptr | float* | [num_rays, C, d2] | 球谐系数矩阵（d2 = sh_degree²） |
| lobes_ptr | float* | [num_rays, 3] | 法线方向向量（世界坐标系） |
| roughness_ptr | float* | [num_rays, 1] | 表面粗糙度参数 |
| ​输出​ | | | |
| output_recon | float* | [num_rays, C] | 重建后的遮挡（各通道） |

使用低差异序列 ​Hammersley ​ 生成均匀分布的样本点 Xi。

基于粗糙度用 importanceSampleGGX，生成符合GGX分布的微表面法线方向。

对每个采样方向 sample_dir 计算球谐系数累加值，逐通道（channels=1）累加结果取平均。

辐射度

irradiance volume一般就是以一定的大小（比如2m的立方体）来把场景分割开，每一个小立方体中去记录光照信息（一般我们管这个收集光照信息的叫probe）。

irradiance_volumes.query_irradiance

1.​输入输出​

​输入参数​：batch.size是H*W

points：世界坐标系中的三维点坐标，形状为 [batch_size, 3]。

normals：对应点的法线方向向量，形状为 [batch_size, 3]。

​输出目标​：返回每个点的辐照度值（Irradiance），形状为 [batch_size, 1, 3]（单通道或RGB三通道）。

​2.法线方向球谐基函数计算​

components = components_from_spherical_harmonics(degree, directions=normals)  # [bs, d2]

根据法向量 normals 计算球谐基函数值（Spherical Harmonics Basis）。

3.​三线性插值获取球谐系数​

irradiance_coefficients = trilinear_interpolation(
    coefficients, aabb, points, normals, degree
)  # [bs, d2, channel]

根据 points 的位置在体素网格中定位相邻8个体素。结合 normals 计算法线感知的插值权重（避免背面体素影响）。对每个通道的球谐系数进行加权融合。

4.​辐照度重建计算​

irradiance_map = (irradiance_coefficients * components[..., None]).sum(1)  # [bs, channel]

将插值后的球谐系数 irradiance_coefficients 与基函数值 components 逐通道相乘并求和。

​数学本质​：

其中 clm​ 为球谐系数，Ylm​ 为基函数值。

mipmap

cubemap.build_mips() # 构建环境光的mipmap

MIPMAP（多重纹理映射）:是一种纹理优化技术，预先生成一系列逐渐缩小的纹理图像（mipmap层级）。当物体远离摄像机时，使用较小的纹理层级可以提高渲染性能并减少走样（aliasing）。

def build_mips(self, cutoff: float = 0.99) -> None:
    # 1. 构建mipmap金字塔
    self.specular = [self.base]  # 初始化为原始分辨率
    while 当前mip未达最小分辨率:
        通过cubemap_mip下采样生成新mip
    
    # 2. 生成漫反射贴图（使用最小mip）
    self.diffuse = 辐照度计算(self.specular[-1])
    
    # 3. 预滤波镜面反射
    for 每个mip层级:
        根据层级索引计算粗糙度（0.08-0.5线性映射）
        执行GGX预滤波
    单独处理最大粗糙度（1.0）

`build_mips是生成环境贴图的Mipmap链，处理漫反射和不同粗糙度的镜面反射。

self.specular镜面反射：采用while循环动态构建mipmap金字塔，从初始分辨率(如512x512)开始持续下采样，终止条件为达到最小分辨率阈值(LIGHT_MIN_RES=16)，最终形成分辨率从高到低的多级纹理链。前N-1个mip层级映射到[0.08, 0.5]区间，最后一个mip硬编码为1.0粗糙度，为每个mip层级预计算镜面反射。

漫反射部分：使用最低分辨率mip生成辐照度图。

pbr_shading光照计算流程

​输入：

light (CubemapLight): 环境光贴图对象，包含漫反射和镜面反射部分。

normals (torch.Tensor): 表面法线向量，形状为 [H, W, 3]。

view_dirs (torch.Tensor): 视线方向向量，形状为 [H, W, 3]。

albedo (torch.Tensor): 漫反射颜色（基色），形状为 [H, W, 3]。

roughness (torch.Tensor): 粗糙度，形状为 [H, W, 1]。

mask (torch.Tensor): 可见性掩码，形状为 [H, W, 1]。

tone (bool): 是否启用色调映射，默认为 False。

gamma (bool): 是否启用 Gamma 校正，默认为 False。

occlusion (Optional[torch.Tensor]): 环境光遮蔽系数，形状为 [H, W, 1]。

irradiance (Optional[torch.Tensor]): 预计算的辐照度，形状为 [H, W, 1]。

metallic (Optional[torch.Tensor]): 金属度，形状为 [H, W, 1]。

brdf_lut (Optional[torch.Tensor]): BRDF查找表，形状为 [1, 256, 256, 2]。

background (Optional[torch.Tensor]): 背景颜色，形状为 [H, W, 3]。

返回:

Dict: 包含渲染结果的字典，包括：

• render_rgb: 最终渲染图像，形状为 [H, W, 3]。

• diffuse_rgb: 漫反射分量，形状为 [H, W, 3]。

• specular_rgb: 镜面反射分量，形状为 [H, W, 3]。

• diffuse_light: 漫反射光照强度，形状为 [H, W, 3]。

​1. 漫反射分量（Diffuse）​​

• 光照采样 ：
  用法线方向 normals 从预计算的辐照度贴图 light.diffuse 中采样漫反射光照：

  diffuse_light = dr.texture(light.diffuse, normals)  # [1, H, W, 3]

• 环境光遮蔽 ：
  若提供 occlusion，混合预计算辐照度 irradiance 增强阴影细节：

  diffuse_light = diffuse_light * occlusion + (1 - occlusion) * irradiance

• 最终漫反射颜色 ：
  结合材质底色 albedo：

  diffuse_rgb = diffuse_light * albedo  # 能量守恒：漫反射光 × 基色[6,9](@ref)

​2. 镜面反射分量（Specular）​​

• 反射方向计算 ：
  根据视角方向 view_dirs 和法线 normals 计算反射方向 ref_dirs：

  ref_dirs = 2 * (normals·view_dirs) * normals - view_dirs

• 镜面光照采样 ：
  根据粗糙度 roughness 选择环境贴图的Mip层级，从 light.specular 采样：

  miplevel = light.get_mip(roughness)  # 粗糙度越高→Mip层级越高（模糊反射）[9](@ref)
  spec = dr.texture(light.specular, ref_dirs, mip_level_bias=miplevel)

• BRDF查找与合成 ：

  1. 用法线视角夹角 NoV 和粗糙度查询BRDF查找表 brdf_lut：

     fg_lookup = dr.texture(brdf_lut, [NoV, roughness])  # 输出 [F0比例, 环境光比例]

  2. 计算基础反射率 F0（金属与非金属差异）：

     F0 = metallic * albedo + (1 - metallic) * 0.04  # 金属：albedo；非金属：0.04[6,9](@ref)

  3. 合成镜面反射：

     reflectance = F0 * fg_lookup[..., 0] + fg_lookup[..., 1]
     specular_rgb = spec * reflectance  # 镜面光 × BRDF响应

3.后处理

1. 颜色合成 ：

   render_rgb = diffuse_rgb + specular_rgb  # 漫反射 + 镜面反射

2. 色调映射 （tone=True）：使用 aces_film() 压缩高动态范围（HDR→LDR），避免过曝。

3. Gamma校正 （gamma=True）：转换线性空间→sRGB空间，符合显示器特性。

4. 遮罩处理 ：不可见区域（mask=False）替换为背景色 background。

五、relight.py

python relight.py \ -m outputs/lego

\ -s datasets/TensoIR/lego/

\ --checkpoint outputs/lego/chkpnt35000.pth

\ --hdri datasets/TensoIR/Environment_Maps/high_res_envmaps_2k/bridge.hdr

\ --eval

\ --gamma

检查点是stage2步骤的输出。

流程

1.加载hdri文件

2.latlong_to_cubemap把一个经纬度格式的环境贴图（hdri）转换为立方体贴图（Cubemap）

1. 遍历 6 个立方体面:s=0~5，分别代表 cubemap 的 6 个面（+X, -X, +Y, -Y, +Z, -Z）。
2. 生成每个像素在当前面内的归一化二维坐标网格，范围约为 [-1, 1]。
3. 计算当前面每个像素的三维方向向量
4. 计算当前面每个像素对应的经纬度贴图（latlong）坐标：把三维方向投影到经纬度贴图上的采样坐标。
5. 用 nvdiffrast 的纹理采样函数从经纬度贴图采样颜色

3.render_set针对训练集/测试集的每个视角，调用渲染函数，输出 relight 后的图片

1. 环境光球导入，保存一张环境贴图
2. 遍历所有视角，渲染出深度图、法线图、Albedo、金属、粗糙度等
3. 用 PBR Shading（物理基础渲染）合成最终 RGB 图像
4. 应用alpha掩码，将背景部分置为0