【代码精读】开山之作MVSNet PyTorch版本超详细分析-CFANZ编程社区

总体结构

对于训练核心的代码有如下几个：

train.py: 整体深度学习框架(参数处理、dataset和DataLoader构建、epoch batch训练、计算loss梯度下降、读取/保存模型等)
models
- module.py: mvsnet所需的网络基础架构和方法(网络组成模块、投影变换homo_wraping、深度回归depth_regression)
- mvsnet.py: MVSNet整体Pipeline(特征提取深度回归残差优化网络定义、mvsnet_loss定义、核心四大步骤: 特征提取，cost volume构建、代价体正则化、深度图refine)
datasets
- data_yao.py: 定义MVSDataset(ref图和src图，投影矩阵，深度图真值，深度假设列表，mask)
utils.py: 一些小工具(logger、不同度量指标、系列wrapper方法)

项目整体文件结构

checkpoints(自己创建): 保存训练好的模型和tensorboard数据可视化所需的数据
outputs(自己创建): test的时候输出的预测深度图和点云融合后的点云文件等
lists: train, valid, test用的scan选择列表
evaluations: dtu数据集官方提供的matlab代码，主要用于测试重建点云的质量

DTU数据集结构

共128个scan

train: 79个
val: 18个
test: 22个

Train

【Cameras】

pair.txt: 只有一个，每个scan通用的

每个场景49个view的配对方式

49   # 场景的总视点数

0    # ref视点
src视点总数 第十个视点 视点选取时匹配的score   第一个视点
10           10          2346.41             1       2036.53 9 1243.89 12 1052.87 11 1000.84 13 703.583 2 604.456 8 439.759 14 327.419 27 249.278 

1
10 9 2850.87 10 2583.94 2 2105.59 0 2052.84 8 1868.24 13 1184.23 14 1017.51 12 961.966 7 670.208 15 657.218 

2
10 8 2501.24 1 2106.88 7 1856.5 9 1782.34 3 1141.77 15 1061.76 14 815.457 16 762.153 6 709.789 10 699.921

train/xxxxx_cam.txt：49个，每个视点有一个相机参数，不同scan是一致的(与Camera根目录下的camera参数文件不一样，代码里用的是train这个)
- 相机外参、相机内参、最小深度、深度假设间隔(之后还要乘以interval_scale才送去用)

【Depths】

深度图 & 深度图可视化

共128个scan
depth_map_00xx.pfm: 每个scan文件夹里49个视角的深度图 (深度以mm为单位)
****depth_visual_00xx.png: 还有49张深度图的png版本被用作mask(二值图，值为1的像素是深度可靠点，后续训练时只计算这些点的loss)

【Rectified】

原图

共128个scan
每个scan文件夹里里共49个视角*7种光照 = 343张图片
命名：rect_[view]_[light]_r5000.png
图片尺寸：640*512

Test

共有22个基准测试场景，对于每一个scan文件夹

pair.txt: 49个场景的配对信息，与train/Cameras/pair.txt是一样的，只是在每个scan里都复制了一份
images/: 该场景下49张不同视角的原始图片
cams/: 每个视点下的相机参数文件(❓不知道为什么有64个)

具体模块

代码中的数据维度

B: batch size 在研究数据维度时可以直接将这维去掉
C: 图像特征维度最开始是3-channels，后来通过特征提取网络变成32维
Ndepth: 深度假设维度，这里是192个不同的深度假设
H: 图像高度，原始是640，经过特征提取网络下采样了四倍，变成160
W: 图像宽度，同上，512 -> 128

dtu_yao/MVSDataset

