9.3.tensorRT高级(4)封装系列-自动驾驶案例项目self-driving-车道线检测
目录
-
- 前言
- 1. 车道线检测
- 总结
前言
杜老师推出的 tensorRT从零起步高性能部署 课程，之前有看过一遍，但是没有做笔记，很多东西也忘了。这次重新撸一遍，顺便记记笔记。
本次课程学习 tensorRT 高级-自动驾驶案例项目self-driving-车道线检测
课程大纲可看下面的思维导图
1. 车道线检测
这节我们学习车道线检测模型的分析，我们的目的是找到车道线检测的 onnx，分析其 onnx 的大致使用逻辑，然后写出最简洁版本的 predict.py，大体可以分为以下三步：
1. 打开车道线检测的 onnx，查看其输入与输出
2. 查看代码，找到 onnx 的预处理，分析得到预处理的逻辑
3. 针对获得的信息，编写 predict.py，尝试写出来
值得注意的是，在这个案例中，由于后处理过于复杂，因此考虑合并到 onnx 中，使得模型尽量的简单
在开始之前，我们先对车道线检测任务进行一个简单的分析
对于常规的框回归任务，例如求取下图中硬币在图像中的位置，cx,cy,w,h，其通常直接输出 4 个标量值进行回归
目前最新的，大家更倾向于使用位置概率点乘其位置作为输出值，属于加权和，如下图所示
这种方法将回归的坐标以 n 个位置概率进行表示，例如对于 cx 的回归，表示为 5 个概率，可以认为对图像划分为 5 块，然后 cx 更有可能落到哪一块上进行表述。例如落在图像中心上时，其中心概率最高。有一种 attention 的味道。像 NanoDet、Alphapose 的后处理都与位置概率类似
车道线检测图如下所示：
对于车道线检测任务，我们是有一些先验知识的，比如车道线一样是位于图像下半部分，图像上半部分是天空无需考虑。另外检测的车道线通常是驾驶区域的 2 条加上两侧总共 4 条车道线；还有车道线点坐标的 y 值是知道的，我们会将图像按行划分为 N 个网格，每条车道线输出的点数就是 N，因此每个点的 y 我们是已知的；唯一不确定的是每个点的 x 坐标，这是需要模型学习出来的
那模型该如何回归这些点的 x 坐标呢？其实是通过位置概率来实现的，我们将图像按列分成 M 个网格，网络需要输出的总数量是 4xNxM，另外我们还要在列方向上增加一个维度，用来判断该点是否存在，因此网络的最终输出就是 4xNx(M+1)
我们来观察下车道线的 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)
输出如下图：
可以看到输出符合我们的预期，输出的车道线检测图如下所示：
那如果要使用 tensorRT 进行推理，你会发现后处理太复杂了，我们需要考虑将后处理放到 onnx 中，我们可以先导出后处理的 onnx 模型，然后把它添加到我们的 onnx 模型中，如下图所示：
总结
本次课程学习了开源项目中的车道线检测案例，主要是对车道线检测模型的 onnx 进行了简单分析，并通过对项目代码的分析将预处理和后处理部分理清楚，然后通过 onnxruntime 进行了简单验证，随后将复杂的后处理部分塞到 onnx 中方便后续在 tensorRT 上执行推理