【代码详解】nvdiffmodeling渲染部分代码解析

目的

是英伟达开源的一个通过深度学习来优化mesh和材质等信息的repo。数据是待优化的mesh网格.obj文件和材质信息.mtl文件，由于是在仿真数据上做实验，所以他们自然也拥有真值mesh和mtl，在需要对应角度的图片的时候直接即可，然后就可以把优化后的mesh的图片和真值对应角度渲染的图片计算loss。

但是在实际的项目中，我们当然是在一个不太好的mesh和纹理上做优化，本身就没有它们的真值（我们都有真值了还优化个锤子），所以无法从mesh和材质的真值来渲染得到对应视角下应该有的图片。因此，需要做的是把我们的数据转化成之后的图片，也要从我们的mesh中得到它需要的信息。这篇文章先梳理一下这个项目的主要逻辑。

代码内容

、是在固定迭代次数之后对当前的mesh和真值的mesh在特定角度渲染一下保存图片。

、是每次算loss之前渲染待优化mesh和真值mesh得到的图片，我们需要固定对mesh渲染的角度，以和特定角度下拍摄的图片真值计算loss。

仿真实验的时候，代码中是随机生成了一些RT矩阵，然后进行操作：

mvp = np.zeros((FLAGS.batch, 4,4),  dtype=np.float32)
campos   = np.zeros((FLAGS.batch, 3), dtype=np.float32)
lightpos = np.zeros((FLAGS.batch, 3), dtype=np.float32)# ==============================================================================================
#  Build transform stack for minibatching
# ==============================================================================================
for b in range(FLAGS.batch):# Random rotation/translation matrix for optimization.r_rot  = util.random_rotation_translation(0.25)r_mv       = np.matmul(util.translate(0, 0, -RADIUS), r_rot)mvp[b]     = np.matmul(proj_mtx, r_mv).astype(np.float32)campos[b]  = np.linalg.inv(r_mv)[:3, 3]lightpos[b] = util.cosine_sample(campos[b])*RADIUS

r_rot是一个随机创建的RT矩阵，外参描述的是世界坐标系变换到相机坐标系。

mvp[b]是batch中第b个batch里的投影矩阵×外参的结果。

是r_mv的逆矩阵（旋转平移矩阵求逆即相机姿态矩阵，描述了相机坐标系如何变换到世界坐标系下）的第1、2、3行的第4列元素，也就是拿到了-T，描述了相机中心在世界坐标系中的位置。

然后创建一个batch×512×512×3的随机背景颜色。

之后把mesh移到中央：

def center_by_reference(base_mesh, ref_aabb, scale):center = (ref_aabb[0] + ref_aabb[1]) * 0.5scale = scale / torch.max(ref_aabb[1] - ref_aabb[0]).item()v_pos = (base_mesh.v_pos - center[None, ...]) * scalereturn Mesh(v_pos, base=base_mesh)

其中是xyz的 box的中心点，1×3，scale是对应的缩放比例，v_pos是所有缩放之后顶点的坐标，N×3。

然后调用函数对mesh的真值进行渲染得到图片，的shape为[, , , 3]，可视化结果如下图。

with torch.no_grad():color_ref = render.render_mesh(glctx, _opt_ref, mvp, campos, lightpos, FLAGS.light_power, iter_res, spp=iter_spp, num_layers=1, background=randomBgColor, min_roughness=FLAGS.min_roughness)

的定义如下：

def render_mesh(ctx,mesh,mtx_in,view_pos,light_pos,light_power,resolution,spp                       = 1,num_layers                = 1,msaa                      = False,background                = None,antialias                 = True,min_roughness             = 0.08):

mesh就是移动到中央之后的mesh，就是mvp矩阵，也就是 p r o j c a m e r a c l i p T w o r l d c a m e r a T m o d e l w o r l d proj_{}^{clip}T_{world}^{}T_{model}^{world} ​​​，就是也就是T，是，是超参里面设置的，是超参里设置得到的分辨率(即)，spp是超参里设置的（默认为1）。

把这些numpy变量都转化为：

    def prepare_input_vector(x):x = torch.tensor(x, dtype=torch.float32, device='cuda') if not torch.is_tensor(x) else xreturn x[:, None, None, :] if len(x.shape) == 2 else xfull_res = resolution*spp# Convert numpy arrays to torch tensorsmtx_in      = torch.tensor(mtx_in, dtype=torch.float32, device='cuda') if not torch.is_tensor(mtx_in) else mtx_inlight_pos   = prepare_input_vector(light_pos)light_power = prepare_input_vector(light_power)view_pos    = prepare_input_vector(view_pos)

然后把mesh的顶点转换到xyz均属于[-1,1]的裁剪空间，这就是把顶点[, , 3]经过mvp矩阵转换后得到的坐标，所以的shape为[, , 4]，其中真值ref mesh和待优化的base mesh是可能不一样的，所以在输出的时候是ref mesh或base mesh的顶点数。

    # clip space transformv_pos_clip = ru.xfm_points(mesh.v_pos[None, ...], mtx_in)

接下来是从前到后渲染所有layer，默认为1，这里先是使用了里面的()，的文档里说在为1时和()函数一样，返回的rast和db的shape均为[, , , 4]，rast的4维是uvzw，u和v是该像素在三维空间的面片中由三个顶点表示的坐标，z是深度值，w是这个三角面片的id，db存储的是导数。然后，调用了.py这个文件里的进行渲染，后面再解析。总之是一个形如[1, 2, , , , 4]的list。

# Render all layers front-to-backlayers = []with dr.DepthPeeler(ctx, v_pos_clip, mesh.t_pos_idx.int(), [resolution*spp, resolution*spp]) as peeler:for _ in range(num_layers):rast, db = peeler.rasterize_next_layer() layers += [(render_layer(rast, db, mesh, view_pos, light_pos, light_power, resolution, min_roughness, spp, msaa), rast)]