MVSDataset(datapath, listfile, mode, nviews, ndepths=192, interval_scale=1.06)
- datapath: 数据集路径
- listfile: 数据列表(用哪些scan训练和测试都是提前定好的)
- mode: train or test
- nviews: 多视点总数(实现中取3=1ref+2src)
- ndepths: 深度假设数(默认假设192种不同的深度)
- interval_scale: 深度间隔缩放因子(数据集文件中定义了深度采样间隔是2.5，再把这个值乘以缩放因子，最终每隔2.5*1.06取一个不同的深度假设)
build_list(): 构建训练样本条目，最终的meta数组中共用27097条数据，每个元素如下：
```
# scan   light_idx      ref_view          src_view
# 场景    光照(0~6)  中心视点(估计它的深度)    参考视点
('scan2', 0, 0, [10, 1, 9, 12, 11, 13, 2, 8, 14, 27])
```
- read_img(): 将图像归一化到0～1(神经网络训练常用技巧，激活函数的取值范围大都是0～1，便于高效计算)
- 79个不同的scan
- 7种不同的光照
- 每个scan有49个不同的中心视点
read_cam_file(): 相机外参、相机内参、最小深度(都为425)、深度假设间隔(都为2.5)
getitem(): 取一组用来训练的数据
- imgs: 1ref + 2src（都归一化到0-1） (3, 3, 512, 640) 3个3channel的512*640大小的图片
- proj_metrices: 3个4*4投影矩阵 $\begin{bmatrix} R_{3,3} \ t_{3,1} \\ 0 \ 1 \end{bmatrix}$ (3, 4, 4)
  - 这里是一个视点就有一个投影矩阵，因为MVSNet中所有的投影矩阵都是相对于一个基准视点的投影关系，所以如果想建立两个视点的关系，他们两个都有投影矩阵，可以大致理解为 $B = P_B^{-1}P_AA$
  - 投影矩阵按理说应该是33的，这里在最后一行补了[0, 0, 0, 1]为了后续方便计算，所以这里投影矩阵维度是44
- depth: ref的深度图 (128, 160)
- depth_values: ref将来要假设的所有深度值 (从425开始每隔2.5取一个数，一共取192个)
  - 2.5还要乘以深度间隔缩放因子
- mask: ref深度图的mask(0-1二值图)，用来选取真值可靠的点 (128, 160)

dtu_yao_eval.py/MVSDataset

参数与训练时完全一致
build_list: 构建视点匹配列表，最终meta长度为1078，每个元素如下，与train相比没有光照变化
```
('scan1', 0, [10, 1, 9, 12, 11, 13, 2, 8, 14, 27])
```
read_cam_file(): 内参除4，最终生成的深度图也下采样4倍
read_img(): 裁掉下方的16个像素，图像尺寸变为1184*1600，裁剪后不需要修改内存
getitem():
- imgs: (5, 3, 1184, 1600) 测试的时候有5张图像，读的时候每张被裁剪掉了下面16像素
- proj_metrics: 5个投影矩阵，注意内参除了4倍
- depth_values: 深度假设范围，仍然是从425开始每隔2.5取一个数，一共192个
- filename: ref所在的文件夹名，如scan1/

train.py

构建训练参数

lrepochs: 训练中采用了动态调整学习率的策略，在第10，12，14轮训练的时候，让learning_rate除以2变为更小的学习率
wd: weight decay策略，作为Adam优化器超参数，实现中并未使用
numdepth: 深度假设数量，一共假设这么多种不同的深度，在里面找某个像素的最优深度
interval_scale: 深度假设间隔缩放因子，每隔interval假设一个新的深度，这个interval要乘以这个scale
loadckpt, logdir, resume: 主要用来控制从上次学习中恢复继续训练的参数
summary_freq: 输出到tensorboard中的信息频率
save_freq: 保存模型频率，默认是训练一整个epoch保存一次模型

################################  args  ################################
mode            train                           <class 'str'>       
model           mvsnet                          <class 'str'>       
dataset         dtu_yao                         <class 'str'>       
trainpath       /Data/MVS/train/dtu/            <class 'str'>       
testpath        /Data/MVS/train/dtu/            <class 'str'>       
trainlist       lists/dtu/train.txt             <class 'str'>       
testlist        lists/dtu/test.txt              <class 'str'>       
epochs          16                              <class 'int'>       
lr              0.001                           <class 'float'>     
lrepochs        10,12,14:2                      <class 'str'>       
wd              0.0                             <class 'float'>     
batch_size      1                               <class 'int'>       
numdepth        192                             <class 'int'>       
interval_scale  1.06                            <class 'float'>     
loadckpt        None                            <class 'NoneType'>  
logdir          ./checkpoints/d192              <class 'str'>       
resume          False                           <class 'bool'>      
summary_freq    20                              <class 'int'>       
save_freq       1                               <class 'int'>       
seed            1                               <class 'int'>       
########################################################################

