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

目的

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

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

代码内容

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

color_opt、color_ref是每次算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里的投影矩阵×外参的结果。
campos是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.5
    scale = scale / torch.max(ref_aabb[1] - ref_aabb[0]).item()
    v_pos = (base_mesh.v_pos - center[None, ...]) * scale
    return Mesh(v_pos, base=base_mesh)

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

然后调用render_mesh函数对mesh的真值进行渲染得到图片，color_ref的shape为[minibatch, full_res, full_res, 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)

render_mesh的定义如下：

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，mtx_in就是mvp矩阵，也就是$pro{j}_{camera}^{clip}{T}_{world}^{camera}{T}_{model}^{world}$，view_pos就是campos也就是T，light_pos是lightpos，lightpower是超参里面设置的，resolution是超参里设置得到的分辨率iter_res(即train_res)，spp是超参里设置的iter_spp（默认为1）。
把这些numpy变量都转化为tensor：

    def prepare_input_vector(x):
        x = torch.tensor(x, dtype=torch.float32, device='cuda') if not torch.is_tensor(x) else x
        return x[:, None, None, :] if len(x.shape) == 2 else x

    full_res = resolution*spp

    # Convert numpy arrays to torch tensors
    mtx_in      = torch.tensor(mtx_in, dtype=torch.float32, device='cuda') if not torch.is_tensor(mtx_in) else mtx_in
    light_pos   = prepare_input_vector(light_pos)
    light_power = prepare_input_vector(light_power)
    view_pos    = prepare_input_vector(view_pos)

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

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

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

# Render all layers front-to-back
    layers = []
    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)]

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

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

接下来要将每一个layer的颜色合成到一起，从远往近，对于每一个color和rast，rast的第4维度是三角面片的id，如果大于0就意味着它应该在这个像素被渲染。
color的最后一项是透明度，把累积颜色和这一层的颜色做线性插值，即accum_col + alpha * (color[…, 0:3] - accum_col)。如果需要反走样，再调用反走样函数。

# Composite BACK-TO-FRONT
    for 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，用平均池化把图片降采样，否则的话就把accum_col返回，可视化如下图，这就是render_mesh的返回结果，只是这个图因为是刚开始训练的时候产生的，所以形状和纹理都还很差。

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

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

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

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

    full_res = resolution*spp

    ################################################################################
    # Rasterize
    ################################################################################

    # Scale down to shading resolution when MSAA is enabled, otherwise shade at full resolution
    if 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') * spp
    else:
        rast_out_s = rast
        rast_out_deriv_s = rast_deriv

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

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

    ################################################################################
    # Interpolate attributes
    ################################################################################

    # Interpolate world space position
    gb_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 space
    assert mesh.v_nrm is not None and mesh.v_tng is not None
    gb_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 coordinate
    assert mesh.v_tex is not None
    gb_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自然也是[minibatch, full_res, full_res, 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 MSAA
    if spp > 1 and msaa:
        color = util.scale_img_nhwc(color, [full_res, full_res], mag='nearest', min='nearest')

    # Return color & raster output for peeling
    return color

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