【三维重建】【深度学习】NeuS代码Pytorch实现–训练阶段代码解析(上)
论文提出了一种新颖的神经表面重建方法，称为NeuS，用于从2D图像输入以高保真度重建对象和场景。在NeuS中建议将曲面表示为有符号距离函数(SDF)的零级集，并开发一种新的体绘制方法来训练神经SDF表示，因此即使没有掩模监督，也可以实现更准确的表面重建。NeuS在高质量的表面重建方面的性能优于现有技术，特别是对于具有复杂结构和自遮挡的对象和场景。本篇博文将根据代码执行流程解析训练阶段具体的功能模块代码。
文章目录
- 【三维重建】【深度学习】NeuS代码Pytorch实现--训练阶段代码解析(上)
- 前言
- Runner控制器初始化
- Dataset数据管理器初始化
- NeuS模型初始化
- 计算相机内参、外参
- 学习率更新
- 图像训练集随机排序
- 随机光线rays生成
- 光线rays的最近、远点
- 光线rays上进行前景粗采样
- 光线rays上进行背景采样
- 总结
前言
在详细解析NeuS网络之前,首要任务是搭建NeuS【win10下参考教程】所需的运行环境，并完成模型的训练和测试，展开后续工作才有意义。
本博文是对NeuS训练阶段涉及的部分功能代码模块进行解析，其他代码模块后续的博文将会陆续讲解。
Runner控制器初始化
Runner作为一个封装好的控制器以方便训练和使用neus模型。
在exp_runner.py文件的class Runner下的 def __init__部分，这个部分是代码的基础部分，基本上所有用到功能代码的模块都在这里呗初始化。
- 读取配置文件：读取./confs目录下用户选择的对应配置文件的内容。
# Runner作为一个封装好的控制器以方便使用neus模型(训练和使用)
# 指定运行设备
self.device = torch.device("cuda" if torch.cuda.is_available() else "cup")
# 配置文件的路径
self.conf_path = conf_path
# 读取配置文件的内容
f = open(self.conf_path)
conf_text = f.read()
# CASE_NAME在配置文件的作用可以理解为占位的字符串,因此这里用case中的内容进行了替换
conf_text = conf_text.replace('CASE_NAME', case)
f.close()
# 将配置内容的格式变换成树形结构的形式
self.conf = ConfigFactory.parse_string(conf_text)
# 同理进行替换
self.conf['dataset.data_dir'] = self.conf['dataset.data_dir'].replace('CASE_NAME', case)
# 训练所需的(指定)数据集的存放位置
self.base_exp_dir = self.conf['general.base_exp_dir']
os.makedirs(self.base_exp_dir, exist_ok=True)
- 初始化数据管理类：管理图像数据集以及其对应的mask数据集和相机内外参数据等。
# 初始化一个data数据类
self.dataset = Dataset(self.conf['dataset'])
- 训练参数设置：关于训练迭代次数、保存模型权重周期、检验模型测试效果周期等。
# -----训练参数设置-----
# 开始训练的迭代epoch序号
self.iter_step = 0
# 结束训练的迭代epoch序号
self.end_iter = self.conf.get_int('train.end_iter')
# 训练过程中保存模型权重的周期
self.save_freq = self.conf.get_int('train.save_freq')
# 训练过程中打印必要信息的周期(loss和学习率)
self.report_freq = self.conf.get_int('train.report_freq')
# 训练过程中合成一个rgb视角图的周期
self.val_freq = self.conf.get_int('train.val_freq')
# 训练过程中生成一个ply模型的周期
self.val_mesh_freq = self.conf.get_int('train.val_mesh_freq')
# 训练过程中的batchsize(rays的个数)
self.batch_size = self.conf.get_int('train.batch_size')
# 理解成图片下采样的倍数
self.validate_resolution_level = self.conf.get_int('train.validate_resolution_level')
# 学习率
self.learning_rate = self.conf.get_float('train.learning_rate')
# 控制学习率变化的参数
self.learning_rate_alpha = self.conf.get_float('train.learning_rate_alpha')
# 是否使用白色背景
self.use_white_bkgd = self.conf.get_bool('train.use_white_bkgd')
# 预热启动区间
self.warm_up_end = self.conf.get_float('train.warm_up_end', default=0.0)
# 退火区间
self.anneal_end = self.conf.get_float('train.anneal_end', default=0.0)
# -----训练参数设置-----
- NeuS网络模型设置：初始化组成NeuS的个功能部分的具体神经网络，是否加载已完成训练预训练权重等。
# -----neus网络模型设置-----
# 计算loss时,sdf的梯度loss占整个loss的权重
self.igr_weight = self.conf.get_float('train.igr_weight')
# 计算loss时,mask的loss占整个loss的权重
self.mask_weight = self.conf.get_float('train.mask_weight')
# 是否在已有的最新模型基础上进行下一步操作
self.is_continue = is_continue
self.model_list = []
# 用于存放所以神经网络模型的参数
params_to_train = []
# nerf网络
self.nerf_outside = NeRF(**self.conf['model.nerf']).to(self.device)
# sdf网络
self.sdf_network = SDFNetwork(**self.conf['model.sdf_network']).to(self.device)
# 偏差网络
self.deviation_network = SingleVarianceNetwork(**self.conf['model.variance_network']).to(self.device)
# 渲染网络
self.color_network = RenderingNetwork(**self.conf['model.rendering_network']).to(self.device)
# 添加各个模型的参数
params_to_train += list(self.nerf_outside.parameters())
params_to_train += list(self.sdf_network.parameters())
params_to_train += list(self.deviation_network.parameters())
params_to_train += list(self.color_network.parameters())
# 设置优化器
self.optimizer = torch.optim.Adam(params_to_train, lr=self.learning_rate)
# 初始化neus神经网络
self.renderer = NeuSRenderer(self.nerf_outside,
self.sdf_network,
self.deviation_network,
self.color_network,
**self.conf['model.neus_renderer'])
# Load checkpoint
latest_model_name = None
# 选择已有的最新模型
if is_continue:
# 加载模型目录下的所有文件(可能包括非权重文件)
model_list_raw = os.listdir(os.path.join(self.base_exp_dir, 'checkpoints'))
model_list = []
# 讲权重文件单独筛选出来
for model_name in model_list_raw:
if model_name[-3:] == 'pth' and int(model_name[5:-4]) <= self.end_iter:
model_list.append(model_name)
# 对权重文件进行排序,并选择最新的权重
model_list.sort()
latest_model_name = model_list[-1]
# 若存在权重文件,neus神经网络加载权重
if latest_model_name is not None:
logging.info('Find checkpoint: {}'.format(latest_model_name))
self.load_checkpoint(latest_model_name)
# -----neus网络模型设置-----
注意：每个功能模块的代码都在控制器的初始化函数中做了初始化，具体每个功能模块代码的使用位置、情况以及代码解析在之后执行过程中将详细讲解。
Dataset数据管理器初始化
源码中定义了Dataset类用来存放图像数据集以及其相对应mask数据集和相机投影矩阵等信息，并能够根据NeuS具体的任务需求产生射线rays，用于后续进行采样。
这里暂时只对Dataset的初始化代码做解析。
# 设置指定的设备
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 配置文件
self.conf = conf
# 数据存放的路径
self.data_dir = conf.get_string('data_dir')
# 相机投影矩阵存放路径(渲染RGB图像和拍摄RGB图像的投影矩阵)
self.render_cameras_name = conf.get_string('render_cameras_name')
self.object_cameras_name = conf.get_string('object_cameras_name')
# 查看是否包含参数camera_outside_sphere,没有返回true
self.camera_outside_sphere = conf.get_bool('camera_outside_sphere', default=True)
# 查看是否包含参数scale_mat_scale,没有返回1.1
self.scale_mat_scale = conf.get_float('scale_mat_scale', default=1.1)
# 加载相机投影矩阵
camera_dict = np.load(os.path.join(self.data_dir, self.render_cameras_name))
self.camera_dict = camera_dict
# 所有图片的路径
self.images_lis = sorted(glob(os.path.join(self.data_dir, 'image/*.png')))
# 图像数量
self.n_images = len(self.images_lis)
# 加载图片数据集,并进行归一化处理
self.images_np = np.stack([cv.imread(im_name) for im_name in self.images_lis]) / 256.0
# 所有图片对用的mask的路径
self.masks_lis = sorted(glob(os.path.join(self.data_dir, 'mask/*.png')))
# 加载mask数据集,并进行归一化处理
self.masks_np = np.stack([cv.imread(im_name) for im_name in self.masks_lis]) / 256.0
# 图片坐标系到世界坐标系的矩阵4×4
self.world_mats_np = [camera_dict['world_mat_%d' % idx].astype(np.float32) for idx in range(self.n_images)]
# 用于坐标系归一化(0~1之间),渲染的场景都位于原点的单位球体内
self.scale_mats_np = []
self.scale_mats_np = [camera_dict['scale_mat_%d' % idx].astype(np.float32) for idx in range(self.n_images)]
# 图像数据集对应的内参
self.intrinsics_all = []
# 图像数据集的外参
self.pose_all = []
for scale_mat, world_mat in zip(self.scale_mats_np, self.world_mats_np):
P = world_mat @ scale_mat
P = P[:3, :4] # 去除最后一层的[0 0 0 1]
# 从相机投影矩阵中拆分出内参和外参(逆)
intrinsics, pose = load_K_Rt_from_P(None, P)
self.intrinsics_all.append(torch.from_numpy(intrinsics).float())
self.pose_all.append(torch.from_numpy(pose).float())
# 图像数据集
self.images = torch.from_numpy(self.images_np.astype(np.float32)).to(self.device) # [n_images, H, W, 3]
# mask数据集
self.masks = torch.from_numpy(self.masks_np.astype(np.float32)).to(self.device) # [n_images, H, W, 3]
# 内参
self.intrinsics_all = torch.stack(self.intrinsics_all).to(self.device) # [n_images, 4, 4]
# 内参的逆
self.intrinsics_all_inv = torch.inverse(self.intrinsics_all) # [n_images, 4, 4]
# 焦距
self.focal = self.intrinsics_all[0][0, 0]
# 外参
self.pose_all = torch.stack(self.pose_all).to(self.device) # [n_images, 4, 4]
# 图像尺寸
self.H, self.W = self.images.shape[1], self.images.shape[2]
# 图像的像素总数
self.image_pixels = self.H * self.W
object_bbox_min = np.array([-1.01, -1.01, -1.01, 1.0])
object_bbox_max = np.array([ 1.01, 1.01, 1.01, 1.0])
object_scale_mat = np.load(os.path.join(self.data_dir, self.object_cameras_name))['scale_mat_0']
# 逆矩阵×矩阵构=>造单位矩阵
object_bbox_min = np.linalg.inv(self.scale_mats_np[0]) @ object_scale_mat @ object_bbox_min[:, None] # [4,1]
object_bbox_max = np.linalg.inv(self.scale_mats_np[0]) @ object_scale_mat @ object_bbox_max[:, None] # [4,1]
self.object_bbox_min = object_bbox_min[:3, 0] # [3] xyz
self.object_bbox_max = object_bbox_max[:3, 0] # [3] xyz
关于object_bbox示意图如下图所示：
个人理解，这里简单说明一下，图像对应的投影矩阵乘上了缩放平移矩阵(P = world_mat @ scale_mat)，因此所有反投影生成的三维空间点都被限制在以原点为中心，X、Y和Z轴取值介于(-1,1)的空间范围之内。
NeuS模型初始化
NeuS神经网络模型是由多个神经网络模型构成的复合型神经网络模型，用于管理多个神经网络模型在不同阶段的使用。
# nerf网络
self.nerf = nerf
# sdf网络
self.sdf_network = sdf_network
# 偏差(标准)网络
self.deviation_network = deviation_network
# 渲染网络
self.color_network = color_network
# 粗采样点数
self.n_samples = n_samples
# 精采样点数
self.n_importance = n_importance
# 背景采样点数
self.n_outside = n_outside
# 理解为下采样倍数
self.up_sample_steps = up_sample_steps
# 扰动
self.perturb = perturb
计算相机内参、外参
关于相机内外参的知识点可以参考博主之前的博文【预备基础知识】关于四大坐标系的部分。
for scale_mat, world_mat in zip(self.scale_mats_np, self.world_mats_np):
P = world_mat @ scale_mat
P = P[:3, :4] # 去除最后一层的[0 0 0 1]
# 从相机投影矩阵中拆分出内参和外参(逆)
intrinsics, pose = load_K_Rt_from_P(None, P)
self.intrinsics_all.append(torch.from_numpy(intrinsics).float())
self.pose_all.append(torch.from_numpy(pose).float())
world_mats_np所表示的内容是下图所示的红色框中的投影矩阵，通过矩阵相乘已经将相机内外参融合。
拆分计算出相机的内参K以及外参Rt：
def load_K_Rt_from_P(filename, P=None):
if P is None:
# 加载相机的参数信息
lines = open(filename).read().splitlines()
if len(lines) == 4:
lines = lines[1:]
lines = [[x[0], x[1], x[2], x[3]] for x in (x.split(" ") for x in lines)]
P = np.asarray(lines).astype(np.float32).squeeze()
# 分解矩阵,将P分解为内参K和外参Rt
out = cv.decomposeProjectionMatrix(P)
# 内参
K = out[0]
# 外参旋转矩阵
R = out[1]
# 外参平移矩阵
t = out[2]
'''
因为分解计算出的K，k22位置上的值不等于1(理论上是必须是1),而是一个接近1的值(eg:1.3或1.5)
因此K/k22来保证k22位置为1
fx 0 0
0 fy 0
0 0 1
'''
K = K / K[2, 2]
# 内参(4×4)
'''
fx 0 0 0
0 fy 0 0
0 0 1 0
0 0 0 1
'''
intrinsics = np.eye(4)
intrinsics[:3, :3] = K
# 外参(4×4)
pose = np.eye(4, dtype=np.float32)
# 转置
pose[:3, :3] = R.transpose()
# 与上面类似,分解计算出的t4接近1,保证t4为理论值1
pose[:3, 3] = (t[:3] / t[3])[:, 0]
return intrinsics, pose
学习率更新
本小节开始正式进行NeuS的训练阶段(Runner.train)，但是介于内容比较丰富，博主挨个讲解代码执行流程中遇到的功能函数。
# 更新学习率
self.update_learning_rate()
def update_learning_rate(self):
# 热启动阶段
if self.iter_step < self.warm_up_end:
# 热启动阶段:learning_factor 从0~1
learning_factor = self.iter_step / self.warm_up_end
# 常规训练阶段
else:
alpha = self.learning_rate_alpha
# progress理解为训练的进度,从0~1
progress = (self.iter_step - self.warm_up_end) / (self.end_iter - self.warm_up_end)
# learning_factor,从1~alpha~1
learning_factor = (np.cos(np.pi * progress) + 1.0) * 0.5 * (1 - alpha) + alpha
for g in self.optimizer.param_groups:
# 更新学习率
g['lr'] = self.learning_rate * learning_factor
常规阶段的learning_factor 示意图如下图所示：
图像训练集随机排序
在exp_runner.py文件的class Runner下的 def train部分，dataset中记录了图像数据集的个数。
# 对图像序号进行随机排序
image_perm = self.get_image_perm()
Runner控制器的定义的函数。
def get_image_perm(self):
# 根据图像数据集总数随机初生成一个数字序号序列
return torch.randperm(self.dataset.n_images)
博主发现在源码中，训练过程中只对图像训练集进行过一次随机排序。
随机光线rays生成
在【NeuS总览】的博文中，已经简单介绍过这个过程。
在exp_runner.py文件的class Runner下的 def train部分。
data = self.dataset.gen_random_rays_at(image_perm[self.iter_step % len(image_perm)], self.batch_size)
Dataset数据管理器的定义的函数，在models/dataset.py文件下。
def gen_random_rays_at(self, img_idx, batch_size):
"""
Generate random rays at world space from one camera.
一个摄影机在世界空间生成随机光线
"""
# 在2D图像上随机选择batch_size个像素点(u,v)
pixels_x = torch.randint(low=0, high=self.W, size=[batch_size])
pixels_y = torch.randint(low=0, high=self.H, size=[batch_size])
# 获得像素点(u,v)颜色和mask的数据
color = self.images[img_idx][(pixels_y, pixels_x)] # [batch_size, 3]
mask = self.masks[img_idx][(pixels_y, pixels_x)] # [batch_size, 3]
# 相机坐标系下的方向向量:内参(逆)×像素坐标系
p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1).float() # [batch_size, 3]
p = torch.matmul(self.intrinsics_all_inv[img_idx, None, :3, :3], p[:, :, None]).squeeze() # [batch_size, 3]
# 单位方向向量:对方向向量做归一化处理
rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True) # [batch_size, 3]
# 世界坐标系下的方向向量：外参(逆)×相机坐标系
rays_v = torch.matmul(self.pose_all[img_idx, None, :3, :3], rays_v[:, :, None]).squeeze() # [batch_size, 3]
#世界坐标系下的光心位置(平移矩阵t)
rays_o = self.pose_all[img_idx, None, :3, 3].expand(rays_v.shape) # [batch_size, 3]
return torch.cat([rays_o.to(self.device), rays_v.to(self.device), color, mask[:, :1]], dim=-1).cuda() # [batch_size, 10(3+3+3+1)]
代码的执行示意图如下图所示，函数返回了光线rays穿过图片的rgb值以及对应像素位置的mask标签、rays_o(光心)和rays_v(单位方向向量)。
光线rays的最近、远点
在exp_runner.py文件的class Runner下的 def train部分。
near, far = self.dataset.near_far_from_sphere(rays_o, rays_d)
Dataset数据管理器的定义的函数，在models/dataset.py文件下。
def near_far_from_sphere(self, rays_o, rays_d):
# rays_d在rays_d的投影,是为了后续做归一化
a = torch.sum(rays_d**2, dim=-1, keepdim=True)
# 向量rays_o(原点到光心)在rays_d(单位方向向量)的投影
b = 2.0 * torch.sum(rays_o * rays_d, dim=-1, keepdim=True)
# mid是rays_o在rays_d的投影的终点(的负数)
mid = 0.5 * (-b) / a
# 以mid为中点,设定最近点near和最远点far
near = mid - 1.0
far = mid + 1.0
return near, far
代码的执行示意图如下图所示，rays_o本身是光心，这里看作原点到光心的向量，求出rays_o在单位方向向量rays_d上的投影，但是这个投影是在rays_d负方向的延长线上，源码做了取反和归一化，将其作为了中点计算出光线rays的最近点和最远点。
这里的near和far不是坐标点(n,3)，而是一个值(n,1)，可以理解成单位向量的比列系数。
光线rays上进行前景粗采样
在models/renderer.py文件的render函数内。
# 粗采样点采样区间以及粗采样点点集位置(均匀采样)
z_vals = torch.linspace(0.0, 1.0, self.n_samples)
z_vals = near + (far - near) * z_vals[None, :] # [batch_size,n_samples]
if perturb > 0:
# 在-0.5~0.5均匀分布的范围内中为每个ray的所有粗采样点随机选取一个统一的扰动系数
t_rand = (torch.rand([batch_size, 1]) - 0.5)
# 对均匀采样的粗采样点进行扰动
z_vals = z_vals + t_rand * 2.0 / self.n_samples # [batch_size,n_samples]
这里的扰动是每个ray都设置一个扰动，所有粗采样点都使用同一个扰动(batch_size个)。
光线rays上进行背景采样
在models/renderer.py文件的render函数内。在无mask分割前后背景的模式下，才会对背景进行采样。
z_vals_outside = None
if self.n_outside > 0:
# 粗采样点采样区间以及粗采样点点集位置(均匀采样)
z_vals_outside = torch.linspace(1e-3, 1.0 - 1.0 / (self.n_outside + 1.0), self.n_outside) # [batch_size,n_outside]
if perturb > 0:
if self.n_outside > 0:
# 背景采样点前后俩点的中点
mids = .5 * (z_vals_outside[..., 1:] + z_vals_outside[..., :-1])
# 远点集
upper = torch.cat([mids, z_vals_outside[..., -1:]], -1)
# 近点集
lower = torch.cat([z_vals_outside[..., :1], mids], -1)
# 在0~1均匀分布的范围内中为每个ray的每个背景采样点随机选取不同的扰动系数
t_rand = torch.rand([batch_size, z_vals_outside.shape[-1]])
# 对均匀采样的背景采样点进行扰动
z_vals_outside = lower[None, :] + (upper - lower)[None, :] * t_rand # [batch_size,n_outside]
# 在far以为的位置进行采样
if self.n_outside > 0:
z_vals_outside = far / torch.flip(z_vals_outside, dims=[-1]) + 1.0 / self.n_samples # [batch_size,n_outside]
这里的扰动是每个ray的每个背景采样点都单独设置一个扰动，所有背脊采样点独立使用一个扰动(batch_size×n_outside个)。
总结
尽可能简单、详细的介绍NeuS训练阶段部分代码：各个类的作用，以及光线rays的产生和在其上进行的前景粗采样和背景采样。后续会讲解训练阶段的其他代码。