构建SummaryWriter(使用tensorboardx进行可视化)
构建MVSDataset和DatasetLoader
构建MVSNet model，mvsnet_loss，optimizer
如果之前有训练模型，从上次末尾或指定的模型继续训练
train()
1. 设置milestone动态调整学习率
2. 对于每个epoch开始训练
  1. 对于每个batch数据进行训练
    1. 计算当前总体step：global_step = len(TrainImgLoader) * epoch_idx + batch_idx
    2. train_sample()
    3. 输出训练中的信息(loss和图像信息)
  2. 每个epoch训练完保存模型
  3. 每轮模型训练完进行测试(这里的测试应该理解为validation，因为用到了7种不同的光照，真正测试是eval，那时候只有一种光照)
    1. DictAverageMeter() 主要存储loss那些信息，方便计算均值输出到fulltest
    2. test_sample()

train_sample()

def train_sample(sample, detailed_summary=False):
    """训练DataLoader中取出的一次数据

    Args:
        sample ([imgs, proj_matrices, depth, depth_values, mask]): 1ref图+2src图，3个投影矩阵，深度图真值，深度假设列表，mask

    Returns:
        [loss, scalar_outputs, image_outputs]: 
            scalar_outputs: loss, abs_depth_error, thresXmm_error
            image_outputs: depth_est, depth_gt, ref_img, mask, errormap
    """
    model.train()           # 切换到train模式
    optimizer.zero_grad()   # 优化器梯度清零开始新一次的训练

    sample_cuda = tocuda(sample)    # 将所有Tensor类型的变量放到cuda计算
    depth_gt = sample_cuda["depth"] # 深度图ground truth数据
    mask = sample_cuda["mask"]      # mask用于将没有深度的地方筛除掉不计算loss

    outputs = model(sample_cuda["imgs"], sample_cuda["proj_matrices"], sample_cuda["depth_values"])     # 将数据放到model中进行训练
    depth_est = outputs["depth"]    # MVSNet得到的深度估计图

    loss = model_loss(depth_est, depth_gt, mask)    # 计算estimation和ground truth的loss，mask用于选取有深度值的位置，只用这些位置的深度真值计算loss
    loss.backward()     # loss函数梯度反传
    optimizer.step()    # 优化器中所有参数沿梯度下降一步

    scalar_outputs = {"loss": loss}         # 这轮训练得到的loss
    image_outputs = {
        "depth_est": depth_est * mask,      # 深度图估计(滤除掉本来就没有深度的位置)
        "depth_gt": sample["depth"],        # 深度图真值
        "ref_img": sample["imgs"][:, 0],    # 要估计深度的ref图
        "mask": sample["mask"]              # mask图(0-1二值图，为1代表这里有深度值)
    }

    if detailed_summary:
        image_outputs["errormap"] = (depth_est - depth_gt).abs() * mask                             # 预测图和真值图的区别部分
        scalar_outputs["abs_depth_error"] = AbsDepthError_metrics(depth_est, depth_gt, mask > 0.5)  # 绝对深度估计误差(整个场景深度估计的偏差平均值) mean[abs(est - gt)]
        scalar_outputs["thres2mm_error"] = Thres_metrics(depth_est, depth_gt, mask > 0.5, 2)        # 整个场景深度估计误差大于2mm的偏差偏差值(认为2mm之内都是估计的可以接受的) mean[abs(est - gt) > threshold]
        scalar_outputs["thres4mm_error"] = Thres_metrics(depth_est, depth_gt, mask > 0.5, 4)
        scalar_outputs["thres8mm_error"] = Thres_metrics(depth_est, depth_gt, mask > 0.5, 8)

    return tensor2float(loss), tensor2float(scalar_outputs), image_outputs

对于训练某个数据的输出如下：deptp_est, depth_gt, errormap, mask, ref_img即对应该函数的输出

eval.py

相机参数读取是内参intrinsics要除4
测试时参考图像用了5个视点

生成所有测试图片的深度图和confidence图
通过光度一致性和几何一致性优化深度图

save_depth(): 通过MVSNet进行test生成深度图的核心步骤

首先构建MVSDataset和Loader
对于每一条训练数据通过模型
- 输入：1ref + 4src，每个视点的投影矩阵，深度假设list
- 输出：深度图，photometric confidence
  - 深度图里的数据都是668.08545, 559.7229这类的真实物理距离(不满足像素的取值所以在mac上直接看是一片空白的)
  - 置信度里的数据是0～1之间的小数
