二十. 在ROS系统上实现基于PyTorch YOLO v5的实时物体检测

一. 背景介绍

    在我前面的博文 十八.在JetsonNano上为基于PyTorch的物体检测网络测速和选型 中,我介绍过在基于Jetson Nano硬件平台和Ubuntu 18.04 with JetPack4.3, CUDA 10.2软件平台测试各种物体检测网络的检测速度的结果, 最终根据测试结果,以及各物体检测网络的性能和精准度等因素,我选择使用基于PyTorch 的YOLO v5. 

    但那个过程只是在开发环境那中,安装Python3,安装PyTorch,执行YOLO v5, 用coco数据集的静态图片跑了跑YOLO v5的过程. 这里我将实现YOLO v5和ROS系统的对接, 获取ROS实时视频图像帧,发送到YOLO v5网络模型进行物体检测,并将检测结果实时显示在处理后的视频图像帧上, 最终实现端到端的实时物体检测. (对应ROS功能包实现源码,见文末)

二. 准备工作

    除最基本的硬件和软件开发环境外,我们需要做:

1. 在基于Jetson nano硬件的Ubuntu 18.04上安装PyTorch,这主要包括torch和torchvision包,其他包PyTorch可选:

    这个安装过程我在 十八.在JetsonNano上为基于PyTorch的物体检测网络测速和选型 中已经介绍过这里不再阐述.
   

2. 可选.由于PyTorch基于Python3, 而ROS的Melodic版本, 其很多支持包仍然是基于Python2.7的,所以我们有个工作,让ROS支持Python3代码. 这个并不难一般安装Python3(我安装的是Python3.6)后,并在*.py文件头部中添加"#!/usr/bin/env python3"即可.

3. 也是较难处理的一步, 和图像打交道必然会用到Open CV,ROS系统也不例外,其处理Image类型消息和cv2之间的类型转换使用cv_bridge库, 但ROS Melodic发型的还是Python2.7的库, 在我们的Python3代码中无法正常调用(core dumped). 所以我买了需要下载cv_bridge源码自己编译Python3版本的cv_bridge.

    具体编译过程,网上有很多,但都是较老的文章,介绍的都是介于catkin编译方法,我试过均失败. **推荐大家参考另一篇基于最新catkin_make编译的方法,我也是基于此文编译成功的, 看这里: ** ROS——在Ubuntu18.04下基于ROS Melodic编译python3的cv_bridge
   

编译过程我也遇到自己的困难,并发表与上面博文的评论里和博主互动.

附上当时的交流截图:

1. 准备自己的物体检测网络代码. 这个需要自己处理;

三.详细过程与思路

    同样准备工作做完, 我们先要捋清楚为了实现我们的目标还需要哪些更细节的步骤, 先想清楚再动手做.  承接上一篇博文(十九. )这里给出一张流程图:

1.订阅ROS消息获取相机Image数据, 这个过程是标准ROS过程,和上篇博文不同的是,上篇博文我采用C++编写,这里我采用Python3编写,大致过程如下:

#网络检测网络初始化
self.detector = Yolov5Detector()
self.detector.detectorInit()
self.detector.detectorWarmUp()

#必要的发布和订阅
self.pubDetectedImage = rospy.Publisher('camera/image_raw/detect', Image, queue_size=1)
self.subCameraInfo = rospy.Subscriber("/camera/camera_info", CameraInfo, callback=self.onCameraInfoCallback, callback_args=None, queue_size=1)
self.subCameraImage = rospy.Subscriber("camera/image_raw", Image, callback=self.onCameraImageCallback, callback_args=None, queue_size=1)

2. ROS消息类型转换为CV2,并处理为物体网络支持类型,然后送入网络

def onCameraImageCallback(self,imageData):
    try:
        cvImage = self.bridge.imgmsg_to_cv2(imageData, "bgr8")
        
        # 图像去畸变, 使用相机内参和畸变系数可以图像去畸变
        #self.current_image_frame = cv2.undistort(cvImage, self.camera_intrinsic_value, self.distortion_coefficients, None, self.camera_intrinsic_value)
        self.current_image_frame = cvImage

        
        image_np = np.swapaxes(self.current_image_frame,0,2)  # [H,W,C] -> [C,W,H]
        image_np = np.swapaxes(image_np,1,2)  #  [C,W,H] -> [C,H,W]
        
        cvImgRet, detect, imageLabelList = self.detector.detectImage(image_np,self.current_image_frame)
    except Exception as e:
        print(e)

3. 对输出结果做cv2转ROS消息处理 并发布检测结果消息

image = self.bridge.cv2_to_imgmsg(cvImgRet, "bgr8")

self.pubDetectedImage.publish(image) #发布检测结果图像帧

self.pubDetectedResult.publish(detectedResult) #发布边框信息,得分,类别

4. 详细的YOLO v5过程不在阐述.

四.实现效果

    ![](https://img-blog.csdnimg.cn/53fd405047b944d698349381ada5843f.png?x-oss-process=image/watermark,type_d3F5LXplbmhlaQ,shadow_50,text_Q1NETiBAb2tnd2Y=,size_20,color_FFFFFF,t_70,g_se,x_16)

最后,应很多网友要求,我上传了这个ROS功能包的源码资源: 基于YOLO v5的物体检测ROS功能包实现