[3D Reconstruction] [Deep Learning] NeuS Code Pytorch Implementation – Test Phase Code Analysis (Part 1)
The paper proposes a novel neural surface reconstruction method, called NeuS, for reconstructing objects and scenes with high fidelity from 2D image input. In NeuS it is proposed to represent surfaces as zero-order sets of signed distance functions (SDF) and to develop a new volume rendering method to train neural SDF representations, thus enabling more accurate surface reconstruction even without mask supervision. . NeuS outperforms existing techniques in high-quality surface reconstruction, especially for objects and scenes with complex structures and self-occlusion. This blog post will analyze the specific functional module code in the testing phase based on the code execution process.
Article directory
Preface
Before analyzing the NeuS network in detail, the first task is to build the operating environment required by NeuS [ win10 reference tutorial ] and complete the training and testing of the model. Only then will it make sense to carry out subsequent work.
This blog post will analyze the functional code modules involved in the NeuS testing phase.
The blogger has analyzed the codes of each functional module in detail in different blog posts. Click [Reference Tutorial under Win10], and the directory link of the blog post is placed in the preface.
The code snippet here is the train function part of the exp_runner.py file. It is part of the training phase in a broad sense. However, since it does not participate in the update of the NeuS network and only performs phased verification of the NeuS network, the blogger put it here Detailed explanation in the blog post.
if self.iter_step % self.save_freq == 0:
self.save_checkpoint()
if self.iter_step % self.val_freq == 0:
self.validate_image()
if self.iter_step % self.val_mesh_freq == 0:
self.validate_mesh()
self.update_learning_rate()
if self.iter_step % len(image_perm) == 0:
image_perm = self.get_image_perm()
save_checkpoint
Member method in the Runner class belonging to the exp_runner.py file, the purpose is to save the NeuS weights of the completed phase training.
def save_checkpoint(self):
checkpoint = {
'nerf': self.nerf_outside.state_dict(), # 各深度学习网络参数权重
'sdf_network_fine': self.sdf_network.state_dict(),
'variance_network_fine': self.deviation_network.state_dict(),
'color_network_fine': self.color_network.state_dict(),
'optimizer': self.optimizer.state_dict(), # 优化器
'iter_step': self.iter_step, # 训练的次数
}
# 创建放置权重模型的文件夹
os.makedirs(os.path.join(self.base_exp_dir, 'checkpoints'), exist_ok=True)
# 保存
torch.save(checkpoint, os.path.join(self.base_exp_dir, 'checkpoints', 'ckpt_{:0>6d}.pth'.format(self.iter_step)))
validate_image
After completing the NeuS model training in stages, you need to render the pictures and compare them with the real training pictures to verify the effect of the model training.
First, the gen_rays_at function is needed to generate the rays of the entire picture (after downsampling), then obtain the farthest point and the closest point of the sampling point (foreground) on the rays , and finally obtain the required results through the renderer function.
def validate_image(self, idx=-1, resolution_level=-1):
# 假设验证图像的序号小于0,随机获取一个图片序号
if idx < 0:
idx = np.random.randint(self.dataset.n_images)
print('Validate: iter: {}, camera: {}'.format(self.iter_step, idx))
if resolution_level < 0:
# 下采样倍数
resolution_level = self.validate_resolution_level
# [W, H, 3]
rays_o, rays_d = self.dataset.gen_rays_at(idx, resolution_level=resolution_level)
H, W, _ = rays_o.shape
# 按照batch_size切分,[W*H,3]=>tuple形式:W*H/batch_size个[batch_size, 3]
rays_o = rays_o.reshape(-1, 3).split(self.batch_size)
rays_d = rays_d.reshape(-1, 3).split(self.batch_size)
out_rgb_fine = []
out_normal_fine = []
for rays_o_batch, rays_d_batch in zip(rays_o, rays_d):
# 最近点和最远点
near, far = self.dataset.near_far_from_sphere(rays_o_batch, rays_d_batch)
# 背景颜色
background_rgb = torch.ones([1, 3]) if self.use_white_bkgd else None
render_out = self.renderer.render(rays_o_batch,
rays_d_batch,
near,
far,
cos_anneal_ratio=self.get_cos_anneal_ratio(),
background_rgb=background_rgb)
def feasible(key): return (key in render_out) and (render_out[key] is not None)
# 前景颜色
if feasible('color_fine'):
out_rgb_fine.append(render_out['color_fine'].detach().cpu().numpy())
# 梯度信息和采样点权重
if feasible('gradients') and feasible('weights'):
n_samples = self.renderer.n_samples + self.renderer.n_importance
# 梯度信息权重加成
normals = render_out['gradients'] * render_out['weights'][:, :n_samples, None] # [batch_size,n_samples，3]
# 采样点是否在球体内
if feasible('inside_sphere'):
# 只保留采样点在球体内的部分
normals = normals * render_out['inside_sphere'][..., None] # [batch_size,n_samples，3]
# normals是带有权重的有效梯度信息
normals = normals.sum(dim=1).detach().cpu().numpy() # [batch_size，3]
out_normal_fine.append(normals)
del render_out
gen_rays_at
The functions defined by the Dataset data manager are in the models/dataset.py file. The blogger [ NeuS Overview ] has briefly introduced this process in his blog post.
def gen_rays_at(self, img_idx, resolution_level=1):
"""
Generate rays at world space from one camera.
一个摄影机在世界空间中生成光线
"""
# 下采样倍数
l = resolution_level
# 获取2D图像上所有的像素点(下采样后的)
tx = torch.linspace(0, self.W - 1, self.W // l)
ty = torch.linspace(0, self.H - 1, self.H // l)
# 生成网格用于生成坐标
pixels_x, pixels_y = torch.meshgrid(tx, ty) # [W, H]
# 相机坐标系下的方向向量:内参(逆)×像素坐标系
p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1) # [W, H, 3]
p = torch.matmul(self.intrinsics_all_inv[img_idx, None, None, :3, :3], p[:, :, :, None]).squeeze() # [W, H, 3]
# 单位方向向量:对方向向量做归一化处理
rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True) # [W, H, 3]
# 世界坐标系下的方向向量：外参(逆)×相机坐标系
rays_v = torch.matmul(self.pose_all[img_idx, None, None, :3, :3], rays_v[:, :, :, None]).squeeze() # [W, H, 3]
# 世界坐标系下的光心位置(外参的逆对应的平移矩阵t)
rays_o = self.pose_all[img_idx, None, None, :3, 3].expand(rays_v.shape) # [W, H, 3]
return rays_o.transpose(0, 1), rays_v.transpose(0, 1) # [H, W, 3]
The execution diagram of the code is shown in the figure below. The function returns rays_o (optical center) and rays_v (unit direction vector).
Pay attention to the difference between the training process and the verification process to generate rays. In the training process, batch_size pixels are randomly selected to generate rays that pass through these pixels, while in the verification process, all pixels of the entire image are selected to generate rays that pass through them. Rays of pixels throughout the image.
validate_mesh
After completing the NeuS model training in stages, it is also necessary to reconstruct the physical model in three dimensions to verify the effect of the model training.
First, the spatial range of reconstruction needs to be demarcated, then the vertex coordinates and face indexes are obtained through the drawing algorithm, and finally the actual three-dimensional model file is output.
def validate_mesh(self, world_space=False, resolution=64, threshold=0.0):
# 获取提取域(方体)的对角线顶点
bound_min = torch.tensor(self.dataset.object_bbox_min, dtype=torch.float32)
bound_max = torch.tensor(self.dataset.object_bbox_max, dtype=torch.float32)
# 面绘制算法获取vertices顶点坐标和triangles面索引
vertices, triangles =\
self.renderer.extract_geometry(bound_min, bound_max, resolution=resolution, threshold=threshold)
os.makedirs(os.path.join(self.base_exp_dir, 'meshes'), exist_ok=True)
if world_space:
# 再次缩放位移
vertices = vertices * self.dataset.scale_mats_np[0][0, 0] + self.dataset.scale_mats_np[0][:3, 3][None]
# 表示和操作三角网格模型
mesh = trimesh.Trimesh(vertices, triangles)
# 保存mesh模型
mesh.export(os.path.join(self.base_exp_dir, 'meshes', '{:0>8d}.ply'.format(self.iter_step)))
logging.info('End')
The figure below shows that bound_min and bound_max delineate the three-dimensional reconstruction range.
As a reminder, the range of 3D reconstruction and the range of rendering into 2D images are different. They each have their own settings, so don’t get confused.
extract_geometry
They are all under the models/renderer.py file. The author of the source code here has made a matryoshka doll. The former extract_geometry is a class member method belonging to the NeuSRenderer class, and the latter is an independent function.
def extract_geometry(self, bound_min, bound_max, resolution, threshold=0.0):
return extract_geometry(bound_min,
bound_max,
resolution=resolution,
threshold=threshold,
query_func=lambda pts: -self.sdf_network.sdf(pts))
Marching_cubes surface rendering algorithm reference , extract_fields is to obtain the sdf value of each point in the three-dimensional reconstruction range.
def extract_geometry(bound_min, bound_max, resolution, threshold, query_func):
print('threshold: {}'.format(threshold))
# 获取提取域多的sdf
u = extract_fields(bound_min, bound_max, resolution, query_func)
# 面绘制算法
# vertices 顶点坐标[N,3] N是根据具有情况而通过算法得出,与其他无关
# triangles 面索引[M,3] 索引指向顶点坐标数组中的对应顶点，3个顶点一个面
vertices, triangles = mcubes.marching_cubes(u, threshold)
# 提取域的对角顶点
b_max_np = bound_max.detach().cpu().numpy() # [3]
b_min_np = bound_min.detach().cpu().numpy() # [3]
# 缩小位移
vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
return vertices, triangles
extract_fields
The function of this function is to obtain appropriate extraction points (voxels) within the three-dimensional reconstruction range, and calculate the corresponding sdf value for each extraction point (voxel).
def extract_fields(bound_min, bound_max, resolution, query_func):
N = 64
# 根据提取域(方体)的对角顶点,获取提取域在各xyz轴的范围(max-min)和单位刻度((max-min)/resolution)
X = torch.linspace(bound_min[0], bound_max[0], resolution).split(N)
Y = torch.linspace(bound_min[1], bound_max[1], resolution).split(N)
Z = torch.linspace(bound_min[2], bound_max[2], resolution).split(N)
# 初始化对应方体的sdf值
u = np.zeros([resolution, resolution, resolution], dtype=np.float32)
with torch.no_grad():
for xi, xs in enumerate(X):
for yi, ys in enumerate(Y):
for zi, zs in enumerate(Z):
# 网格化
xx, yy, zz = torch.meshgrid(xs, ys, zs) # [N,N,N]
# [N^3,3]
pts = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1)
# 找到对应点的sdf
val = query_func(pts).reshape(len(xs), len(ys), len(zs)).detach().cpu().numpy()
# 为方体正确的赋sdf值
u[xi * N: xi * N + len(xs), yi * N: yi * N + len(ys), zi * N: zi * N + len(zs)] = val
return u
The execution diagram of the code is shown in the figure below. The orange squares are the extraction points (voxels). Smaller extraction points (voxels) can be divided into smaller extraction points (voxels) in more detail according to the division requirements.
Summarize
Introduce part of the code in the NeuS test phase as simply and in detail as possible: the process of validate_image rendering pictures and validate_mesh rebuilding the model. The remaining code in the testing phase will be explained later.