Anytime Stereo Image Depth Estimation on Mobile Devices 文章及代码调试

AnyNet是Yang Wang等人2016年提出的一种双目深度计算网络。最近做项目有用到该网络，其中碰到一些小坑为大家展示一下。第一部分是论文中重要地方的翻译和讲解；第二部分是源码修改及调试。
原文地址：https://arxiv.org/abs/1810.11408
源码地址：https://github.com/mileyan/AnyNet

翻译

摘要

机器人立体深度估计的许多应用都需要实时生成精确的视差图，但计算条件非常苛刻。当前最先进的算法迫使用户在缓慢生成精确映射和快速生成不精确映射之间做出选择，此外，这些方法通常需要太多参数，无法在电源或内存受限的设备上使用。针对这些缺点，我们提出了一种新的视差预测方法。与之前的工作相比，我们的端到端学习方法可以在推理时权衡计算量和准确性。深度估计是分阶段进行的，在此期间，可以在任何时间查询模型输出结果，得到其当前的最佳估计。我们的最终模型可以在10-35帧的范围内处理1242×375分辨率的图像，误差只会有微小的增加——可训练参数比其他渣渣少了两个数级。

介绍

前两段掠过
在本文中，我们提出了一种随时计算的视差估计方法，并提出了一种在速度和精度之间动态权衡的模型(见图1)。例如，一架高速飞行的自主无人机可以在高频率上询问我们的3D深度估计方法。如果一个物体出现在它的飞行路径上，它将能够迅速地感知它，并通过降低速度或执行规避操作做出相应的反应。当以低速飞行时，延迟并不那么重要，同样的无人机可以计算出更高的分辨率和更精确的3D深度地图，实现在拥挤场景中进行高精度导航或详细绘制环境地图等任务。

image

AnyNet预测时间线的例子。随着时间的推移，深度估计变得越来越准确。该算法可在任何时候轮询以返回当前深度图的最佳估计值。最初的估计可能足以触发避障机动，而后来的图像包含了足够的细节，可以进行更高级的路径规划程序

卷积网络深度估计的计算复杂度通常与图像分辨率成三次方，与被认为是的最大视差成线性关系。记住这些特征，我们不断细化深度图，同时始终确保分辨率或最大视差范围足够低，以确保最小的计算时间。我们从低分辨率(1/16)估算全视差范围的深度图开始。立方复杂度允许我们在几毫秒内计算出这个初始深度图(大部分时间花在初始特征提取和下采样上)。从这个低分辨率估计值开始，我们通过上行采样逐步增加视差图的分辨率，然后修正现在在高分辨率下很明显的误差。尽管使用了更高的分辨率，这些更新仍然很快，因为可以假设剩余的视差被限制在几个像素内，允许我们限制最大的视差，和相关的计算，仅10% - 20%的全范围。这些连续的更新除了初始的低分辨率设置外，完全避免了全范围的视差计算，并确保所有计算都是重复使用的，这使得我们的方法有别于大多数现有的多尺度网络结构。此外，我们的算法可以在任何时间轮询，以检索当前最佳估计深度图。可以实现广泛的可能帧率范围(在TX2模块上是10-35FPS)，同时在高延迟设置中仍然保留精确的视差估计。我们的整个网络可以在所有尺度上使用一个联合损失进行端到端训练，我们称之为AnyNet。
我们在多个基准数据集上评估了AnyNet，得到了各种令人鼓舞的结果:首先，AnyNet使用最先进的方法获得了具有竞争力的准确性，同时可训练参数比其他渣渣少了一个数量级。这对于资源受限的嵌入式设备尤其有影响。其次，我们发现深度卷积网络能够从粗糙的视差图预测残差。最后，包含最终空间传播模型(SPNet)的大大提高了视差图的质量，在计算成本(和参数存储需求)比现有方法低的情况下产生了最先进的结果。
图2显示了AnyNet体系结构的示意图布局。一个输入图像对首先通过U-Net特征提取器，它计算几个输出分辨率的特征地图(比例1/16,1/8,1/4)。在第一阶段，只计算最低尺度的特征(1/16)，并通过视差网络(图4)生成一个低分辨率的视差图(stage 1)。视差图估计右输入图像中每个像素的水平偏移量，它可以计算一个深度图。由于输入分辨率低，整个stage1的计算只需要几毫秒。如果允许更多的计算时间，我们进入stage2，在U-Net中继续计算，以获得更大尺度(1/8)的特征。而不是计算一个完整的视差图在这个更高的分辨率，在stage2，我们简单地修正已经计算的视差图从阶段1。首先，我们放大视差图来匹配阶段2的分辨率。然后我们计算一个残差贴图，其中包含小的修正，指定每个像素应该增加或减少多少视差贴图。如果时间允许，stage3与第stage2的过程相似，将分辨率从1/8到1/4再翻倍。stage4使用SPNet细化了阶段3的视差图。

在这里插入图片描述

