仿真环境中自动驾驶（CV视觉-转向盘角度）

 本文参考自2016年英伟达发表的论文《End to End Learning for Self-Driving Cars》

'''
本章内容:
1、阅读自动驾驶论文
2、数据采集
3、根据论文搭建自动驾驶神经网络
4、训练模型
5、在仿真环境中进行自动驾驶
'''

import traceback
from io import BytesIO
import argparse
import base64
import socketio
import torch
import torch.nn as nn
import cv2
import csv
import torch.backends.cudnn as cudnn
import os
import time
import numpy as np
import eventlet
from PIL.Image import Image
from torchvision.transforms import transforms
from torch.utils.tensorboard import SummaryWriter
from torch.utils.data.sampler import SubsetRandomSampler


# 数据采集

class AutoDriveDataset(torch.utils.data.Dataset):
    def __init__(self, csv_path, image_dir, transform=None):
        super(AutoDriveDataset, self).__init__()
        self.samples = []
        with open(csv_path) as f:
            reader = csv.reader(f)
            for line in reader:
                self.samples.append(line)
        # 图片文件夹路径
        self.image_dir = image_dir
        # 图像预处理
        self.transform = transform

    def __getitem__(self, index):
        example = self.samples[index]
        # 图片路径
        center = example[0]
        # 方向盘的角度
        angle = example[3]
        # 获取图片信息
        image = cv2.imread(center)
        # image = preprocess_bgr(image)

        if self.transform:
            image = self.transform(image)

        return image, float(angle)

    def __len__(self):
        return len(self.samples)


# 模型定义
class AutoDriveNet(nn.Module):
    def __len__(self):  # 原始图片320*120
        super(AutoDriveNet, self).__init__()
        # 神经网络 外部输入的数据为 3*66*200
        # 卷积核5x5 半径2: 66 - 2*2--> 62 -步长->31
        self.conv_layer = nn.Sequential(
            nn.Conv2d(3, 24, kernel_size=5, stride=2),
            nn.ReLU(),
            nn.Conv2d(24, 36, kernel_size=5, stride=2),
            nn.ReLU(),
            nn.Conv2d(36, 48, kernel_size=5, stride=2),
            nn.ReLU(),
            nn.Conv2d(48, 64, kernel_size=5, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=5, stride=2),
            nn.Dropout(p=0.5)  # 防止过拟合
        )

        # 线性层
        self.linear_layers = nn.Sequential(
            nn.Linear(64 * 1 * 18, 100),
            nn.ReLU(),
            nn.Linear(100, 50),
            nn.ReLU(),
            nn.Linear(50, 10),
            nn.Linear(10, 1)
        )

    def forward(self, x):
        # 正向传播/前向传播 n 3*66*200
        out = x.view(x.size(0), 3, 66, 200)
        out = self.conv_layer(out)
        out = out.view(out.size(0), -1)
        # 将数据打平
        out = self.linear_layers(out)
        return out


# 模型训练


class AverageLoss(object):
    '''计算每轮训练的平均loss'''

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


