目的
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矩阵,也就是 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_{camera}^{clip}T_{world}^{camera}T_{model}^{world} projcameraclipTworldcameraTmodelworld,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这三张图了。最后的结果有多好呢?