在本节的其余部分中，我们将更详细地描述模型的各个组件。

1. U-Net特征提取网络: 图3详细说明了U-Net特征提取器，该特征提取器同时应用于左右图像。U-Net体系结构以不同的分辨率(1/16、1/8、1/4)计算特征图，在第1-3阶段stage作为输入，只有在需要时才计算。对原始输入图像进行最大pooling或strided convolution的下采样，然后用卷积滤波器进行处理。低分辨率的feature map捕获全局上下文，而高分辨率feature map捕获局部细节。在尺度为1/8和1/4时，最后的卷积层包含了之前计算的低尺度特征。

   
   
   在这里插入图片描述

2. 视差计算网络：它的输入为U-Net输出的特征图。我们使用这个组件来计算初始视差图(stage1)以及为后续校正计算剩余的视差图(stage2和3)。视差网络首先计算cost volume。这里的cost volume是指左侧图像中的一个像素与右侧图像中对应的一个像素之间的相似性。如果输入特征图尺寸为，那么cost volume尺寸为。其中(i, j, k)项描述左侧图像的像素(i, j)与右侧图像的像素(i, j−k)的匹配程度。M为考虑的最大视差。我们可以将左侧图像中的每个像素(i, j)表示为矢量，我们同样定义右视图矢量。M为本文考虑的最大视差。因此cost volume可认为是两个向量 之间的1范数距离。即： 。Cost volume可能由于图像中物体模糊、遮挡或模糊匹配而产生的错误。因此本文添加3D卷积层细化cost volume。提高其精确度。如果左图与右图相似。可得出他们的视差k，如果不相似则按照kendall等人的建议计算加权平均值。

   在这里插入图片描述

3. SpNet:为了进一步改善我们的结果，我们增加了最后的第四个阶段，其中我们使用一个空间传播网络(SPNet)来改进我们的视差预测。SPNet通过应用一个局部滤波器来锐化视差图，局部滤波器的权值是通过对左侧输入图像应用一个小CNN来预测的。我们表明，这种改进以相对较少的额外成本显著改善了我们的结果。（你可以把它看作是一个锐化过程）

视差计算网络公式太多点（我又懒得打了），以上内容足够了解网络内部参数细节了，如果需要了解具体运作方式还是得去看原文剩下的两段

源码Debug

文章源码是在ubuntu中运行的，我这里用的windows+WSL。首先从博文顶端下载代码后解压。

作者在github中的教程步骤是：

1. 处理Scene Flow数据集，准备好数据后训练sh ./create_dataset.sh。需要将Scene Flow存入D：\sampler, 大致看了代码，需要下载Scene Flow中的：FlyingThings3D、Driving、Monkaa。三个数据集大概50G吧（注意要下载带视差图的），而且下载及慢，迅雷+会员可以解决上述问题。。。
2. spn网络make：cd model/spn_1 sh make.sh，需要bash，我用的window子系统进行配置，可以借鉴我的另一篇配置WSL，配置到gcc能用就行。或者使用linux，或者干脆舍弃第四层输出
3. 训练：python main.py --maxdisp 192 --with_spn，作者推荐最大视差192，打开spn。
4. 微调：python finetune.py --maxdisp 192 --with_spn --datapath 具体地址/training，作者使用kitti数据集，打开spn。
5. 载入预训练模型：python finetune.py --maxdisp 192 --with_spn --datapath path-to-kitti2012/training/ --save_path results/kitti2012 --datatype 2012 --pretrained checkpoint /kitti2012_ck/checkpoint.tar \ --split_file checkpoint/kitti2012_ck/split.txt --evaluate
   可以仔细看一下上面的参数，有一个--pretrained checkpoint /kitti2012_ck/checkpoint.tar，在这里我们可以载入作者的训练结果。下载地址作者已经给出。

我建议直接跑作者的预训练模型：

1. 首先，下载以上数据集，建议kitti，因为比较小。
2. 其次，下载作者的训练结果，使用经过微调后的checkpoint.tar
3. 直接通过finetune.py看结果：python finetune.py --maxdisp 192 --datapath F:/data_scene_flow/training/ --save_path results/kitti2015 --datatype 2015 --pretrained checkpoint/kitti2015_ck/checkpoint.tar --evaluate 我用的scene flow的示例数据集，关闭spn（如果没编译spn就无法使用）。最终代码跑起来，得到每一个epoch的损失即准确度。
4. 作者给出的代码是没有可视化环节的！，因此我写了一个输出可视化py文件，它可以输出你的样本的视差图：
   首先是可视化函数show.py，它存入utils文件夹中：

#作者：Rayne
#作用：打印Output信息
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
def save_data(tensor_batch):
    for i,tensor in enumerate(tensor_batch):
        slice=tensor.cpu().detach().numpy()
        img_save = np.clip(slice, 0,2**16)
        img_save = (img_save * 256.0).astype(np.uint16).squeeze()
        image=Image.fromarray(img_save)
        image.save('data/stage'+str(i)+'.png')