每一层都渲染好之后，要开始混合了。先考虑背景，如果背景是空的话那最简单，直接初始化一个的RGB矩阵，如果有背景的话，那就要求背景必须和一样大，如果spp这个缩放比例大于1的话还要对背景插值。

# Clear to background layerif background is not None:assert background.shape[1] == resolution and background.shape[2] == resolutionif spp > 1:background = util.scale_img_nhwc(background, [full_res, full_res], mag='nearest', min='nearest')accum_col = backgroundelse:accum_col = torch.zeros(size=(1, full_res, full_res, 3), dtype=torch.float32, device='cuda')

接下来要将每一个layer的颜色合成到一起，从远往近，对于每一个color和rast，rast的第4维度是三角面片的id，如果大于0就意味着它应该在这个像素被渲染。

color的最后一项是透明度，把累积颜色和这一层的颜色做线性插值，即 + alpha * (color[…, 0:3] - )。如果需要反走样，再调用反走样函数。

# Composite BACK-TO-FRONTfor color, rast in reversed(layers):alpha     = (rast[..., -1:] > 0) * color[..., 3:4]accum_col = torch.lerp(accum_col, color[..., 0:3], alpha)if antialias:accum_col = dr.antialias(accum_col.contiguous(), rast, v_pos_clip, mesh.t_pos_idx.int()) # TODO: need to support bfloat16

最后，如果spp大于1，用平均池化把图片降采样，否则的话就把返回，可视化如下图，这就是的返回结果，只是这个图因为是刚开始训练的时候产生的，所以形状和纹理都还很差。

# Downscale to framebuffer resolution. Use avg pooling out = util.avg_pool_nhwc(accum_col, spp) if spp > 1 else accum_colreturn out

现在，再回过头来看一下做了什么。

def render_layer(rast,rast_deriv,mesh,view_pos,light_pos,light_power,resolution,min_roughness,spp,msaa):

先是把变为指定的大小。MSAA是多重采样抗锯齿，寻找出物体边缘部分的像素，然后对它们进行缩放处理。

    full_res = resolution*spp################################################################################# Rasterize################################################################################# Scale down to shading resolution when MSAA is enabled, otherwise shade at full resolutionif spp > 1 and msaa:rast_out_s = util.scale_img_nhwc(rast, [resolution, resolution], mag='nearest', min='nearest')rast_out_deriv_s = util.scale_img_nhwc(rast_deriv, [resolution, resolution], mag='nearest', min='nearest') * sppelse:rast_out_s = rastrast_out_deriv_s = rast_deriv

然后基于v_pos即mesh顶点位置、即rast、即顶点的id来做空间中的插值。之后，再计算每个面的法向量并编个id，因此它们的shape都为[, 3]。是对每个面的法向量做插值，实际上得到的就是，图片中每一个像素对应在三维空间里的面片的法向量。类似地，是对顶点的法向量做插值，是对顶点的切向量做插值。它们的shape均为[, , , 3]。可视化结果如下图。

然后，对每个顶点的纹理做插值，的shape为[, , , 2]，的shape是[, , , 4]。

    ################################################################################# Interpolate attributes################################################################################# Interpolate world space positiongb_pos, _ = interpolate(mesh.v_pos[None, ...], rast_out_s, mesh.t_pos_idx.int())# Compute geometric normals. We need those because of bent normals trick (for bump mapping)v0 = mesh.v_pos[mesh.t_pos_idx[:, 0], :]v1 = mesh.v_pos[mesh.t_pos_idx[:, 1], :]v2 = mesh.v_pos[mesh.t_pos_idx[:, 2], :]face_normals = util.safe_normalize(torch.cross(v1 - v0, v2 - v0))face_normal_indices = (torch.arange(0, face_normals.shape[0], dtype=torch.int64, device='cuda')[:, None]).repeat(1, 3)gb_geometric_normal, _ = interpolate(face_normals[None, ...], rast_out_s, face_normal_indices.int())# Compute tangent spaceassert mesh.v_nrm is not None and mesh.v_tng is not Nonegb_normal, _ = interpolate(mesh.v_nrm[None, ...], rast_out_s, mesh.t_nrm_idx.int())gb_tangent, _ = interpolate(mesh.v_tng[None, ...], rast_out_s, mesh.t_tng_idx.int()) # Interpolate tangents# Texure coordinateassert mesh.v_tex is not Nonegb_texc, gb_texc_deriv = interpolate(mesh.v_tex[None, ...], rast_out_s, mesh.t_tex_idx.int(), rast_db=rast_out_deriv_s)

拿到了图片中每个像素对应的mtl纹理图上的坐标后，就可以给它们上色了，color的shape自然也是[, , , 4]。可视化结果如下图。

    ################################################################################# Shade################################################################################color = shade(gb_pos, gb_geometric_normal, gb_normal, gb_tangent, gb_texc, gb_texc_deriv, view_pos, light_pos, light_power, mesh.material, min_roughness)################################################################################# Prepare output################################################################################# Scale back up to visibility resolution if using MSAAif spp > 1 and msaa:color = util.scale_img_nhwc(color, [full_res, full_res], mag='nearest', min='nearest')# Return color & raster output for peelingreturn color

通过的方法，待优化的mesh可以渲染出特定视角下的图片，mesh真值也可以渲染出这个视角下的图片，那接下来就是顺理成章地计算loss和反向传播了，迭代若干次，就可以学习到更好的mesh和漫反射Kd、镜面反射Ks、法向量这三张图了。最后的结果有多好呢？