9.3.tensorRT高级(4)封装系列-自动驾驶案例项目self-driving-车道线检测

前言

杜老师推出的 tensorRT从零起步高性能部署 课程，之前有看过一遍，但是没有做笔记，很多东西也忘了。这次重新撸一遍，顺便记记笔记。

本次课程学习 tensorRT 高级-自动驾驶案例项目self-driving-车道线检测

课程大纲可看下面的思维导图

1. 车道线检测

这节我们学习车道线检测模型的分析，我们的目的是找到车道线检测的 onnx，分析其 onnx 的大致使用逻辑，然后写出最简洁版本的 predict.py，大体可以分为以下三步：

1. 打开车道线检测的 onnx，查看其输入与输出

2. 查看代码，找到 onnx 的预处理，分析得到预处理的逻辑

3. 针对获得的信息，编写 predict.py，尝试写出来

值得注意的是，在这个案例中，由于后处理过于复杂，因此考虑合并到 onnx 中，使得模型尽量的简单

在开始之前，我们先对车道线检测任务进行一个简单的分析

对于常规的框回归任务，例如求取下图中硬币在图像中的位置，cx,cy,w,h，其通常直接输出 4 个标量值进行回归

图1 常规框回归

目前最新的，大家更倾向于使用位置概率点乘其位置作为输出值，属于加权和，如下图所示

图2 位置概率

这种方法将回归的坐标以 n 个位置概率进行表示，例如对于 cx 的回归，表示为 5 个概率，可以认为对图像划分为 5 块，然后 cx 更有可能落到哪一块上进行表述。例如落在图像中心上时，其中心概率最高。有一种 attention 的味道。像 NanoDet、Alphapose 的后处理都与位置概率类似

车道线检测图如下所示：

对于车道线检测任务，我们是有一些先验知识的，比如车道线一样是位于图像下半部分，图像上半部分是天空无需考虑。另外检测的车道线通常是驾驶区域的 2 条加上两侧总共 4 条车道线；还有车道线点坐标的 y 值是知道的，我们会将图像按行划分为 N 个网格，每条车道线输出的点数就是 N，因此每个点的 y 我们是已知的；唯一不确定的是每个点的 x 坐标，这是需要模型学习出来的

那模型该如何回归这些点的 x 坐标呢？其实是通过位置概率来实现的，我们将图像按列分成 M 个网格，网络需要输出的总数量是 4xNxM，另外我们还要在列方向上增加一个维度，用来判断该点是否存在，因此网络的最终输出就是 4xNx(M+1)

我们来观察下车道线的 onnx 模型，如下图所示：

图3 onnx模型

可以看到 onnx 模型的输入是 1x3x288x800，其中输入图像的高度是 288，宽度是 800，输出是 1x201x18x4，其中 4 代表 4 条车道线，18 代表将图像下半部分划分为 18 行（即 N=18），201 代表将图像下半部分划分为 201 列（即 M=200）

我们分析总结可以得到如下信息：

1. 输入是：1x3x288x800

2. 输出是：1x201x18x4

3. 对于车道线检测任务而言有一些定义或者说是先验

• 只需要识别 4 条线
• 对于车道线基本是在地面上的，因此 y 方向可以从图像中心开始，也就是 anchor 起始坐标是图像中心到图像底部
• 对于车道线的检测，因为线是连续的，因此这里可以转变为离散的点的检测，对于一根线可以设计为 18 个点来描述
• 因此回归一个点，其 y 坐标已知，x 坐标需要回归出来
• 对于 x 的回归，采用了位置概率来表示，划分为 200 个网格表示其坐标
• 对于车道线的点是否存在这个问题，采用第 201 个概率表示，若这个点不存在，则 201 个点位置的值是最大的

我们再分析项目中的 image_processor/lane_engine.cpp 代码可以得出具体的预处理和后处理所做的工作：(详细分析请参照视频)

预处理部分：

• 图像的预处理直接是 image / 255.0
• 图像需要从 BGR 到 RGB
• 图像直接 resize 到 288x800

后处理部分：

• 对 0-200 维度进行 softmax，此时得到的是位置概率
• 对位置概率和位置索引点乘相加，得到 location，此时 location 是 18x4
• 对原始输出的最大值进行判断，决定该点是否存在
• 最后通过过滤得到 4 根线的坐标

我们可以简单的写个 demo 来验证下，代码如下：

import onnxruntime
import cv2
import numpy as np
import matplotlib.pyplot as plt
import scipy

session = onnxruntime.InferenceSession("workspace/ultra_fast_lane_detection_culane_288x800.onnx", provider_options=["CPUExecutionProvider"])

image = cv2.imread("workspace/imgs/dashcam_00.jpg")
show  = image.copy()
image = cv2.resize(image, (800, 288))
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image_tensor = (image / 255.0).astype(np.float32)
image_tensor = image_tensor.transpose(2, 0, 1)[None]

prob = session.run(["200"], {"input.1": image_tensor})[0][0]

print(prob.shape)

out_j = prob
prob = scipy.special.softmax(out_j[:-1, :, :], axis=0)
idx = np.arange(200) + 1
idx = idx.reshape(-1, 1, 1)
loc = np.sum(prob * idx, axis=0)

print(loc.shape)

# 201 x 18 x 4, 201 维度上找最大值
out_j = np.argmax(out_j, axis=0)
loc[out_j == 200] = 0

col_sample = np.linspace(0, 800 - 1, 200)
col_sample_w = col_sample[1] - col_sample[0]
ys = np.array([121, 131, 141, 150, 160, 170, 180, 189, 199, 209, 219, 228, 238, 248, 258, 267, 277, 287])

xs = loc * col_sample_w * show.shape[1] / 800
ys = ys * show.shape[0] / 288

colors = [(0, 255, 0), (255, 0, 0), (255, 0, 0), (0, 255, 0)]

for iline in range(4):
    for x, y in zip(xs[:, iline], ys):
        if x == 0:
            continue

        cv2.circle(show, (int(x), int(y)), 5, colors[iline], -1, 16)

cv2.imwrite("lane.jpg", show)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

输出如下图：

图4 输出

可以看到输出符合我们的预期，输出的车道线检测图如下所示：

图5 车道线检测效果图

那如果要使用 tensorRT 进行推理，你会发现后处理太复杂了，我们需要考虑将后处理放到 onnx 中，我们可以先导出后处理的 onnx 模型，然后把它添加到我们的 onnx 模型中，如下图所示：

图6 复杂后处理放onnx

总结

本次课程学习了开源项目中的车道线检测案例，主要是对车道线检测模型的 onnx 进行了简单分析，并通过对项目代码的分析将预处理和后处理部分理清楚，然后通过 onnxruntime 进行了简单验证，随后将复杂的后处理部分塞到 onnx 中方便后续在 tensorRT 上执行推理