【ROS】python rospy定义一个节点实现图像目标检测结果的发布和在Rviz上的实时显示

写在前面

  之前，我们调用了torch中的yolov5模型，把机器人上相机采集的RGB图像放入模型中去做目标检测，最近想实现把检测结果能够实时在Rviz上做一个展示的效果。其实思路很简单，就是自己定义一个Publisher把目标检测的结果放进去，然后在Rviz里针对它去做一个订阅，实现效果放在文章最后了。接下来，说一下具体的解决方法。

1. 明确要发布的消息类型

  参考Camera的数据类型，我们需要发布的应该是一个Image类型的msg。Image数据类型的具体结构查阅可以参考作者上一篇文章，这里也还是放一个我读到的Camera的一个Image的msg：

  对于Image类型的数据来说，我们不需要去先new一个类，再自己去写header和定义宽、高之类的，这里我们只需要给出data数据，我们可以用cv_bridge.cv2_to_imgmsg(img, "bgr8")去完成cv2读取的图像格式向ros的图像格式的转换，具体就是：

from cv_bridge import CvBridge

cv_bridge = CvBridge()
img_msg = cv_bridge.cv2_to_imgmsg(img, "bgr8")

  更详细的也可以参考下面的链接。https://blog.csdn.net/gezongbo/article/details/122706911

2. YOLOv5模型的调用和结果返回

  关于YOLOv5模型的调用，Pytorch官方的hub里面给出了很多模型，我们这里用的是YOLOv5x，使用方法也非常简单，具体步骤可以参考下面两个链接：
  https://pytorch.org/hub/ultralytics_yolov5/
  https://github.com/ultralytics/yolov5/issues/36
  需要注意的是，如果直接将检测后得到的result返回，返回值应该是一个Detections类的实例；而我们需要返回的值是检测完之后的带检测框的图像，方法为调用该类中的render()方法，在第二个链接的Base64 Results部分也有使用，想深入了解的话也可以去yolov5模型的代码里去查看，应该在yolov5/models/common.py中。模型地址：
https://github.com/ultralytics/yolov5
简单来说，如果传入模型的只有一张图片，那么我们获取这张带检测结果的图像只需要：

import torch

# Model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
# Image
img = cv2.imread('test.png')
img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
# Inference
results = yolo(img)
# Get the Image with boundingbox
result_img = results.render()[0]

3. 如何写一个Publisher

  关于如何写一个发布器Publisher, 可以参考这一篇文章。https://blog.csdn.net/Mr_Poohhhh/article/details/103797043

大概就是：

class MyPublisher:
    def __init__(self):
        # 定义一个发布器
        self.__pub_ = rospy.Publisher('topic to be published',
        Pub msg type, queue_size)
	
    def my_public(self,data):
		# data为需要publish的数据	
        self.__pub_.publish(data)

def main():
	# 这里自己处理好好要发布的data
    data = Pub msg type() 
    publisher = MyPublisher()
    publisher.my_public(data)
if __name__ == "__main__":
    main()

4. 完整代码

#!/usr/bin/env python

import torch
import rospy
import numpy as np
import copy
import cv2
from cv_bridge import CvBridge
from sensor_msgs.msg import Image
import message_filters

class YoloRGB(object):
    def __init__(self):
        self.cv_bridge = CvBridge()
        # /locobot/yolo_result/img 是我们自己定义的节点名字
        self.yolo_img_pub = rospy.Publisher("/locobot/yolo_result/img", Image, queue_size=5)
    
    def img_pub(self, img=None):
        img_msg = self.cv_bridge.cv2_to_imgmsg(img, "bgr8")
        self.yolo_img_pub.publish(img_msg)

def main():
	# 输入图片（这里可以自己做替换）
	img = cv2.imread('test.png')
	img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
	
	# 检测部分
	model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
	results = yolo(img)
	result_img = results.render()[0]
	
	# 这里写一个while(true)是因为程序运行完成了之后这个publisher就订阅不到了
	# 具体效果自己把握
	yolo_pub = YoloRGB()
	while(true):
		# 发布imgmsg
		yolo_pub.img_pub(result_img)
    
if __name__ == "__main__":
    main()

5. 在Rviz订阅我们自己发布的话题

  在我们自己的程序跑起来之后，可以打开Rviz找到我们自己定义的话题名，然后选择加入到我们的界面中。具体步骤为：
  左边Displays窗口点Add按钮。

  切换到By topic下找到我们自己定义命名的topic。

  可以看到我们新加入的Image能够成功显示了。

  最后的效果就是这样啦。