# 模型训练
class ModelTrain:
    def __init__(self, CSV_PATH=None, IMAGE_PATH=None):
        if not CSV_PATH:
            self.CSV_PATH = "d:/DATA/ml/driving/3/driving_log.csv"
        if not IMAGE_PATH:
            self.IMAGE_PATH = "d:/DATA/ml/driving/3/IMG"

    def train(self):
        transform = transforms.Compose([transforms.Lambda(lambda x: (x / 127.5) - 1.0)])
        # 数据集
        data_set = AutoDriveDataset(self.CSV_PATH, self.IMAGE_PATH, transform)
        total_size = len(data_set)
        print(f'数据总量:{total_size}')

        # 将数据进行打散
        # 获取所有的数据的索引
        indexes = list(range(total_size))
        print(indexes)
        np.random.seed(1234)
        np.random.shuffle(indexes)
        print(indexes)
        # 划分数据集，训练集0.8,测试集0.2
        splitRatio = 0.8
        trainSize = int(total_size * splitRatio)
        # 训练集的索引
        trainIndexes = indexes[0:trainSize]
        # 验证集的索引
        valIndexes = indexes[trainSize:]
        # 随机采样
        trainSampler = SubsetRandomSampler(trainIndexes)
        valSampler = SubsetRandomSampler(valIndexes)
        # 数据加载器
        trainDataLoader = torch.utils.data.DataLoader(data_set, batch_size=400, sampler=trainSampler)
        valDataLoader = torch.utils.data.DataLoader(data_set, batch_size=400, sampler=valSampler)

        writer = SummaryWriter()

        # 获取当前支持设备
        if torch.cuda.is_available():
            device = torch.device('cuda:0')
            cudnn.benchmark = True
        else:
            device = torch.device('cpu')

        # 为了接着上一次的训练结果继续训练
        if os.path.exists('models/best-{}.pt'):
            model = torch.load('models/best-.pt')
        else:
            model = AutoDriveNet()

        model.to(device)

        lossFunction = nn.MSELoss().to(device)
        optimzer = torch.optim.Adam(model.parameters(), lr=1e-4)
        best_loss = 10000

        for epoch in range(300):
            model.train()
            lossEpoch = AverageLoss()
            loss = None
            for inputs, labels in trainDataLoader:
                # 为了提高运行效率，将数据的计算移动到GPU中
                inputs = inputs.float().to(device)
                lables = lables.float().to(device)
                # 1.正向传播
                pred = model(inputs)
                # 2.计算损失
                loss = lossFunction(pred, lables, torch.unsqueeze(1))
                # 3.反向传播
                loss.backward()
                # 4.更新参数
                optimzer.step()
                # 5.梯度清空
                optimzer.zero_grad()
                # 记录损失值
                lossEpoch.update(loss.items(), inputs.size(0))
            # 保存最后的loss
            if loss.item() < best_loss:
                best_loss = loss.item()
                value = int(time.time())
                torch.save(model, f'models/best-{value}.pt')
            # 监控损失值的变化
            writer.add_scalar('MSE_LOSS', lossEpoch.avg, epoch)
            print(f'Epoch:{epoch},MSE_LOSS:{lossEpoch.avg}')
        print('train done!')
        torch.save(model, 'last.pt')
        writer.close()


if __name__ == '__main__':

        eventlet.monkey_patch(socket=True, select=True)
        sio = socketio.Server(async_mode='eventlet')
        app = socketio.WSGIApp(sio)
        model = None
        prev_image_array = None
        transformations = transforms.Compose([transforms.Lambda(lambda x:(x/127.5)-1.0)])
        IMAGE_HEIGHT, IMAGE_WIDTH, IMAGE_CHANNELS = 66, 200, 3


        def preprocess(image):
            image = image[60:-25, :, :]

            image = cv2.resize(image, (IMAGE_WIDTH, IMAGE_HEIGHT))
            image = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
            return image

        MAX_SPEED = 13
        MIN_SPEED = 3
        speed_limit = MAX_SPEED
        device = torch.device('cuda')

        @sio.on('telemetry')
        def telemetry(sid,data):
            if data:
                # 当前方向角度
                steering_angle=float(data['steering_angle'])
                # 当前的油门
                throttle = float(data['throttle'])
                # 当前的车速
                speed = float(data['speed'])
                # The current image from the center camera of the car
                image = Image.open(BytesIO(base64.b64decode(data['image'])))

                try:
                    with torch.no_grad():
                        image = np.asarray(image)
                        image = preprocess(image)
                        image = np.array([image])
                        image = transformations(image)
                        image = torch.Tensor(image).to(device)
                        # 预测方向盘角度
                        steering_angle = model(image).view(-1).cpu().data.numpy()
                        global  speed_limit
                        if speed>speed_limit:
                            speed_limit=MIN_SPEED
                        else:
                            speed_limit=MAX_SPEED

                        throttle=1.0-steering_angle**2-(speed/speed_limit)**2
                        print(f'{steering_angle} {throttle} {speed}')
                        send_control(steering_angle,throttle)
                except Exception as e:
                    print(traceback.format_exc())
            else:
                sio.emit('manul',data={},skip_sid=True)
        @sio.on('connect')
        def connect(sid,environ):
            print(f'connect: {sid}')
            send_control(0,0)

        def send_control(steering_angle,throttle):
            sio.emit('steer',
                     data={
                         "steering_angle":steering_angle.__str__(),
                         "throttle":throttle.__str__()
                     },skip_sid=True)
        '''
        pip install eventlet==0.29.1 --force-reinstall
        pip install python-socketio==4.5.1 --force-reinstall
        pip install python-engineio==3.11.2 --force-reinstall
        '''
        if __name__ == '__main__':
            model_path='best.pt'
            model = torch.load(model_path)
            model.eval()
            eventlet.wsgi.server(eventlet.listen(('',8090)),app)