def imshow(tensor_batch):

    for i,tensor in enumerate(tensor_batch):
        img_cpu = tensor.cpu().detach().numpy()
        img_save = np.clip(img_cpu, 0, 2**16)
        img_save = (img_save ).astype(np.uint16)
        img_save=img_save.squeeze()
        plt.subplot(3, 3, i + 1)
        plt.imshow(img_save)
        plt.title('stage '+str(i), fontsize=12)
    plt.show()
    

这里有个坑，outputs输出必须过滤小于零的数，并且乘256.否则输出的就是雪花图。
其次是可视化代码，把它放在主目录下：

import argparse
import torch
import torch.nn.parallel
import torch.utils.data
import models.anynet
from utils.show import save_data,imshow

from PIL import Image
from dataloader import preprocess
import random


parser = argparse.ArgumentParser(description='AnyNet with Flyingthings3d')

parser.add_argument('--maxdisplist', type=int, nargs='+', default=[12, 3, 3])
parser.add_argument('--datapath', default='F:/data_scene_flow/training/',
                    help='datapath')
parser.add_argument('--epochs', type=int, default=10,
                    help='number of epochs to train')
parser.add_argument('--train_bsize', type=int, default=6,
                    help='batch size for training (default: 12)')
parser.add_argument('--test_bsize', type=int, default=4,
                    help='batch size for testing (default: 8)')
parser.add_argument('--save_path', type=str, default='results/pretrained_anynet',
                    help='the path of saving checkpoints and log')
parser.add_argument('--resume', type=str, default=None,
                    help='resume path')
parser.add_argument('--lr', type=float, default=5e-4,
                    help='learning rate')
parser.add_argument('--with_spn', action='store_true', help='with spn network or not')
parser.add_argument('--print_freq', type=int, default=5, help='print frequence')
parser.add_argument('--init_channels', type=int, default=1, help='initial channels for 2d feature extractor')
parser.add_argument('--nblocks', type=int, default=2, help='number of layers in each stage')
parser.add_argument('--channels_3d', type=int, default=4, help='number of initial channels of the 3d network')
parser.add_argument('--layers_3d', type=int, default=4, help='number of initial layers of the 3d network')
parser.add_argument('--growth_rate', type=int, nargs='+', default=[4,1,1], help='growth rate in the 3d network')
parser.add_argument('--spn_init_channels', type=int, default=8, help='initial channels for spnet')


args = parser.parse_args()
def main():
    global args
    model = models.anynet.AnyNet(args)

    # print(model)
    model = torch.nn.DataParallel(model).cuda()
    #载入预训练模型
    checkpoint = torch.load('results/finetune_anynet/checkpoint.tar')
    model.load_state_dict(checkpoint['state_dict'])
    imgL, imgR = data_load('data/left_1.png', 'data/right_1.png')
    imgL=torch.unsqueeze(imgL,0)
    imgR=torch.unsqueeze(imgR,0)
    imgL = imgL.float().cuda()
    imgR = imgR.float().cuda()
    outputs=model(imgL,imgR)
    save_data(outputs)
    imshow(outputs)


def data_load(left_img_dir,right_img_dir):
    left_img=Image.open(left_img_dir).convert('RGB')
    right_img=Image.open(right_img_dir).convert('RGB')

    w, h = left_img.size
    th, tw = 256, 512
    # 变为256，512
    x1 = random.randint(0, w - tw)
    y1 = random.randint(0, h - th)
    left_img = left_img.crop((x1, y1, x1 + tw, y1 + th))
    right_img = right_img.crop((x1, y1, x1 + tw, y1 + th))

    # dataL = dataL.crop((w - 1232, h - 368, w, h))
    # dataL = np.ascontiguousarray(dataL, dtype=np.float32) / 256
    processed = preprocess.get_transform(augment=False)
    left_img = processed(left_img)
    right_img = processed(right_img)
    return left_img, right_img


main()
# imgL,imgR=data_load('data/left.png','data/right.png')

我的输入时Scene flow中的一组数据：

在这里插入图片描述

在这里插入图片描述

在这里插入图片描述

将你的文件存入data中，修改好名字后运行即可得出以下结果：

在这里插入图片描述

呃。。。。还行吧，并没有作者给出的图片那么准确。。。。

补充

损失函数

论文没提到损失函数，我从源码中得到内容如下：
由于具有四层输出，因此损失函数定义了一个损失权重：

每一层的损失函数采用smooth L1作为该层输出损失量：

暂时想到的就是这些了，有时间会继续补充后续内容。如果有什么问题欢迎留言，欢迎大家提出自己宝贵的意见。希望可以帮到你~

应用

作者给出的网络速度确实还算可以，跑起来很快，得到视差图可以根据摄像头内参输出深度图。下一步打算输出模型至ONNX，用c++封装后在板子上跑跑试试实时输出效果。