import torch
from tqdm import tqdm
import numpy as np
from PIL import Image
torch.backends.cuda.preferred_linalg_library("magma")
SH_C0 = 0.28209479177387814
SH_C1_x = 0.4886025119029199
SH_C1_y = 0.4886025119029199
SH_C1_z = 0.4886025119029199
SH_C2_xy = 1.0925484305920792
SH_C2_xz = 1.0925484305920792
SH_C2_yz = 1.0925484305920792
SH_C2_zz = 0.31539156525252005
SH_C2_xx_yy = 0.5462742152960396
SH_C3_yxx_yyy = 0.5900435899266435
SH_C3_xyz = 2.890611442640554
SH_C3_yzz_yxx_yyy = 0.4570457994644658
SH_C3_zzz_zxx_zyy = 0.3731763325901154
SH_C3_xzz_xxx_xyy = 0.4570457994644658
SH_C3_zxx_zyy = 1.445305721320277
SH_C3_xxx_xyy = 0.5900435899266435
def evaluate_sh(f_dc, f_rest, points, c2w):
sh = torch.empty((points.shape[0], 16, 3),
device=points.device, dtype=points.dtype)
sh[:, 0] = f_dc
sh[:, 1:, 0] = f_rest[:, :15] # R
sh[:, 1:, 1] = f_rest[:, 15:30] # G
sh[:, 1:, 2] = f_rest[:, 30:45] # B
view_dir = points - c2w[:3, 3].unsqueeze(0) # [N, 3]
view_dir = view_dir / (view_dir.norm(dim=-1, keepdim=True) + 1e-8)
x, y, z = view_dir[:, 0], view_dir[:, 1], view_dir[:, 2]
xx, yy, zz = x * x, y * y, z * z
xy, xz, yz = x * y, x * z, y * z
Y0 = torch.full_like(x, SH_C0) # [N]
Y1 = - SH_C1_y * y
Y2 = SH_C1_z * z
Y3 = - SH_C1_x * x
Y4 = SH_C2_xy * xy
Y5 = SH_C2_yz * yz
Y6 = SH_C2_zz * (3 * zz - 1)
Y7 = SH_C2_xz * xz
Y8 = SH_C2_xx_yy * (xx - yy)
Y9 = SH_C3_yxx_yyy * y * (3 * xx - yy)
Y10 = SH_C3_xyz * x * y * z
Y11 = SH_C3_yzz_yxx_yyy * y * (4 * zz - xx - yy)
Y12 = SH_C3_zzz_zxx_zyy * z * (2 * zz - 3 * xx - 3 * yy)
Y13 = SH_C3_xzz_xxx_xyy * x * (4 * zz - xx - yy)
Y14 = SH_C3_zxx_zyy * z * (xx - yy)
Y15 = SH_C3_xxx_xyy * x * (xx - 3 * yy)
Y = torch.stack([Y0, Y1, Y2, Y3, Y4, Y5, Y6, Y7, Y8, Y9, Y10, Y11, Y12, Y13, Y14, Y15],
dim=1) # [N, 16]
return torch.sigmoid((sh * Y.unsqueeze(2)).sum(dim=1))
def project_points(pc, c2w, fx, fy, cx, cy):
w2c = torch.eye(4, device=pc.device)
R = c2w[:3, :3]
t = c2w[:3, 3]
w2c[:3, :3] = R.t()
w2c[:3, 3] = -R.t() @ t
PC = ((w2c @ torch.concatenate(
[pc, torch.ones_like(pc[:, :1])], dim=1).t()).t())[:, :3]
x, y, z = PC[:, 0], PC[:, 1], PC[:, 2] # Camera space
uv = torch.stack([fx * x / z + cx, fy * y / z + cy], dim=-1)
return uv, x, y, z
def inv2x2(M, eps=1e-12):
a = M[:, 0, 0]
b = M[:, 0, 1]
c = M[:, 1, 0]
d = M[:, 1, 1]
det = a * d - b * c
safe_det = torch.clamp(det, min=eps)
inv = torch.empty_like(M)
inv[:, 0, 0] = d / safe_det
inv[:, 0, 1] = -b / safe_det
inv[:, 1, 0] = -c / safe_det
inv[:, 1, 1] = a / safe_det
return inv
def build_sigma_from_params(scale_raw, q_raw):
scale = torch.exp(scale_raw).clamp_min(1e-6)
q = q_raw / (q_raw.norm(dim=-1, keepdim=True) + 1e-9)
R = quat_to_rotmat(q)
S = torch.diag_embed(scale)
return R @ S @ S @ R.transpose(1, 2)
def quat_to_rotmat(quat):
x, y, z, w = quat.unbind(dim=-1)
xx, yy, zz = x * x, y * y, z * z
xy, xz, yz = x * y, x * z, y * z
xw, yw, zw = x * w, y * w, z * w
R = torch.stack([
1 - 2 * (yy + zz), 2 * (xy - zw), 2 * (xz + yw),
2 * (xy + zw), 1 - 2 * (xx + zz), 2 * (yz - xw),
2 * (xz - yw), 2 * (yz + xw), 1 - 2 * (xx + yy)
], dim=-1).reshape(quat.shape[:-1] + (3, 3))
return R
def scale_intrinsics(H, W, H_src, W_src, fx, fy, cx, cy):
scale_x = W / W_src
scale_y = H / H_src
fx_scaled = fx * scale_x
fy_scaled = fy * scale_y
cx_scaled = cx * scale_x
cy_scaled = cy * scale_y
return fx_scaled, fy_scaled, cx_scaled, cy_scaled
@torch.no_grad()
def render(pos, color, opacity_raw, sigma, c2w, H, W, fx, fy, cx, cy,
near=2e-3, far=100, pix_guard=64, T=16, min_conis=1e-6,
chi_square_clip=9.21, alpha_max=0.99, alpha_cutoff=1/255.):
uv, x, y, z = project_points(pos, c2w, fx, fy, cx, cy)#相机坐标系下的x,y,z,以及画面里的uv坐标，x,y,z的形状都是(829825,)，uv的形状是(829825,2)
in_guard = (uv[:, 0] > -pix_guard) & (uv[:, 0] < W + pix_guard) & (
uv[:, 1] > -pix_guard) & (uv[:, 1] < H + pix_guard) & (z > near) & (z < far)#这是用来过滤掉视锥体之外的点云的一些条件，根据这些条件最后筛选出了375009个点
uv = uv[in_guard]
pos = pos[in_guard]
color = color[in_guard]
opacity = torch.sigmoid(opacity_raw[in_guard]).clamp(0, 0.999)
z = z[in_guard]
x = x[in_guard]
y = y[in_guard]
sigma = sigma[in_guard]
idx = torch.nonzero(in_guard, as_tuple=False).squeeze(1)
# Project the covariance
Rcw = c2w[:3, :3]
Rwc = Rcw.t()
invz = 1 / z.clamp_min(1e-6)
invz2 = invz * invz
J = torch.zeros((pos.shape[0], 2, 3), device=pos.device, dtype=pos.dtype)
J[:, 0, 0] = fx * invz
J[:, 1, 1] = fy * invz
J[:, 0, 2] = -fx * x * invz2
J[:, 1, 2] = -fy * y * invz2
tmp = Rwc.unsqueeze(0) @ sigma @ Rwc.t().unsqueeze(0) # Eq. 5
sigma_camera = J @ tmp @ J.transpose(1, 2)
sigma_camera = 0.5 * (sigma_camera + sigma_camera.transpose(1, 2)) # Enforce symmetry 椭球投影到相机画面上之后的椭圆的协方差矩阵
# Ensure positive definiteness
evals, evecs = torch.linalg.eigh(sigma_camera)#专门为对称矩阵求特征值和特征向量的函数
evals = torch.clamp(evals, min=1e-6, max=1e4)
sigma_camera = evecs @ torch.diag_embed(evals) @ evecs.transpose(1, 2)
#筛选出协方差矩阵正定的球的uv,color,opacity,z,sigma_camera,idx
keep = torch.isfinite(
sigma_camera.reshape(sigma.shape[0], -1)).all(dim=-1)
uv = uv[keep]
color = color[keep]
opacity = opacity[keep]
z = z[keep]
sigma_camera = sigma_camera[keep]
idx = idx[keep]
# Global depth sorting,按深度值升序排列
order = torch.argsort(z, descending=False)
uv = uv[order]
u = uv[:, 0]
v = uv[:, 1]
color = color[order]
opacity = opacity[order]
sigma_camera = sigma_camera[order]
evals = evals[order]
idx = idx[order]
# Tiling
major_variance = evals[:, 1].clamp_min(1e-12).clamp_max(1e4) # [N]
radius = torch.ceil(3.0 * torch.sqrt(major_variance)).to(torch.int64)#半径=（3*根号最大特征值）向下取整
#根据半径把椭圆变成一个正圆，求出正圆的uv坐标范围
umin = torch.floor(u - radius).to(torch.int64)
umax = torch.floor(u + radius).to(torch.int64)
vmin = torch.floor(v - radius).to(torch.int64)
vmax = torch.floor(v + radius).to(torch.int64)
#筛选出完整地在画面范围内的正圆的u,v,color,opacity,sigma_cameral,umin,umax,vmin,vmax,idx，筛选后剩下351908个
on_screen = (umax >= 0) & (umin < W) & (vmax >= 0) & (vmin < H)
if not on_screen.any():
raise Exception("All projected points are off-screen")
u, v = u[on_screen], v[on_screen]
color = color[on_screen]
opacity = opacity[on_screen]
sigma_camera = sigma_camera[on_screen]
umin, umax = umin[on_screen], umax[on_screen]
vmin, vmax = vmin[on_screen], vmax[on_screen]
idx = idx[on_screen]
umin = umin.clamp(0, W - 1)
umax = umax.clamp(0, W - 1)
vmin = vmin.clamp(0, H - 1)
vmax = vmax.clamp(0, H - 1)
# Tile index for each AABB，T是瓦片的大小，默认为16
#计算正圆的范围是在横竖第几个瓦片上,取整数部分
umin_tile = (umin // T).to(torch.int64) # [N]
umax_tile = (umax // T).to(torch.int64) # [N]
vmin_tile = (vmin // T).to(torch.int64) # [N]
vmax_tile = (vmax // T).to(torch.int64) # [N]
# Number of tiles each gaussian intersects
#每一个正圆的AABB横竖覆盖多少个瓦片
n_u = umax_tile - umin_tile + 1 # [N]
n_v = vmax_tile - vmin_tile + 1 # [N]
# Max number of tiles
#所有的正圆AABB们中横着覆盖的最大的瓦片数量和竖着覆盖的最大瓦片数量
max_u = int(n_u.max().item())
max_v = int(n_v.max().item())
nb_gaussians = umin_tile.shape[0]
span_indices_u = torch.arange(max_u, device=pos.device, dtype=torch.int64) # [max_u] [1,2,3...39] (39,)
span_indices_v = torch.arange(max_v, device=pos.device, dtype=torch.int64) # [max_v] [1,2,3...39] (39,)
#umin_tile的形状为(351908,),umin_tile[:, None, None]其中None可以扩展维度，相当于unsqueeze,umin_tile[:, None, None]的形状为(351908,1,1)
#span_indices_u的形状为(39,),同理，span_indices_u[None, :, None]的形状为(1,39,1)
#umin_tile[:, None, None] + span_indices_u[None, :, None]的时候，\
#umin_tile[:, None, None]会通过广播机制复制，自动扩展为(351908,39,1)
#同理，span_indices_u[None, :, None]会变为(351908,39,1)
#此时形状相同，即可相加
#tile_u说的应该是在每一个AABB的左上角画一个矩形，其尺寸和最大的AABB一致，这个矩形应该能覆盖每一个AABB
#每一个这样的矩形所覆盖的瓦片的横着的索引
#同理tile_v就不用说了
tile_u = (umin_tile[:, None, None] + span_indices_u[None, :, None]
).expand(nb_gaussians, max_u, max_v) # [N, max_u, max_v]
tile_v = (vmin_tile[:, None, None] + span_indices_v[None, None, :]
).expand(nb_gaussians, max_u, max_v) # [N, max_u, max_v]
#mask说的是在351908个覆盖AABB的39x39瓦片大小的大矩形里的所有瓦片中，如果这个瓦片也是AABB包围盒的一部分，就赋值为True,否则为False
mask = (span_indices_u[None, :, None] < n_u[:, None, None]
) & (span_indices_v[None, None, :] < n_v[:, None, None]) # [N, max_u, max_v]
#flat_tile_u就是每一个AABB所覆盖的瓦片的横着的相对于整张图片左上角的索引,一共有3838210个
flat_tile_u = tile_u[mask] # [0, 0, 1, 1, 2, ...]
flat_tile_v = tile_v[mask] # [0, 1, 0, 1, 2]
nb_tiles_per_gaussian = n_u * n_v # [N],每一个AABB的瓦片总和，形状为(351908,)
#print(sum(nb_tiles_per_gaussian)),求和可以得到3838210
gaussian_ids = torch.repeat_interleave(
torch.arange(nb_gaussians, device=pos.device, dtype=torch.int64),
nb_tiles_per_gaussian) # [0, 0, 0, 0, 1 ...] 在每一个AABB中，每一个瓦片对应的高斯的编号
nb_tiles_u = (W + T - 1) // T #整张图片横着可以划分多少个瓦片，98
flat_tile_id = flat_tile_v * nb_tiles_u + flat_tile_u # [0, 0, 0, 0, 1 ...]整张图里所有的AABB的瓦片的每一个瓦片的id\
#这个id是按从左到右从上到下的顺序来的，如下所示
# 0 1 2 3
# 4 5 6 7
# 。。。
#只是会有相同的瓦片属于不同的高斯，但是它们的flat_tile_id相同
idx_z_order = torch.arange(nb_gaussians, device=pos.device, dtype=torch.int64)
M = nb_gaussians + 1
#idx_z_order[gaussian_ids]和gaussian_ids应该是相同的，可以用下面的代码检查一下,耗时9分35秒，一共3838210个
'''
if len(idx_z_order[gaussian_ids]) != len(gaussian_ids):
print(False)
else:
for i in tqdm(range(len(gaussian_ids))):
if gaussian_ids[i] != idx_z_order[gaussian_ids][i]:
print(False)
break
print(True)
'''
comp = flat_tile_id * M + idx_z_order[gaussian_ids]#给每一个瓦片一个特殊编号
#编号是从近到远从左到右从上到下,如下图所示
# /-3-7-11--15---19--/|
# / 2 6 10 14 18 / |
# /_1_5__9__13___17__/ |
# |0 4 8 12 16 | |
# |202428 32 36 | |
# |404448 ... | /
# |__________________|/
#这只是一个示意图，实际算出来的编号比这个要长的多，因为乘以了M
comp_sorted, perm = torch.sort(comp)
gaussian_ids = gaussian_ids[perm]
tile_ids_1d = torch.div(comp_sorted, M, rounding_mode='floor')#还原出已排好序的瓦片的编号
# tile_ids_1d [0, 0, 0, 1, 1, 2, 2, 2, 2]
# nb_gaussian_per_tile [3, 2, 4]
# start [0, 3, 5]
# end [3, 5, 9]
unique_tile_ids, nb_gaussian_per_tile = torch.unique_consecutive(tile_ids_1d, return_counts=True)#消除连续重复元素
start = torch.zeros_like(unique_tile_ids)
#nb_gaussian_per_tile[:-1]意思是去掉nb_gaussian_per_tile的最后一个元素
#cumsum意思是前缀累加,比如a=[a0,a1,a2],则cumsum(a)=[a0,a0+a1,a0+a1+a2]
#这里算出来的是编号为0~某一个id总共有多少个高斯
#到此为止需要注意，一个高斯其实是一个椭球，投影到画面上就是一个椭圆
#它是先用一个正圆来粗略地替代椭圆，然后用正圆的AABB，即一个正方形，来粗略地替代一个正圆
#所以它是用一个正方形来代表一个高斯，很多瓦片可能其实并没有被椭圆覆盖，也没有被正圆覆盖
#但是只要它被正方形覆盖了，我们就说这个瓦片上有一个高斯
start[1:] = torch.cumsum(nb_gaussian_per_tile[:-1], dim=0)
#start和end说的是某一个编号为unique_tile_id的瓦片上有从几到几的高斯，这里说的这个高斯不是单独的高斯，是根据瓦片重复出来的高斯
end = start + nb_gaussian_per_tile
inverse_covariance = inv2x2(sigma_camera)
inverse_covariance[:, 0, 0] = torch.clamp(inverse_covariance[:, 0, 0], min=min_conis)
inverse_covariance[:, 1, 1] = torch.clamp(inverse_covariance[:, 1, 1], min=min_conis)
final_image = torch.zeros((H * W, 3), device=pos.device, dtype=pos.dtype)
# Iterate over tiles
for tile_id, s0, s1 in zip(unique_tile_ids.tolist(), start.tolist(), end.tolist()):
current_gaussian_ids = gaussian_ids[s0:s1]
#取出这个瓦片的左上角坐标x0,y0,和右下角x1,y1
txi = tile_id % nb_tiles_u
tyi = tile_id // nb_tiles_u
x0, y0 = txi * T, tyi * T
x1, y1 = min((txi + 1) * T, W), min((tyi + 1) * T, H)
if x0 >= x1 or y0 >= y1:
continue
#给瓦片里的每个像素标好id,即pixel_idx_1d
#id是先从左到右然后从上到下标的
#可以理解成把整张图先从左到右然后从上到下把所有的像素的id都标出来，然后再截取其中的一个瓦片的id,
#因此pixel_idx_1d是不连续的，每16个跳跃一次
#整张图标id的顺序如下所示
# 123
# 456
# 789
#从第一行开始从左到右，标完一行后开始下一行，再从左到右，再下一行，从左到右，再下一行，从左到右...
xs = torch.arange(x0, x1, device=pos.device, dtype=pos.dtype)
ys = torch.arange(y0, y1, device=pos.device, dtype=pos.dtype)
pu, pv = torch.meshgrid(xs, ys, indexing='xy')
px_u = pu.reshape(-1) # [T * T]
px_v = pv.reshape(-1)
pixel_idx_1d = (px_v * W + px_u).to(torch.int64)
#取出这个瓦片上的所有高斯球投影后的圆心uv坐标和颜色、不透明度、协方差矩阵的逆
gaussian_i_u = u[current_gaussian_ids] # [N]
gaussian_i_v = v[current_gaussian_ids] # [N]
gaussian_i_color = color[current_gaussian_ids] # [N, 3]
gaussian_i_opacity = opacity[current_gaussian_ids] # [N]
gaussian_i_inverse_covariance = inverse_covariance[current_gaussian_ids] # [N, 2, 2]
#q的形状为（292，256），也就是这个瓦片上的292个高斯对这个瓦片上的16x16=256个像素的马氏距离的平方分别是多少
#马氏距离的平方，也就是高斯分布函数的指数部分
du = px_u.unsqueeze(0) - gaussian_i_u.unsqueeze(-1) # [N, T * T]
dv = px_v.unsqueeze(0) - gaussian_i_v.unsqueeze(-1) # [N, T * T]
A11 = gaussian_i_inverse_covariance[:, 0, 0].unsqueeze(-1) # [N, 1]
A12 = gaussian_i_inverse_covariance[:, 0, 1].unsqueeze(-1)
A22 = gaussian_i_inverse_covariance[:, 1, 1].unsqueeze(-1)
q = A11 * du * du + 2 * A12 * du * dv + A22 * dv * dv # [N, T * T]
# chi_square_clip为卡方阈值，马氏距离平方服从 \(\chi^2(2)\) 卡方‑2 分布
# 原始每个瓦片 N 个高斯，T×T 像素，全部算exp(-0.5*q)开销很大。
# 很多像素离高斯椭圆很远，\(\exp(-0.5q)\)几乎等于 0，对像素颜色几乎没有贡献。
#卡方测试，判断像素是否落在 2D 高斯椭圆范围内；椭圆外像素贡献置 0，减少无用计算，模拟原版 3DGS 的高斯包围盒裁剪。
inside = q <= chi_square_clip
# g=e^q,由马氏距离的平方得到这个瓦片上的292个高斯对这个瓦片上的16x16=256个像素的权重分别是多少，g就是这个权重
g = torch.exp(-0.5 * torch.clamp(q, max=chi_square_clip)) # [N, T * T]
g = torch.where(inside, g, torch.zeros_like(g))
#alpha_i：该高斯在此像素处实际起作用的不透明度
#opacity：优化变量，3D 高斯本身的不透明度，每个点一个标量
#alpha_i = opacity * 权重，然后再clamp,不要让alpha_i过大
#alpha_i的形状是(292,256),是每个高斯对每个像素的值
#i是指瓦片内第i个高斯
alpha_i = (gaussian_i_opacity.unsqueeze(-1) * g).clamp_max(alpha_max) # [N, T * T]
#如果该像素处高斯有效不透明度大于等于阈值alpha_cutoff，保留alpha_i；否则直接置为 0，不参与混合
#GPU 每个线程处理一个高斯‑像素对；当 alpha 过小，直接跳过该像素的颜色累加，减少算术与内存写入
#把有效不透明度低于阈值的高斯‑像素对直接清零，抛弃贡献可以忽略不计的项，减少后续 alpha 混合计算量；属于性能优化手段，会引入极轻微视觉误差。
alpha_i = torch.where(alpha_i >= alpha_cutoff, alpha_i, torch.zeros_like(alpha_i))
one_minus_alpha_i = 1 - alpha_i # [N, T * T]
#T_i是透射率
#T0 = (1-alpha0)
#T1 = (1-alpha0)*(1-alpha1)
#T2 = (1-alpha0)*(1-alpha1)(1-alpha2)
#Ti = (1-alpha0)*(1-alpha1)(1-alpha2)...(1-alpha_i)
#T_i的形状是(292,256)
T_i = torch.cumprod(one_minus_alpha_i, dim=0)
#对于某一个像素，把最后一个高斯的透射率去掉，最前面补一个1.0的透射率
T_i = torch.concatenate([
torch.ones((1, alpha_i.shape[-1]), device=pos.device, dtype=pos.dtype),
T_i[:-1]], dim=0)
alive = (T_i > 1e-4).float()
w = alpha_i * T_i * alive # [N, T * T]
#给瓦片里的像素上色
final_image[pixel_idx_1d] = (w.unsqueeze(-1) * gaussian_i_color.unsqueeze(1)).sum(dim=0)
return final_image.reshape((H, W, 3)).clamp(0, 1)
if __name__ == "__main__":
pos = torch.load('trained_gaussians/kitchen/pos_7000.pt').cuda()#(829825,3)
opacity_raw = torch.load('trained_gaussians/kitchen/opacity_raw_7000.pt').cuda()#(829825,)不透明度
f_dc = torch.load('trained_gaussians/kitchen/f_dc_7000.pt').cuda()#(829825,3)0阶球谐基函数系数，三通道
f_rest = torch.load('trained_gaussians/kitchen/f_rest_7000.pt').cuda()#(829825,45)1~15阶球谐基函数系数，三通道*15阶=45
scale_raw = torch.load('trained_gaussians/kitchen/scale_raw_7000.pt').cuda()#(829825,3)椭球在三个方向的缩放系数，由协方差矩阵分解后得到
q_raw = torch.load('trained_gaussians/kitchen/q_rot_7000.pt').cuda()#(829825,4)椭球的旋转量，四元数，由协方差矩阵分解后再转换得到
cam_parameters = np.load('out_colmap/kitchen/cam_meta.npy', allow_pickle=True).item()#(829825,5)相机内参：fx fy width height model:PINHOLE针孔相机
orbit_c2ws = torch.load('camera_trajectories/kitchen_orbit.pt').cuda()#(360,4,4)360个相机姿态，每个姿态用4x4的矩阵表示
sigma = build_sigma_from_params(scale_raw, q_raw)#(829825,3,3)每个椭球的3x3协方差矩阵
with torch.no_grad():
for i, c2w_i in tqdm(enumerate(orbit_c2ws)):
c2w = c2w_i
H = cam_parameters['height'] // 2
W = cam_parameters['width'] // 2
H_src = cam_parameters['height']
W_src = cam_parameters['width']
fx, fy = cam_parameters['fx'], cam_parameters['fy']
cx, cy = W_src / 2, H_src / 2
fx, fy, cx, cy = scale_intrinsics(H, W, H_src, W_src, fx, fy, cx, cy)
color = evaluate_sh(f_dc, f_rest, pos, c2w)
img = render(pos, color, opacity_raw, sigma, c2w, H, W, fx, fy, cx, cy)
Image.fromarray((img.cpu().detach().numpy() * 255).astype(np.uint8)).save(f'novel_views/frame_{i:04d}.png')