/SLAM_YOLOv5

Primary LanguageC++

SLAM-YOLOv5

做物体SLAM的时候本来想找到一个和SLAM结合的，并且是c++版本的目标检测系统，github上面一搜，不是很多，大部分都是离线检测，再不然就是ROS版本或者Python版本的YOLO（不是说不行，但是个人感觉速度很奇怪，速度也不快），满足我要求的几乎没有，因此干脆花了点时间自己东拼西凑一个出来。

代码链接：https://github.com/ChenJiahao031008/SLAM-YOLOv5.git

环境配置

由于深度学习不可避免的涉及到显卡安装等各类问题，这里把我自己的配置罗列下（安装过程没啥困难，也比较基础，就不展开了），因为自己刚配的电脑，比较新，大家可以适当放低配置什么的。

  • 系统 Ubuntu16.04

  • RTX 3060 Laptop显卡 Driver 470.42.01

  • CUDA 11.1（显卡太新以至于11以下的不兼容）

  • cudnn 8.1.0.77

  • OpenCV 3.4.4 （无cuda版本，因为不支持cuda11）

  • LibTorch： libtorch-cxx11-abi-shared-with-deps-1.9.0+cu111

  • （可选）pytorch ：1.9.0-cu111版本

运行实例

注意，运行前修改CMakeLists.txt文件中libtorch库的位置（或者CMAKE时候指定），修改YOLOv5Detector.cc中的硬地址等

mkdir build && cd build
cmake ..
make -j10
./Examples/RGB-D/rgbd_tum  Vocabulary/ORBvoc.txt Examples/RGB-D/TUM1.yaml  Dataset/rgbd_dataset_freiburg1_desk Dataset/rgbd_dataset_freiburg1_desk/associate.txt

改进思路

YOLO这块主要参考了https://github.com/Nebula4869/YOLOv5-LibTorch。

核心**比较简单，就是把Python版本的训练好的.pt文件转换ONNX推理框架，生成TorchScript文件，这里主要参考了代码：https://github.com/ultralytics/yolov5/tree/v5.0

运行指令：

export PYTHONPATH="$PWD"
# 在有pytorch的环境下 
python models/export.py --device 0

其中，需要下载yolov5s.pt在主文件夹下（不放的话也会自动下），在工程目录下生成yolov5s.torchscript.pt，然后讲生成的TorchScript文件放到SLAM框架下使用。

在SLAM中，单线程加入YOLO是比较简单的，只需要对CMakeLists文件做一些改动即可。注意，进行推理可能用到C++14标准，需要切换标志符。同时，这也会对SLAM系统产生一点影响（不过问题不大）

提示：关闭警告可以在CMakeLists中加入：

#添加的部分，关闭警告
add_definitions(-w)

多线程配置

为了提高性能，可以对YOLO另外开一个线程，这块涉及到了ORB-SLAM2的多线程系统设计，简单说一下：

多线程入口，仿照其他线程设计即可。看代码

#ifdef USE_YOLO_DETECTOR
    std::cout << "[INFO] USE_YOLO_DETECTOR." << std::endl;
    mpDetector = new YOLOv5Detector(mpFrameDrawer);
    mptDetector = new thread(&ORB_SLAM2::YOLOv5Detector::Run, mpDetector);
#endif

注意，这里仅仅是把YOLOv5嵌入了SLAM系统，在viewer线程中展示出来，并没有做任何额外的处理。

因此多线程的配合是比较简单的，YOLOv5Detector类只和FrameDrawer交互，从后者中读取需要处理的图像，并传回结果图像（中间过程有一些互锁环节）。而FrameDrawer会在viewer线程中被调用和结果展示。

如果大家感兴趣可以继续扩展，就写到这里了。