将模型输出的两张图分别保存成pfm

reproject_with_depth(): 将ref的点投影到src上，再投影回来

参数：ref的深度图和内外参，src的深度图和内外参
返回值：重投影回来的深度图，重投影回来的x和y坐标，在src上的x和y坐标尺寸都是(128, 160)

check_gemoetric_consistency(): 几何一致性检验，调用上面的方法进行重投影，重投影后像素偏移<1 && 深度差<1%则通过校验

参数：ref的深度图和内外参，src的深度图和内外参
返回值：
- mask: 通过几何检验的mask图
- depth_reprojected: 重投影后的深度图
- x2d_src： ref这些像素在src上的坐标
- y2d_src： ref这些像素在src上的坐标

filter_depth(): 通过光度一致性约束和几何一致性约束filter上一步得到的深度图

photo_mask: 置信度图>0.8
geometric_mask: 至少3个src满足上面的几何一致性校验(重投影后像素偏移<1 && 深度差<1%)
filter每张ref的x y depth，并赋予颜色
最终融合生成最后的点云

module.py

ConvBnReLU, ConvBn, ConvBnReLU3D, ConvBn3D均为基础的网络结构，原始论文中例如特征提取中的一层即为这里的一个基础模块
BasicBlock, Hourglass3d在代码运行中并没使用

homo_warping: 将src图的特征体，根据ref和src的投影矩阵，投影到ref视角下

def homo_warping(src_fea, src_proj, ref_proj, depth_values):
    """投影变换：将src图像的特征投影到ref图像上，融合深度假设范围，得到warped volume

    Args:
        src_fea (src图像的特征)): [B, C, H, W]      此时的C已经是32维了
        src_proj (src图像的投影矩阵): [B, 4, 4]
        ref_proj (参考图像的投影矩阵)): [B, 4, 4]
        depth_values (深度假设范围数组): [B, Ndepth]

    Returns:
        [B, C, Ndepth, H, W] 最终得到的可以理解为src特征图按照不同的深度间隔投影到ref后构建的warped volume
    """
    batch, channels = src_fea.shape[0], src_fea.shape[1]
    num_depth = depth_values.shape[1]
    height, width = src_fea.shape[2], src_fea.shape[3]

    with torch.no_grad():       # 阻止梯度计算，降低计算量，保护数据
        proj = torch.matmul(src_proj, torch.inverse(ref_proj))      # src * ref.T
        rot = proj[:, :3, :3]  # [B,3,3] 取左上角三行三列得到旋转变换
        trans = proj[:, :3, 3:4]  # [B,3,1] 取最后一列的上面三行得到平移变换

        y, x = torch.meshgrid([torch.arange(0, height, dtype=torch.float32, device=src_fea.device),
                               torch.arange(0, width, dtype=torch.float32, device=src_fea.device)])   # 按照ref图像维度构建一张空的平面，之后要做的是根据投影矩阵把src中的像素映射到这张平面上
        y, x = y.contiguous(), x.contiguous()                   # 保证开辟的新空间是连续的(数组存储顺序与按行展开的顺序一致，transpose等操作是跟原tensor共享内存的)
        y, x = y.view(height * width), x.view(height * width)   # 将维度变换为图像样子
        xyz = torch.stack((x, y, torch.ones_like(x)))           # [3, H*W]
        xyz = torch.unsqueeze(xyz, 0).repeat(batch, 1, 1)       # [B, 3, H*W] unsqueeze先将维度变为[1, 3, H*W], repeat是为了将batch的维度引入进来
        rot_xyz = torch.matmul(rot, xyz)                        # [B, 3, H*W] 先将空白空间乘以旋转矩阵
        rot_depth_xyz = rot_xyz.unsqueeze(2).repeat(1, 1, num_depth, 1) \
                                                    * depth_values.view(batch, 1, num_depth, 1)  # [B, 3, Ndepth, H*W] 再引入Ndepths维度，并将深度假设值填入这个维度
        proj_xyz = rot_depth_xyz + trans.view(batch, 3, 1, 1)                   # [B, 3, Ndepth, H*W] 旋转变换后的矩阵+平移矩阵 -> 投影变换后的空白平面
        proj_xy = proj_xyz[:, :2, :, :] / proj_xyz[:, 2:3, :, :]                # [B, 2, Ndepth, H*W] xy分别除以z进行归一
        proj_x_normalized = proj_xy[:, 0, :, :] / ((width - 1) / 2) - 1         # [B, Ndepth, H*W] x方向按照宽度进行归一
        proj_y_normalized = proj_xy[:, 1, :, :] / ((height - 1) / 2) - 1        # y方向按照高度进行归一 @Q 这两步不太知道是干什么的？
        proj_xy = torch.stack((proj_x_normalized, proj_y_normalized), dim=3)    # [B, Ndepth, H*W, 2] 再把归一化后的x和y拼起来
        grid = proj_xy

    warped_src_fea = F.grid_sample(src_fea, grid.view(batch, num_depth * height, width, 2), \
                                    mode='bilinear', padding_mode='zeros')      # 按照grid中的映射关系，将src的特征图进行投影变换
    warped_src_fea = warped_src_fea.view(batch, channels, num_depth, height, width)     # 将上一步编码到height维的深度信息独立出来

    return warped_src_fea

