【三维重建】【深度学习】NeuS代码Pytorch实现--训练阶段代码解析(上)

【三维重建】【深度学习】NeuS代码Pytorch实现–训练阶段代码解析(上)

论文提出了一种新颖的神经表面重建方法，称为NeuS，用于从2D图像输入以高保真度重建对象和场景。在NeuS中建议将曲面表示为有符号距离函数(SDF)的零级集，并开发一种新的体绘制方法来训练神经SDF表示，因此即使没有掩模监督，也可以实现更准确的表面重建。NeuS在高质量的表面重建方面的性能优于现有技术，特别是对于具有复杂结构和自遮挡的对象和场景。本篇博文将根据代码执行流程解析训练阶段具体的功能模块代码。

---

---

前言

在详细解析NeuS网络之前,首要任务是搭建NeuS【win10下参考教程】所需的运行环境，并完成模型的训练和测试，展开后续工作才有意义。
本博文是对NeuS训练阶段涉及的部分功能代码模块进行解析，其他代码模块后续的博文将会陆续讲解。

博主将各功能模块的代码在不同的博文中进行了详细的解析，点击【win10下参考教程】，博文的目录链接放在前言部分。

---

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网络模型设置-----
# 是否是train模式,复制重要的py文件到指定保存路径下
if self.mode[:5] == 'train':
    self.file_backup()

• 加载网络模型权重以及在拷贝重要的文件(训练阶段)

def load_checkpoint(self, checkpoint_name):
     # 加载指定权重文件
     checkpoint = torch.load(os.path.join(self.base_exp_dir, 'checkpoints', checkpoint_name), map_location=self.device)
     # 加载各神经网络模块
     self.nerf_outside.load_state_dict(checkpoint['nerf'])
     self.sdf_network.load_state_dict(checkpoint['sdf_network_fine'])
     self.deviation_network.load_state_dict(checkpoint['variance_network_fine'])
     self.color_network.load_state_dict(checkpoint['color_network_fine'])
     # 加载优化器
     self.optimizer.load_state_dict(checkpoint['optimizer'])
     # 加载保存时间迭代的epoch序号
     self.iter_step = checkpoint['iter_step']

     logging.info('End')

def file_backup(self):
    # 通过配置文件,拷贝指定文件所在的文件目录
    dir_lis = self.conf['general.recording']
    # 在训练结果数据保存路径下创建新的文件夹,用于保存拷贝的重要文件
    os.makedirs(os.path.join(self.base_exp_dir, 'recording'), exist_ok=True)

    # 从指定目录(./和./models/)下拷贝py文件到特定目录((./exp/...../recording)
    for dir_name in dir_lis:
        cur_dir = os.path.join(self.base_exp_dir, 'recording', dir_name)
        os.makedirs(cur_dir, exist_ok=True)
        files = os.listdir(dir_name)
        for f_name in files:
            if f_name[-3:] == '.py':
                copyfile(os.path.join(dir_name, f_name), os.path.join(cur_dir, f_name))
    # 从指定目录(./confs/)下拷贝配置文件到特定目录(./exp/...../recording)
    copyfile(self.conf_path, os.path.join(self.base_exp_dir, 'recording', 'config.conf'))

注意：每个功能模块的代码都在控制器的初始化函数中做了初始化，具体每个功能模块代码的使用位置、情况以及代码解析在之后执行过程中将详细讲解。

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']
# 逆矩阵×矩阵=>单位矩阵[4,4]
# [4,4]×[4,1]=>[4,1]
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)的空间范围之内。

计算相机内参、外参

关于相机内外参的知识点可以参考博主之前的博文【预备基础知识】关于四大坐标系的部分。

 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所表示的内容是下图所示的红色框中的投影矩阵，通过矩阵相乘已经将相机内外参融合。

注意：在图像产生rays时，上述公式也解释了代码中为什么用是(u,v,1)，即z方向是1而不是其他值。

拆分计算出相机的内参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模型初始化

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

学习率更新

本小节开始正式进行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负方向的延长线上(橙色线段)，源码对投影做了取反和归一化(既红色线段)，将投影(红色线段)长度(标量)作为中点(mid)，并计算出光线rays的最近点(near)和最远点(far)，near和far不是坐标点(n,3)，而是一个值(n,1)，可以理解成单位向量(rays_d)的长度比列系数，near可能是个负值。

个人理解：前面也提到了，反投影重建的目标物体被限制在半径为1的球体内，但需要注意的是对应的光心(相机位置rays_o)可能会在球体外。博主是直接print打印查看的rays_o内容，然后计算的模长。

光线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的产生和在其上进行的前景粗采样和背景采样。后续会讲解训练阶段的其他代码。