写在前面
之前,我们调用了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
大概就是:文章来源:https://www.toymoban.com/news/detail-403518.html
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能够成功显示了。
最后的效果就是这样啦。
文章来源地址https://www.toymoban.com/news/detail-403518.html
到了这里,关于【ROS】python rospy定义一个节点实现图像目标检测结果的发布和在Rviz上的实时显示的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!