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

写在前面

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

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

  参考Camera的数据类型,我们需要发布的应该是一个Image类型的msg。Image数据类型的具体结构查阅可以参考作者上一篇文章,这里也还是放一个我读到的Camera的一个Image的msg:
【ROS】python rospy定义一个节点实现图像目标检测结果的发布和在Rviz上的实时显示_第1张图片
  对于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按钮。
【ROS】python rospy定义一个节点实现图像目标检测结果的发布和在Rviz上的实时显示_第2张图片
  切换到By topic下找到我们自己定义命名的topic。
【ROS】python rospy定义一个节点实现图像目标检测结果的发布和在Rviz上的实时显示_第3张图片
  可以看到我们新加入的Image能够成功显示了。
【ROS】python rospy定义一个节点实现图像目标检测结果的发布和在Rviz上的实时显示_第4张图片
  最后的效果就是这样啦。

你可能感兴趣的:(python,目标检测,机器人)