一 . 基本介绍
二、模拟平台安装和基本使用
三、基于OpenCV自动驾驶控制
3.1基于HSV提取空间的特定颜色区域
3.2基于canny提取算子的边缘轮廓
3.对区域定位感兴趣
3.4基于霍夫变换的线段检测
3.动作控制：转向角
四、基于深度学习的自动驾驶控制
4.1算法原理
4.2数据采集
4.3 模型训练和验证
4.3.1 代码结构组织
4.3.2 训练
4.3.3批量验证
4.4模型单张图片测试
4.5系统集成
树莓派小车+神经网络计算棒NCS2）-toc" style="margin-left:0px;">五.开发真正的自动驾驶汽车(树莓派汽车) 神经网络计算棒NCS2）
5.1 硬件环境
5.2汽车基本控制和摄像头测试
5.2.1 基本运动控制
5.2.2 树莓派和PC图像传输在计算机之间
5.3数据采集
5.4训练和验证
5.5部署基于树莓派和神经计算棒的部署
5.5.1 OpenVINO介绍
5.5.2部署流程
5.5.3导出onnx模型
套件-toc" style="margin-left:80px;">5.5.4 在Windows10上安装OpenVINO套件
5.5.5转换onnx模型为IR模型
5.5.6 神经计算棒推理在树莓派上实现
六.小结
参考文献
一 . 基本介绍
众所周知，自动驾驶技术已成为汽车行业的热门词汇。与传统汽车相比，自动驾驶汽车具有许多显著的优势，一是自动驾驶汽车可以缓解道路拥堵，提高交通安全；二是自动驾驶汽车可以为用户提供更方便、更好的旅行体验；三是自动驾驶术将带动整个汽车产业链发生巨变，催生交通出行新生态。如今，无论是科技巨头还是传统车企都在钻研自动驾驶的相关技术，国外典型代表有特斯拉、谷歌、Uber，国内的有百度、蔚来、华为，小鹏，滴滴，小米等。
自动驾驶汽车是一项艰巨而复杂的任务，涉及硬件、软件、AI算法于一体，开发一套完整的自动驾驶汽车系统需要多个领域的工程师、研究者协同合作才可能完成。那么我们个人是否能够搭建一套类似的简易自动驾驶小车，领略下自动驾驶的乐趣呢？ 接下来的一系列内容就围绕这个任务来实现。本文将会一步步的讲解如何搭建一套比较真实的、基于深度学习算法驱动的自动驾驶小车。
本文使用纯视觉方案实现一个端到端的能够在规定道路上行驶的自动驾驶小车，通过USB摄像头拍摄路面图像，实时分析并规划车辆转向角度。
。如果没有树莓派小车的读者可以通过模拟平台来学习。实际上不管是平台仿真还是真实树莓派小车环境，整个实现思路和方法都是一样的，只不过在树莓派小车上实现难度要大一些，因为需要额外的硬件支持，并且需要自己去调试、组装小车和布置道路环境。
为了开发方便，本文使用语言进行全流程开发，如果对Python语言不熟悉的读者可以先学习下Python基础语法再来学习本文内容。
本文使用作为基本的图像处理工具，使用作为深度学习框架。同样的，如果不熟悉这两个工具的读者可以先去官网学习下OpenCV和Pytorch的基本教程。
本文所有代码和数据均开源，下载地址（包含了模拟器、所有训练数据、完整代码）：
链接：https://pan.baidu.com/s/1yhIQpKCJlUqbU6xWzHiY9A 提取码：ln2a
具体效果如下视频所示（第1个是在模拟环境中通过深度学习实现，第2个是真实的使用树莓派搭建的深度学习自动驾驶小车）：
基于深度学习的自动驾驶小车
基于纯视觉的端到端自动驾驶小车，使用pytorch和神经网络计算棒实现 完整教程请在csdn上搜索“钱彬的博客 一文掌握基于深度学习的自动驾驶小车开发”
二、模拟平台安装和基本使用
下载地址：Releases · tawnkramer/gym-donkeycar · GitHub
该地址中提供的模拟器是基于Unity开发的，是经过删减过后的可执行程序，不再需要额外安装unity，下载下来后就可以直接运行。目前覆盖windows、Linux、Mac共3个版本。本文为了简单教学，只讲解如何在windows平台上运行和使用该模拟器。
具体的，下载Windows平台对应的DonkeySimWin.zip压缩包,解压后内容如下所示：
双击运行其中的donkey_sim.exe即可启动模拟器。
主界面如下所示：
该模拟器中提供了很多不同的赛道，在模拟器左侧是相关设置，可以设置不同的视角等。这里我们选择最简单的赛道generated road，因为这个赛道没有障碍物且跟我们真实高速公路环境比较像，上手比较容易。我们先设置左侧的Settings如下：
注意我们设置了paceCar并且勾选了manualDriving，这样我们就可以自己手动操作小车了，不需要使用内置的自动驾驶模式。
接下来单击generated road，进入具体的场景：
在场景中，如果我们前面主界面使用了手工模式（paceCar处勾选manualDriving），那么我们就可以通过键盘来操控小车进行体验了。与一般的赛车游戏类似（qq飞车、跑跑卡丁车等），W键表示前进，A表示左转，D表示右转，S表示后退。
在该模拟器中，控制小车的主要是两个参数：油门（W和S键）和转向角度（A和D键），这个与我们真实驾驶的汽车基本一致：挂挡+踩油门来控制前进动力，打方向盘控制车辆转向。为了能够实现自动驾驶，我们首先要能够根据这两个参数去控制模拟器里面小车的运行。我们怎么样通过Python代码来控制这个模拟器呢？
这个模拟器的好处就在于预留了Python控制接口，我们只需要安装一个驱动库就可以直接驱动模拟器里面的小车运行（提前安装好Git工具）：
pip install git+https://github.com/tawnkramer/gym-donkeycar安装好以后我们可以运行下面的python代码来实现小车的控制（注意：运行下面的代码前先启动模拟器，并停留在模拟器主界面上）：
# 导入库
import gym
import gym_donkeycar
import numpy as np
import cv2
# 设置模拟器环境
env = gym.make("donkey-generated-roads-v0")
# 重置当前场景
obv = env.reset()
# 运行100帧
for t in range(100):
# 定义控制动作
action = np.array([0.3,0.5]) # 动作控制，0.3表示转向，0.5表示油门
# 执行动作
obv, reward, done, info = env.step(action)
# 取一张图像保存
if t == 20:
img = cv2.cvtColor(obv,cv2.COLOR_RGB2BGR)
cv2.imwrite('test.jpg',img)
# 运行完以后重置当前场景
obv = env.reset()我们先分析下这段代码。下面这行代码用于设置模拟器环境，简单来说就是启用哪张地图：
env = gym.make("donkey-generated-roads-v0")在这个模拟器里面我们可以用到的地图如下所示：
- "donkey-warehouse-v0"
- "donkey-generated-roads-v0"
- "donkey-avc-sparkfun-v0"
- "donkey-generated-track-v0"
- "donkey-roboracingleague-track-v0"
- "donkey-waveshare-v0"
- "donkey-minimonaco-track-v0"
- "donkey-warren-track-v0"
- "donkey-thunderhill-track-v0"
- "donkey-circuit-launch-track-v0"
接下来的代码里面，我们运行了100帧，每帧都用固定的控制参数来执行：右转0.3、前进0.5。这两个字段就是我们前面提到的转向和油门值。下面给出这两个值的具体定义：
接下来使用np.array封装这两个参数，然后通过env.step来执行单步动作。执行完动作以后会返回一些信息，其中我们需要重点关注obs这个返回参数，这个参数表示当前位于小车正中间行车记录仪摄像头返回的一帧图像 ，图像宽160像素，高120像素，3通道RGB图像。如下图所示：
由于本文主要使用摄像头图像数据来控制小车运行，因此上述代码中我们抽取了一张图像并保存到本地用来分析并测试算法。
通过上述代码，我们就可以使用python调整两个参数[油门值、转向值]来控制小车的运行，并且可以得到小车每次运行后的图像数据。实现了这样一个逻辑，我们自然就可以通过建立自动驾驶模型，逐帧分析图像，然后控制小车的这两个参数来实现小车的自动驾驶。
本小节内容重点使读者重新熟悉下python基本使用方法，同时熟悉下这个小车驾驶模拟器，接下来我们将正式进入自动驾驶算法研发环节。
三、基于OpenCV的自动驾驶控制
在学习自动驾驶前，我们先看看传统算法是怎么解决上面这个任务的。只有综合比较了传统算法和深度学习算法，我们才能真正体会到深度学习的强大能力。
本小节，我们将使用传统图像处理算法进行行道线检测等步骤来控制小车运行在行道线内。一方面帮助读者巩固下基本的opencv图像处理技术，另一方面可以更清晰的认识这个任务难点，为后面实现基于深度学习的自动驾驶做好铺垫。
具体的，我们希望通过算法来控制小车，最终让这个小车稳定运行在行车道内。这里面涉及到两方面：。。行道线检测的目的就是希望能够根据检测到的行道线位置来计算最终应该转向的角度，从而控制小车始终运行在当前车道线内。
由于道路环境比较简单，针对我们这个任务，我们进一步简化我们的控制变量，我们只控制转向角度，对于油门值我们在运行时保持低匀速，这样我们的重点就可以放在一个变量上面—。
3.1基于HSV空间的特定颜色区域提取
颜色过滤是目前经常被使用到的图像处理技巧之一，例如天气预报抠像等，经常会使用绿幕作为背景进行抠图。本小节使用颜色过滤来初步提取出行道线。
从模拟平台的图像数据上进行分析，小车左侧是黄实线，右侧是白实线。我们希望小车一直运行在这两根线之间。因此，我们首先要定位出这两根线。我们可以通过颜色空间变换来定位这两根线。
为了方便将黄色线和白色线从图像中过滤出来，我们需要将图像从RGB空间转换到HSV空间再处理。
这里首先我们解释下RGB和HSV颜色空间的区别。
RGB 是我们接触最多的颜色空间，由三个通道表示一幅图像，分别为红色(R)，绿色(G)和蓝色(B)。这三种颜色的不同组合可以形成几乎所有的其他颜色。RGB 颜色空间是图像处理中最基本、最常用、面向硬件的颜色空间，比较容易理解。RGB 颜色空间利用三个颜色分量的线性组合来表示颜色，任何颜色都与这三个分量有关，而且这三个分量是高度相关的，所以连续变换颜色时并不直观，想对图像的颜色进行调整需要更改这三个分量才行。自然环境下获取的图像容易受自然光照、遮挡和阴影等情况的影响，即对亮度比较敏感。而 RGB 颜色空间的三个分量都与亮度密切相关，即只要亮度改变，三个分量都会随之相应地改变，而没有一种更直观的方式来表达。但是人眼对于这三种颜色分量的敏感程度是不一样的，在单色中，人眼对红色最不敏感，蓝色最敏感，所以 RGB 颜色空间是一种均匀性较差的颜色空间。如果颜色的相似性直接用欧氏距离来度量，其结果与人眼视觉会有较大的偏差。对于某一种颜色，我们很难推测出较为精确的三个分量数值来表示。所以，RGB 颜色空间适合于显示系统，却并不适合于图像处理。
基于上述理由，在图像处理中使用较多的是 HSV 颜色空间，它比 RGB 更接近人们对彩色的感知经验。非常直观地表达颜色的色调、鲜艳程度和明暗程度，方便进行颜色的对比。在 HSV 颜色空间下，比 BGR 更容易跟踪某种颜色的物体，常用于分割指定颜色的物体。HSV 表达彩色图像的方式由三个部分组成：
- Hue（色调、色相）
- Saturation（饱和度、色彩纯净度）
- Value（明度）
其中Hue用角度度量，取值范围为0～360°，表示色彩信息，即所处的光谱颜色的位置，如下图所示。
如果我们想要过滤出黄色线，那么我们就可以将色调范围控制在[30~90]之间即可。注意，在OpenCV中色调范围是[0~180]，因此上述黄色范围需要缩小1倍，即[15~45]。检测白色行道线也是采用类似的原理。
下面我们用代码实现一下：
import cv2
import numpy as np
#读取图像并转换到HSV空间
frame = cv2.imread('test.jpg')
hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
# 黄色线检测
lower_blue = np.array([15, 40, 40])
upper_blue = np.array([45, 255, 255])
yellow_mask = cv2.inRange(hsv, lower_blue, upper_blue)
cv2.imwrite('yellow_mask.jpg',yellow_mask)
# 白色线检测
lower_blue = np.array([0, 0, 200])
upper_blue = np.array([180, 30, 255])
white_mask = cv2.inRange(hsv, lower_blue, upper_blue)
cv2.imwrite('white_mask.jpg',white_mask)效果如下图所示：
可以看到，黄色区域和白色区域基本都检测出来了，当然还存在不少干扰区域，需要进一步处理。
3.2基于canny算子的边缘轮廓提取
目前我们仅获得了行道线区域，为了后续能够方便的计算行道线角度，我们需要得到行道线具体的轮廓/线段信息，即从区域中提取出线段。这里我们使用Canny算法实现。
Canny边缘检测是从不同视觉对象中提取有用的结构信息并大大减少要处理的数据量的一种技术，于1986年被提出，目前已广泛应用于各种计算机视觉系统。
Canny算法具体包括5个步骤：
1) 使用高斯滤波器，以平滑图像，滤除噪声。
2) 计算图像中每个像素点的梯度强度和方向。
3) 应用非极大值（Non-Maximum Suppression）抑制，以消除边缘检测带来的杂散响应。
4) 应用双阈值（Double-Threshold）检测来确定真实的和潜在的边缘。
5) 通过抑制孤立的弱边缘最终完成边缘检测。
具体实现细节我们不再详细剖析，在OpenCV中集成了canny算法，只需要一行代码即可实现canny边缘检测。
# 黄色线边缘提取
yellow_edge = cv2.Canny(yellow_mask, 200, 400)
cv2.imwrite('yellow_edge.jpg', yellow_edge)
# 白色线边缘提取white
white_edge = cv2.Canny(white_mask, 200, 400)
cv2.imwrite('white_edge.jpg', white_edge)代码中200和400这两个参数表示canny算子的低、高阈值，按照opencv教程一般可以不用修改。
最终效果如下所示：
3.3感兴趣区域定位
在利用OpenCV对图像进行处理时，通常会遇到一个情况，就是只需要对部分感兴趣区域（Region Of Interest, ROI）进行处理。例如针对我们这个模拟平台上的智能小车任务来说，对于黄色行道线，我们只关注图像左下部分，而对于白色行道线，我们只关注图像右下部分即可。至于图像其他部分因为我们通过人工分析知道，这些区域我们并不需要处理。因此，针对黄色边缘我们只需要提取图像左下块区域，针对白色边缘我们只需要提取图像右下块区域。
def region_of_interest(edges, color='yellow'):
'''
感兴趣区域提取
'''
height, width = edges.shape
mask = np.zeros_like(edges)
# 定义感兴趣区域掩码轮廓
if color == 'yellow':
polygon = np.array([[(0, height * 1 / 2),
(width * 1 / 2, height * 1 / 2),
(width * 1 / 2, height),
(0, height)]], np.int32)
else:
polygon = np.array([[(width * 1 / 2, height * 1 / 2),
(width, height * 1 / 2),
(width, height),
(width * 1 / 2, height)]], np.int32)
# 填充感兴趣区域掩码
cv2.fillPoly(mask, polygon, 255)
# 提取感兴趣区域
croped_edge = cv2.bitwise_and(edges, mask)
return croped_edge最终效果如下图所示：
到这里我们看到行道线区域基本定位的比较“干净”了。
3.4基于霍夫变换的线段检测
到目前，我们抽取出了比较精确的行道线轮廓，但是对于实际的自动驾驶任务来说还没有完成目标任务要求，我们要对行道线轮廓再进一步处理，得到行道线的具体线段信息（每条线段的起始点坐标）。本小节我们使用霍夫变换来完成这个任务。霍夫变换，英文名称Hough Transform，作用是用来检测图像中的直线或者圆等几何图形的。
具体的，一条直线的表示方法有好多种，最常见的是y=mx+b的形式。结合我们这个任务，对于最终检测出的感兴趣区域，怎么把图片中的直线提取出来。基本的思考流程是：如果直线 y=mx+b 在图片中，那么图片中，必需有N多点在直线上（像素点代入表达式成立），只要有这条直线上的两个点，就能确定这条直线。该问题可以转换为：求解所有的(m,b)组合。
设置两个坐标系，左边的坐标系表示的是(x,y)值，右边的坐标系表达的是(m,b)的值，即直线的参数值。那么一个(x,y)点在右边对应的就是一条线，左边坐标系的一条直线就是右边坐标系中的一个点。这样，右边左边系中的交点就表示有多个点经过(k,b)确定的直线。但是，该方法存在一个问题，(m,b)的取值范围太大。
为了解决(m,b)取值范围过大的问题，在直线的表示方面用 xcosθ+ysinθ=p 的规范式代替一般表达式，参数空间变成(θ,p)，0=<θ<=2PI。这样图像空间中的一个像素点在参数空间中就是一条曲线（三角函数曲线）。
霍夫线段检测算法原理步骤如下：
- 初始化(θ,p)空间，N(θ,p)=0 。（N(θ,p)表示在该参数表示的直线上的像素点的个数）
- 对于每一个像素点(x,y)，在参数空间中找出令 xcosθ+ysinθ=p 的(θ,p)坐标，N(θ,p)+=1
- 统计所有N(θ,p)的大小，取出N(θ,p)>threasold的参数 。（threadsold是预设的阈值）
OpenCV中封装好了基于霍夫变换的直线检测方法HoughLinesP，下面我们就来使用它进行线段检测。
def detect_line(edges):
'''
基于霍夫变换的直线检测
'''
rho = 1 # 距离精度：1像素
angle = np.pi / 180 #角度精度：1度
min_thr = 10 #最少投票数
lines = cv2.HoughLinesP(edges,
rho,
angle,
min_thr,
np.array([]),
minLineLength=8,
maxLineGap=8)
return lines我们可以打印返回的lines查看内容：
[[[ 1 94 47 62]]
[[143 94 156 103]]
[[103 67 119 77]]
[[ 1 86 41 60]]
[[101 52 158 56]]
[[104 69 159 100]]
[[ 5 52 22 53]]
[[129 63 140 63]]
[[ 87 50 110 52]]
[[ 0 88 17 77]]
[[ 88 55 134 89]]
[[ 2 94 36 70]]
[[ 17 50 29 50]]
[[ 23 73 42 60]]
[[ 90 56 110 70]]
[[ 1 56 16 51]]
[[128 55 148 56]]
[[ 0 89 8 84]]
[[ 88 56 112 75]]
[[151 101 159 104]]
[[ 30 73 43 61]]]返回的每组值都是一条线段表示线段起始位置(x_start,y_start,x_end,y_end)。可以看到小线段很多，我们对这些小线段做一下聚类和平均：
def average_lines(frame, lines, direction='left'):
'''
小线段聚类
'''
lane_lines = []
if lines is None:
print(direction + '没有检测到线段')
return lane_lines
height, width, _ = frame.shape
fits = []
for line in lines:
for x1, y1, x2, y2 in line:
if x1 == x2:
continue
# 计算拟合直线
fit = np.polyfit((x1, x2), (y1, y2), 1)
slope = fit[0]
intercept = fit[1]
if direction == 'left' and slope < 0:
fits.append((slope, intercept))
elif direction == 'right' and slope > 0:
fits.append((slope, intercept))
if len(fits) > 0:
fit_average = np.average(fits, axis=0)
lane_lines.append(make_points(frame, fit_average))
return lane_lines这里需要注意，由于图像的y坐标跟我们数学上经常遇到的y坐标方向是相反的（图像的y坐标轴正向是朝下的），因此，左侧黄色实线斜率是负值，右侧白色实线斜率是正值。上述代码我们将所有小线段的斜率和截距进行了平均，并且使用make_points函数重新计算了该平均线对应到图像上的起始坐标位置，make_points函数如下所示：
def make_points(frame, line):
'''
根据直线斜率和截距计算线段起始坐标
'''
height, width, _ = frame.shape
slope, intercept = line
y1 = height
y2 = int(y1 * 1 / 2)
x1 = max(-width, min(2 * width, int((y1 - intercept) / slope)))
x2 = max(-width, min(2 * width, int((y2 - intercept) / slope)))
return [[x1, y1, x2, y2]]上述函数最后返回的是坐标数值，这样看线段的坐标值不是很直观，我们可以写个脚本显式的观察这些线段：
def display_line(frame, lines, line_color=(0, 0, 255), line_width=2):
'''
在原图上展示线段
'''
line_img = np.zeros_like(frame)
if lines is not None:
for line in lines:
for x1, y1, x2, y2 in line:
cv2.line(line_img, (x1, y1), (x2, y2), line_color, line_width)
line_img = cv2.addWeighted(frame, 0.8, line_img, 1, 1)
return line_img上述代码我们将行道线按照一定权重与原图进行合成，方便我们查看最终效果。
最终检测效果如下所示：
从效果上看我们准确的将两条行道线检测了出来。接下来就是根据这两条行道线进行自动驾驶方向控制。
3.5动作控制：转向角
针对前面的测试图片，我们可以有效的检测出两条行道线（左侧黄色线和右侧白色线），但是在真实的运行过程中，可能会出现3种情况：
：这种情况一般是直线车道且车辆稳定运行在行道线内，这时候我们只需要根据检测出的两条行道线微调整角度即可。
：这种情况在转弯处容易出现，或者在车辆开始大范围偏离时出现，这时候我们的策略应该是向能够检测到的这条行道线方向前进。
：这种情况应该停下小车。
因此，针对三种情况我们需要不同的处理方式。代码如下所示：
# 计算转向角
x_offset = 0
y_offset = 0
if len(yellow_lane)>0 and len(white_lane)>0: # 检测到2条线
_, _, left_x2, _ = yellow_lane[0][0]
_, _, right_x2, _ = white_lane[0][0]
mid = int(width / 2)
x_offset = (left_x2 + right_x2) / 2 - mid
y_offset = int(height / 2)
elif len(yellow_lane)>0 and len(yellow_lane[0])==1: # 只检测到黄色行道线
x1, _, x2, _ = yellow_lane[0][0]
x_offset = x2 - x1
y_offset = int(height / 2)
elif len(white_lane)>0 and len(white_lane[0])==1: # 只检测到白色行道线
x1, _, x2, _ = white_lane[0][0]
x_offset = x2 - x1
y_offset = int(height / 2)
else: # 一条线都没检测到
print('检测不到行道线，退出程序')
break
angle_to_mid_radian = math.atan(x_offset / y_offset)
angle_to_mid_deg = int(angle_to_mid_radian * 180.0 / math.pi)
steering_angle = angle_to_mid_deg/45.0
action = np.array([steering_angle, 0.3]) # 油门值恒定到这里我们就可以开始启动程序了。完整代码如下所示：
# 导入系统库
import cv2
import numpy as np
import math
import gym
import gym_donkeycar
# 导入自定义库
from tools import region_of_interest,detect_line,make_points,average_lines,display_line
def main():
'''
主函数
'''
# 设置模拟器环境
env = gym.make("donkey-generated-roads-v0")
# 重置当前场景
obv = env.reset()
# 开始启动
action = np.array([0, 0.3]) # 动作控制，第1个转向值，第2个油门值
# 执行动作
obv, reward, done, info = env.step(action)
# 获取图像
frame = cv2.cvtColor(obv, cv2.COLOR_RGB2BGR)
# 运行1000次动作
for t in range(1000):
# 转换图像到HSV空间
height, width, _ = frame.shape
hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
# 黄色区域检测
lower_blue = np.array([15, 40, 40])
upper_blue = np.array([45, 255, 255])
yellow_mask = cv2.inRange(hsv, lower_blue, upper_blue)
# 白色区域检测
lower_blue = np.array([0, 0, 200])
upper_blue = np.array([180, 30, 255])
white_mask = cv2.inRange(hsv, lower_blue, upper_blue)
# 黄色线边缘提取
yellow_edge = cv2.Canny(yellow_mask, 200, 400)
# 白色线边缘提取
white_edge = cv2.Canny(white_mask, 200, 400)
# 黄色线感兴趣区域提取
yellow_roi = region_of_interest(yellow_edge, color='yellow')
# 白色线感兴趣区域提取
white_roi = region_of_interest(white_edge, color='white')
# 黄色线段检测
yellow_lines = detect_line(yellow_roi)
yellow_lane = average_lines(frame, yellow_lines, direction='left')
#yellow_show = display_line(frame, yellow_lane)
# 白色线段检测
white_lines = detect_line(white_roi)
white_lane = average_lines(frame, white_lines, direction='right')
#white_show = display_line(frame, white_lane, line_color=(255, 0, 0))
# 计算转向角
x_offset = 0
y_offset = 0
if len(yellow_lane)>0 and len(white_lane)>0: # 检测到2条线
_, _, left_x2, _ = yellow_lane[0][0]
_, _, right_x2, _ = white_lane[0][0]
mid = int(width / 2)
x_offset = (left_x2 + right_x2) / 2 - mid
y_offset = int(height / 2)
elif len(yellow_lane)>0 and len(yellow_lane[0])==1: # 只检测到黄色行道线
x1, _, x2, _ = yellow_lane[0][0]
x_offset = x2 - x1
y_offset = int(height / 2)
elif len(white_lane)>0 and len(white_lane[0])==1: # 只检测到白色行道线
x1, _, x2, _ = white_lane[0][0]
x_offset = x2 - x1
y_offset = int(height / 2)
else: # 一条线都没检测到
print('检测不到行道线，退出程序')
break
angle_to_mid_radian = math.atan(x_offset / y_offset)
angle_to_mid_deg = int(angle_to_mid_radian * 180.0 / math.pi)
steering_angle = angle_to_mid_deg/45.0
action = np.array([steering_angle, 0.3]) # 油门值恒定
# 执行动作
obv, reward, done, info = env.step(action)
# 重新获取图像
frame = cv2.cvtColor(obv, cv2.COLOR_RGB2BGR)
# 运行完以后重置当前场景
obv = env.reset()
if __name__ == '__main__':
'''
主函数入口
'''
main()
下面是实际执行效果：
基于OpenCV的自动驾驶小车模拟平台控制
可以看到我们通过简单的图像处理方法实现了这样模拟平台上的小车控制。
：如果当前生成的赛道有“十字路口交叉”（每次重新进入赛道其生成的赛道都是随机绘制的），那么在运行的时候可能会出现失败、跑出赛道的现象。因为这种十字路口我们在程序中没有考虑。如何规避这个问题，有兴趣的读者可以自行研究。
本文更多的关注基于深度学习的图像处理技术，对于传统的图像处理算法（例如霍夫变换等）本文不再深入分析，读者如果对这些传统图像处理算法不熟悉的可以自行再查阅资料深入研究。
截止到目前为止，我们借助上面这个基于opencv的自动驾驶模拟平台，我们重新巩固了Python、opencv图像处理的基本使用方法，了解了自动驾驶项目的难点，对整个处理流程有了更进一步的认识。需要说明的是，尽管我们上述操作步骤是针对我们这个自动驾驶模拟平台的，但是以上步骤同样适用于很多其他图像处理任务，很多传统的图像处理任务都涵盖颜色空间变换、特定颜色物体提取、感兴趣区域过滤、霍夫变换等步骤，因此掌握上述常规的图像处理技术是非常重要的。
四、基于深度学习的自动驾驶控制
在上一节中我们通过OpenCV图像处理技术实现了一个简易的自动驾驶小车。但是很明显，这辆自动驾驶小车的适应性很差，当图像中有相同颜色的干扰物出现时，那么对于这辆自动驾驶的小车来说就是顶级灾难。另外，我们需要大量人工定义的参数，例如行道线颜色（黄色或白色）、颜色阈值、霍夫变换阈值等，而且一旦地图环境换了，所有这些参数我们都得重新手工调整，这些参数之间又有一定的耦合性，参数调整很麻烦。很显然，这种处理方法普适性不好。
那么能不能丢给机器一大堆图片，让机器自己去学习如何从当前图像中分析出小车应该转向的合适角度？如果没有接触过深度学习，那么乍一听这个想法简直是天方夜谭，然而深度学习确实做到了。这就是为什么近十年深度学习在图像处理领域取得了全面成功。深度学习能够从大量图像数据中自行学习高层次语义特征，完成媲美人类甚至超越人类的推理水平，整个学习过程不用人为干预，我们要做的就是“喂”一堆图片并且设定好需要优化的目标函数即可。当我们“喂”的图片越多、种类越丰富，那么最终机器学习到的驾驶水平越强，而且适应性越好。
本小节开始我们将正式进入基于深度学习的自动驾驶领域。
4.1 算法原理
本项目实现思路参考2016年英伟达发表的论文《End to End Learning for Self-Driving Cars》。这篇文章提出的方法核心思想就是使用神经网络自动提取图像特征，从传统的 image -> features -> action变成了image -> action。该论文使用了深度网络结构，大大增强了图像特征提取能力，最终取得了不错的效果，其训练的模型不论是普通道路还是高速路，不论有道路标线还是没有道路标线都非常有效，解决了传统算法泛化性能差的问题。本文方法的测试性能非常好，在16年自动驾驶研究火热时，是一篇影响力很大的文章，即使放到现在，也是作为自动驾驶入门必读的Paper。
整个算法原理很简单，是对真实人类操作的一个模拟。对于我们人类驾驶员来说，假设我们正在驾驶这辆车，我们的执行流程跟上面算法也是一样的。首先我们用眼睛观看路面，然后我们的大脑根据当前眼睛看到的路面情况“下意识”的转动方向盘，转动一个我们认为合适的角度，从而避免车辆开出路面。这篇论文算法实现原理也是这样，具体如下图所示：
通过中间摄像头采集图像，然后图像输入到预先训练好的CNN网络，这个网络的输出是一个转向角度（可以理解为方向盘的转向角度），有了这个角度就可以控制小车按照这个角度进行转向。
有了这样一个模式，我们就只需要想办法训练这个CNN模型，针对每帧图像，都有一个我们认为合适的转向角度输出，即输入图像，输出一个回归值。具体模型结构如下图所示：
整个模型结构并不复杂，就是一堆的普通的CNN卷积神经网络模块按照顺序堆叠，最后使用全连接网络输出回归值。这个模型一共包含30层，由于其输入精度比较低(66x200)，因此推理速度也是比较快的，借助GPU可以实现实时推理。具体的，图像首先经过Normalization标准化，然后经过5组卷积层处理，最后拉平以后通过4个全连接层输出一个回归值，这个回归值就是我们项目中的转向角。
这里我们会遇到一个问题，训练上述深度神经网络我们需要大量的数据，即每帧图像以及对应的最佳转向值，这些数据怎么来呢？这篇论文里提出了一个方法，既然是模拟人类行为，那么只要让驾驶水平高超的“老司机”在相关赛道上进行手动驾驶，驾驶时一边记录每帧图像同时记录当前帧对应的操控的转向角，这样一组组数据记录下来就是我们认为的“最佳”训练数据。训练时，将模型预测的角度与给定图像帧的期望转向角度进行比较，误差通过反向传播反馈到CNN训练过程中，如下图所示。从图可以看出，这个过程在一个循环中重复，直到误差（本例中使用均方误差）足够低，这意味着模型已经学会了如何合理地转向。事实上，这是一个非常典型的图像分类训练过程，只不过这里预测输出是数值（回归值）而不是对象类别（分类概率）。
可以想象，如果能够完全的训练好这个模型，那么最终模型的输出结果是非常接近人类驾驶经验的。这篇论文通过大量实验证明，上述模型能够直接从拍摄的路面图像中有效的学习到最终的转向角，省去了传统算法颜色区域检测、感兴趣区域选择、霍夫变换等一系列复杂的耦合步骤。这篇论文做了一组实验，通过收集不到一百小时的少量训练数据进行训练，最后得到的模型足以支持在各种条件下操控车辆，比如高速公路、普通公路和居民区道路，以及晴天、多云和雨天等天气状况。
需要说明的是，这个模型的输出仅有一个转向角度，这样容易学习成功。如果输出变量再多一些（例如油门值、摄像头角度、行人避障等），那么这个模型还需要再进一步优化，感兴趣的读者可以借鉴近两年的论文进行深入研究（毕竟真正的自动驾驶是very very复杂的，要考虑的情况非常多）。 接下来我们就按照这个算法流程进行实现。
4.2 数据采集
针对我们采用的自动驾驶模拟平台，为了能够采集到每帧图像及对应的最佳转向角度，我们可以使用前面第2节方法编写控制代码通过键盘控制小车（低匀速运行，仅仅只需要控制转向角度），然后记录每帧数据即可。这种模式是真实自动驾驶使用的，但是需要我们自己把自己练成经验充足的“老司机”，然后再去教会算法怎么驾驶。这样比较麻烦，这里可以有一种“偷懒”的办法。我们使用前面调参调的不错的OpenCV自动驾驶版本，使用OpenCV算法自动驾驶，然后记录每帧图像及对应角度。尽管这个OpenCV自动驾驶水平本身也一般（没有一直控制在两条行道线的绝对正中间），但是胜在能够基本稳定在行道线内。本文只是一个自动驾驶入门项目，可以采用这样的方法收集数据，来快速验证深度学习自动驾驶可行性。真实项目的话还是需要向“老司机”学习的。
完整采集代码如下：
# 导入系统库
import cv2
import numpy as np
import math
import gym
import gym_donkeycar
# 导入自定义库
from tools import region_of_interest, detect_line, make_points, average_lines, display_line
def main():
'''
主函数
'''
# 设置模拟器环境
env = gym.make("donkey-generated-roads-v0")
# 重置当前场景
obv = env.reset()
# 开始启动
action = np.array([0, 0.3]) # 动作控制，第1个转向值，第2个油门值
# 执行动作
obv, reward, done, info = env.step(action)
# 获取图像
frame = cv2.cvtColor(obv, cv2.COLOR_RGB2BGR)
# 运行4000次动作
pic_index = 0
for t in range(4000):
# 转换图像到HSV空间
height, width, _ = frame.shape
hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
# 黄色区域检测
lower_blue = np.array([15, 40, 40])
upper_blue = np.array([45, 255, 255])
yellow_mask = cv2.inRange(hsv, lower_blue, upper_blue)
# 白色区域检测
lower_blue = np.array([0, 0, 200])
upper_blue = np.array([180, 30, 255])
white_mask = cv2.inRange(hsv, lower_blue, upper_blue)
# 黄色线边缘提取
yellow_edge = cv2.Canny(yellow_mask, 200, 400)
# 白色线边缘提取
white_edge = cv2.Canny(white_mask, 200, 400)
# 黄色线感兴趣区域提取
yellow_roi = region_of_interest(yellow_edge, color='yellow')
# 白色线感兴趣区域提取
white_roi = region_of_interest(white_edge, color='white')
# 黄色线段检测
yellow_lines = detect_line(yellow_roi)
yellow_lane = average_lines(frame, yellow_lines, direction='left')
#yellow_show = display_line(frame, yellow_lane)
# 白色线段检测
white_lines = detect_line(white_roi)
white_lane = average_lines(frame, white_lines, direction='right')
#white_show = display_line(frame, white_lane, line_color=(255, 0, 0))
# 计算转向角
x_offset = 0
y_offset = 0
if len(yellow_lane) > 0 and len(white_lane) > 0: # 检测到2条线
_, _, left_x2, _ = yellow_lane[0][0]
_, _, right_x2, _ = white_lane[0][0]
mid = int(width / 2)
x_offset = (left_x2 + right_x2) / 2 - mid
y_offset = int(height / 2)
elif len(yellow_lane) > 0 and len(yellow_lane[0]) == 1: # 只检测到黄色行道线
x1, _, x2, _ = yellow_lane[0][0]
x_offset = x2 - x1
y_offset = int(height / 2)
elif len(white_lane) > 0 and len(white_lane[0]) == 1: # 只检测到白色行道线
x1, _, x2, _ = white_lane[0][0]
x_offset = x2 - x1
y_offset = int(height / 2)
else: # 一条线都没检测到
print('检测不到行道线，退出程序')
break
angle_to_mid_radian = math.atan(x_offset / y_offset)
angle_to_mid_deg = int(angle_to_mid_radian * 180.0 / math.pi)
steering_angle = angle_to_mid_deg / 45.0
action = np.array([steering_angle, 0.1]) # 油门值恒定
# 记录当前图像和转向角度
img_path = "log/{:d}_{:.4f}.jpg".format(pic_index, steering_angle)
cv2.imwrite(img_path, frame)
pic_index += 1
# 执行动作
obv, reward, done, info = env.step(action)
# 重新获取图像
frame = cv2.cvtColor(obv, cv2.COLOR_RGB2BGR)
# 运行完以后重置当前场景
print('结束本次采集')
obv = env.reset()
if __name__ == '__main__':
'''
主函数入口
'''
main()
我们将每帧图像对应的“最佳”转向角度以文件名的形式保存，最终收集图片如下所示：
图片名采用“图片帧号_转向角度.jpg”的形式命名。上述代码每次跑完会在log目录下生成4000多张图片。由于每次的地图都是随机生成的，因此我们可以多跑几次，多收集一些数据。
最终共采集10个文件夹图片，总共4万张图片：
接下来我们需要对这些图片进行整理，拆分数据集用于训练和验证。
详细脚本代码create_data_lists.py如下：
# 导入系统库
import os
import random
def creat_data_list(dataset_path, file_list, mode='train'):
'''
创建txt文件列表
'''
with open(os.path.join(dataset_path, (mode + '.txt')), 'w') as f:
for (imgpath, angle) in file_list:
f.write(imgpath + ' ' + str(angle) + '\n')
print(mode + '.txt 已生成')
def getFileList(dir, Filelist, ext=None):
"""
获取文件夹及其子文件夹中文件列表
输入 dir: 文件夹根目录
输入 ext: 扩展名
返回: 文件路径列表
"""
newDir = dir
if os.path.isfile(dir):
if ext is None:
Filelist.append(dir)
else:
if ext in dir[-3:]:
Filelist.append(dir)
elif os.path.isdir(dir):
for s in os.listdir(dir):
newDir = os.path.join(dir, s)
getFileList(newDir, Filelist, ext)
return Filelist
def main():
'''
主函数
'''
# 设置参数
org_img_folder = './data/simulate' # 数据集根目录
train_ratio = 0.8 # 训练集占比
# 检索jpg文件
jpglist = getFileList(org_img_folder, [], 'jpg')
print('本次执行检索到 ' + str(len(jpglist)) + ' 个jpg文件\n')
file_list = list()
# 解析转向值
for jpgpath in jpglist:
print(jpgpath)
curDataDir = os.path.dirname(jpgpath)
basename = os.path.basename(jpgpath)
angle = (basename[:-4]).split('_')[-1]
imgPath = os.path.join(curDataDir, basename).replace("\\", "/")
file_list.append((imgPath, angle))
# 切分数据
random.seed(256)
random.shuffle(file_list)
train_num = int(len(file_list) * train_ratio)
train_list = file_list[0:train_num]
val_list = file_list[train_num:]
# 创建列表文件
creat_data_list(org_img_folder, train_list, mode='train')
creat_data_list(org_img_folder, val_list, mode='val')
if __name__ == "__main__":
'''
程序入口
'''
main()上述代码我们查找每个log文件夹下的jpg文件，然后解析出对应的转向值。将这些值最后分别保存到train.txt和val.txt文件中。在代码里面，我们设定训练集占比0.8，剩下的0.2则为验证集。
生成的train.txt和val.txt文件每行内容表示一个样本，由图片路径和转向值组成，中间用空格隔开，如下所示：
./data/simulate/log3/1932_0.0667.jpg 0.0667到这里，我们就把基本数据准备工作做好了。
接下来我们将使用Pytoch框架实现深度学习算法进行训练、验证。
4.3 模型训练和验证
4.3.1 代码结构组织
为了方便读者阅读、运行和修改代码，本文深度学习部分将采用比较简单的代码组织方式。完整结构如下图所示：
项目根目录下有7个.py文件和3个文件夹，下面对各个文件和文件夹进行简单说明。
- create_data_lists.py：生成数据列表，解析图像路径和转向角，然后写入txt文件列表供后续PyTorch调用；
- datasets.py：用于构建数据集加载器，主要沿用Pytorch标准数据加载器格式进行封装；
- models.py：模型结构文件，存储自动驾驶模型的结构定义；
- utils.py：工具函数文件，所有项目中涉及到的一些自定义函数均放置在该文件中；
- train.py：用于训练算法；
- eval.py：用于模型评估；
- test.py：用于单张样本测试，运用训练好的模型为单张图像进行测试，查看角度转向是否正确；
- data：用于存放训练和测试数据集以及文件列表；
- results：用于存放运行结果，包括训练好的模型以及单张测试样本；
- runs：训练时，由tensorboard记录损失值变化；
- auto_drive:调用训练好的模型实现自动驾驶；
读者可以下载本文代码和数据集进行查看和运行，整个代码运行顺序如下：
（1）运行create_data_lists.py文件用于为数据集生成文件列表； （2）运行train.py进行算法训练，训练结束后在results文件夹中会生成checkpoint.pth模型文件； （3）运行eval.py文件对测试集进行评估，输出均方误差； （4）运行test.py文件对results文件夹下名为test.jpg的图像进行角度预测；
（5）运行auto_drive.py调用训练好的模型实现自动驾驶；
这里需要说明的是，本文不希望整个项目成为负担很重的“工程”，所以在代码编写上尽量手工化、傻瓜化，简洁化，方便读者后期自行魔改。
4.3.2 训练
首先定义数据采集器，位于datasets.py文件中，代码如下：
# 导入系统库 import os import numpy as np import cv2 # 导入PyTorch库 import torch from torch.utils.data import Dataset class AutoDriveDataset(Dataset): """ 数据集加载器 """ def __init__(self, data_folder, mode, transform=None): """ :参数 data_folder: # 数据文件所在文件夹根路径(train.txt和val.txt所在文件夹路径) :参数 mode: 'train' 或者 'val' :参数 normalize_type: 图像归一化处理方式 """ self.data_folder = data_folder self.mode = mode.lower() self.transform = transform assert self.mode in {'train', 'val'} # 读取图像列表路径 if self.mode == 'train': file_path=os.path.join(data_folder, 'train.txt') else: file_path=os.path.join(data_folder, 'val.txt') self.file_list=list() with open(file_path, 'r') as f: files = f.readlines() for file in files: if file.strip() is None: continue self.file_list.append([file.split(' ')[0],float(file.split(' ')[1])]) def __getitem__(self, i): """ :参数 i: 图像检索号 :返回: 返回第i个图像和标签 """ # 读取图像 img = cv2.imread(self.file_list[i][0]) img = cv2.cvtColor(img,cv2.COLOR_BGR2HSV) if self.transform: img = self.transform(img) # 读取标签 label = self.file_list[i][1] label = torch.from_numpy(np.array([label])).float() return img, label def __len__(self): """ 为了使用PyTorch的DataLoader,必须提供该方法. :返回: 加载的图像总数 """ return len(self.file_list)
上述代码比较简单，我们构造了AutoDriveDataset类用于作为自动驾驶小车数据读取类，从train.txt和val.txt中根据每行内容得到每个样本的图像路径和对应的真值标签。这里需要注意下颜色空间，我们最终是使用HSV空间进行训练的，因此需要做一下转化。
有了数据读取类以后我们就下来定义模型，具体代码如下：
# 导入PyTorch库
import torch.nn as nn
import torch.nn.functional as F
class AutoDriveNet(nn.Module):
'''
端到端自动驾驶模型
'''
def __init__(self):
"""
初始化
"""
super(AutoDriveNet, self).__init__()
self.conv_layers = nn.Sequential(nn.Conv2d(3, 24, 5, stride=2),
nn.ELU(),
nn.Conv2d(24, 36, 5, stride=2),
nn.ELU(),
nn.Conv2d(36, 48, 5, stride=2),
nn.ELU(), nn.Conv2d(48, 64, 3),
nn.ELU(), nn.Conv2d(64, 64, 3),
nn.Dropout(0.5))
self.linear_layers = nn.Sequential(
#nn.Linear(in_features=64 * 2 * 33, out_features=100),
nn.Linear(in_features=64 * 8 * 13, out_features=100),
nn.ELU(),
nn.Linear(in_features=100, out_features=50),
nn.ELU(),
nn.Linear(in_features=50, out_features=10),
nn.Linear(in_features=10, out_features=1))
def forward(self, input):
'''
前向推理
'''
input = input.view(input.size(0), 3, 120, 160)
output = self.conv_layers(input)
output = output.view(output.size(0), -1)
output = self.linear_layers(output)
return output这里需要注意的是我们的模型跟论文里的稍微有点不一样，主要是因为我们的图像尺寸是120x160的，而论文里使用的是66x200。因此，我们对应的输入需要调整下，另外，在最后全连接层也相应的在维度上要调整。对于实际项目来说，现在很多的摄像头都是使用3:4分辨率的，例如树莓派摄像头典型的分辨率是480x640，因此，修改过后的模型更具有普遍性，方便后面迁移到真实环境训练。
整个模型比较简单，前面是多个cnn，最后接几个全连接网络，输入是3通道图像，输出是一个转向回归值。
训练脚本代码train.py如下：
# 导入torch库
import torch.backends.cudnn as cudnn
import torch
from torch import nn
import torchvision.transforms as transforms
from torch.utils.tensorboard import SummaryWriter
# 导入自定义库
from models import AutoDriveNet
from datasets import AutoDriveDataset
from utils import *
def main():
"""
训练.
"""
# 数据集路径
data_folder = './data/simulate'
# 学习参数
checkpoint = None # 预训练模型路径，如果不存在则为None
batch_size = 400 # 批大小
start_epoch = 1 # 轮数起始位置
epochs = 1000 # 迭代轮数
lr = 1e-4 # 学习率
# 设备参数
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
ngpu = 4 # 用来运行的gpu数量
cudnn.benchmark = True # 对卷积进行加速
writer = SummaryWriter() # 实时监控 使用命令 tensorboard --logdir runs 进行查看
# 初始化模型
model = AutoDriveNet()
# 初始化优化器
optimizer = torch.optim.Adam(params=filter(lambda p: p.requires_grad,
model.parameters()),
lr=lr)
# 迁移至默认设备进行训练
model = model.to(device)
criterion = nn.MSELoss().to(device)
# 加载预训练模型
if checkpoint is not None:
checkpoint = torch.load(checkpoint)
start_epoch = checkpoint['epoch'] + 1
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
# 单机多卡训练
if torch.cuda.is_available():
model = nn.DataParallel(model, device_ids=list(range(ngpu)))
# 定制化的dataloader
transformations = transforms.Compose([
transforms.ToTensor(), # 通道置前并且将0-255RGB值映射至0-1
# transforms.Normalize(
# mean=[0.485, 0.456, 0.406], # 归一化至[-1,1] mean std 来自imagenet 计算
# std=[0.229, 0.224, 0.225])
])
train_dataset = AutoDriveDataset(data_folder,
mode='train',
transform=transformations)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=0,
pin_memory=True)
# 开始逐轮训练
for epoch in range(start_epoch, epochs + 1):
model.train() # 训练模式：允许使用批样本归一化
loss_epoch = AverageMeter() # 统计损失函数
n_iter = len(train_loader)
# 按批处理
for i, (imgs, labels) in enumerate(train_loader):
# 数据移至默认设备进行训练
imgs = imgs.to(device)
labels = labels.to(device)
# 前向传播
pre_labels = model(imgs)
# 计算损失
loss = criterion(pre_labels, labels)
# 后向传播
optimizer.zero_grad()
loss.backward()
# 更新模型
optimizer.step()
# 记录损失值
loss_epoch.update(loss.item(), imgs.size(0))
# 打印结果
print("第 " + str(i) + " 个batch训练结束")
# 手动释放内存
del imgs, labels, pre_labels
# 监控损失值变化
writer.add_scalar('MSE_Loss', loss_epoch.avg, epoch)
print('epoch:' + str(epoch) + ' MSE_Loss:' + str(loss_epoch.avg))
# 保存预训练模型
torch.save(
{
'epoch': epoch,
'model': model.module.state_dict(),
'optimizer': optimizer.state_dict()
}, 'results/checkpoint.pth')
# 训练结束关闭监控
writer.close()
if __name__ == '__main__':
'''
程序入口
'''
main()本文使用4卡进行训练（Nvidia T4）。在训练过程中我们使用tensorboard进行损失函数变化监控，可以使用下面的命令安装tensorboard：
pip install tensorboard在训练过程中，可以使用下面的命令启动tensorboard查看：
tensorboard --logdir=runs训练时batch_size设置为100，如果读者的显卡out of memory则该小这个batch_size。训练结果如下所示：
可以看到整个训练过程还是比较平稳正常的，在最终epoch=1000的时候基本处在一个比较好的收敛位置。
训练完成后的文件大小为9.20M，这个模型体量不大，适合嵌入式部署。
4.3.3 批量验证
使用eval.py文件进行精度验证：
# 导入系统库
import time
# 导入PyTorch库
import torch
from torch import nn
import torch.backends.cudnn as cudnn
import torchvision.transforms as transforms
# 导入自定义库
from datasets import AutoDriveDataset
from models import AutoDriveNet
from utils import *
def main():
# 测试集目录
data_folder = "./data/simulate"
# 定义运行的GPU数量
ngpu = 1
#cudnn.benchmark = True
# 定义设备运行环境
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 加载预训练模型
checkpoint = torch.load("./results/checkpoint.pth")
model = AutoDriveNet()
model = model.to(device)
model.load_state_dict(checkpoint['model'])
# 多GPU封装
if torch.cuda.is_available():
model = nn.DataParallel(model, device_ids=list(range(ngpu)))
# 定制化的dataloader
transformations = transforms.Compose([
transforms.ToTensor(), # 通道置前并且将0-255RGB值映射至0-1
# transforms.Normalize(
# mean=[0.485, 0.456, 0.406], # 归一化至[-1,1] mean std 来自imagenet 计算
# std=[0.229, 0.224, 0.225])
])
val_dataset = AutoDriveDataset(data_folder,
mode='val',
transform=transformations
)
val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=1,
pin_memory=True)
# 定义评估指标
criterion = nn.MSELoss().to(device)
# 记录误差值
MSEs = AverageMeter()
# 记录测试时间
model.eval()
start = time.time()
with torch.no_grad():
# 逐批样本进行推理计算
for i, (imgs, labels) in enumerate(val_loader):
# 数据移至默认设备进行推理
imgs = imgs.to(device)
labels = labels.to(device)
# 前向传播
pre_labels = model(imgs)
# 计算误差
loss = criterion(pre_labels, labels)
MSEs.update(loss.item(), imgs.size(0))
# 输出平均均方误差
print('MSE {mses.avg: .3f}'.format(mses=MSEs))
print('平均单张样本用时 {:.3f} 秒'.format((time.time()-start)/len(