depth_regression: 深度回归，根据之前假设的192个深度经过网络算完得到的不同概率，乘以深度假设，求得期望（最后在深度假设维度做了加法，所以运算后深度假设这一维度就没了）这个期望即是最终估计的最优深度，对应论文中的公式（3） $\sum_{d=d_{min}}^{d_{max}} d \times P(d)$

mvsnet.py

FeatureNet: 特征提取网络
- 将三通道的特征转换为32维的高维深度特征，同时图像进行了4倍下采样
- 输入：[3, H, W]
- 输出：[32, H/4, W/4] (32, 160, 128)
CostRegNet: cost volume正则化网络
- 先一路卷积降维，再一路反卷积升维，过程中把每步卷积和反卷积对应的volume都累加起来传播
- 输入：[B, C, D, H/4, W/4]
- 输出：[B, 1, D, H/4, W/4] (B, 1, 192, 160, 128)
RefineNet: 深度图边缘优化残差网络
- 输入: [B, 4, H/4, W/4] 4是因为img有三通道，depth有1通道
- 输出: [B, 1, H/4, W/4] (B, 1, 160, 128)
MVSNet Pipeline
1. feature extraction
  1. 输入：每张图片[B, 3, H, W]
  2. 输出：特征图[B, 32, H/4, W/4]
  3. 通过特征提取网络之后，原始图像的3-channel变为32维的高位特征，并且图像尺寸缩减到原来的1/4
2. differential homograph, build cost volume
  1. 将ref的32维特征和ref投影过来的高维特征累积构成原始cost volume
  2. 通过公式(2) $\frac{\sum_{i=1}^N(V_i - \bar{V_i})^2}{N}$ 计算方差得到最后的cost volume(在实现里通过 $\frac{\sum_{i=1}^N V_i^2}{N} - \bar{V_i}^2$ 公式简化计算)
  3. 最终的cost volume维度是[B, 32, 192, H/4, W/4]
3. cost volume regularization
  1. 首先通过代价体正则化网络进行进一步抽象，最终得到的维度是[B, 1, 192, H/4, W/4]
  2. 通过squeeze将维度为1的维度去除掉，得到[B, 192, H/4, W/4]
  3. 通过Softmax函数，将深度维度的信息压缩为0～1之间的分布，得到概率体probability volume
  4. 通过深度回归depth regression，得到估计的最优深度图 [B, H, W]
  5. 最后进行光度一致性校验，最终得到跟深度图尺寸一样的置信度图：简单来说就是选取上面估计的最优深度附近的四个点，再次通过depth regression得到深度值索引，再通过gather函数从192个深度假设层中获取index对应位置的数据
4. depth map refinement：将原图和得到的深度图合并送入优化残差网络，输出优化后的深度图
mvsnet_loss
- 根据公式(4)计算loss
- 由于是有监督学习，loss就是估计的深度图和ground truth深度图差一差的绝对值
- 唯一要注意的是，数据集中的mask终于在这发挥作用了，我们只选取mask>0.5，也就是可视化中白色的部分计算loss，只有这部分的点深度是valid的