使用Python去调用Yolo

  在之前使用yolo进行预测的时候,大多数通过命令行去执行

1
yolo detect predict model=yolov8n.engine source='0' device=0

  然后它就会打开摄像头,你就可以实时看到yolo处理过后的摄像头的画面,但是它背后是怎么去工作的呢?我之前去使用过yolov5,我的印象当中他有一个predict.py的文件,修改它,再执行就可以了。但是yolov8似乎并没有类似的现成的文件,但是封装性更强了,没有v5的什么一堆配置文件,支持命令行和API的直接调用。
  那就让我们着手去看,一步一步深挖它到底怎么从命令行去调用API最后实现预测的,想省事可以直接跳过,看怎么用API,但是我还是建议跟着看一遍,去打开对应的文件瞧瞧,理解能更深刻一点。
  在我们输入命令的时候,由于我使用的是Anaconda的虚拟环境,它会去环境当中寻找是否有这个命令,也就是是否有yolo这个命令,他一般会存在于:Anaconda/env/yolov8/bin/这个文件路径当中,其中env代表的是环境,这个文件夹下面会存放你创建的所有环境,yolov8是我自己创建的环境,你会在bin当中寻找到yolo这个文件(这里面装的是命令行入口脚本),打开它会有:

1
2
3
4
5
6
7
8
#!/home/djf/anaconda3/envs/yolov8/bin/python
# -*- coding: utf-8 -*-
import re
import sys
from ultralytics.cfg import entrypoint
if __name__ == "__main__":
sys.argv[0] = re.sub(r"(-script\.pyw?|\.exe)?$", "", sys.argv[0])
sys.exit(entrypoint())

  当我们使用命令yolo的时候,他会去执行这个文件,而我们也可以看到,这里面的entrypoint是关键点,我们也可以看到这个函数是从ultralytics.cfg当中引入的。
  ultralytics是python的一个包,是要使用yolov8的时候下的,让我们顺藤摸瓜摸过去。在Anaconda的虚拟环境下,安装的包一般都会装在:Anaconda/env/yolov8/lib/python3.10/site-package文件夹里面,找到ultralytics的文件夹,打开,再进去cfg文件夹,你会看到他的初始化文件,打开你就看得到entrypoint函数了,这个函数正是用来解析命令行输入的。
  在这里面的代码都可以看看,相关的函数或类都可以依照相同的方法一路摸过去。
  在entrypoint函数里面,最重要的是这几行代码,它告诉我们怎么去调用yolov8来实现检测,以及如何读取返回的结果用于进一步处理

1
2
3
4
from ultralytics import YOLO

model = YOLO(model, task=task) #这里面的model就是命令行当中输入的模型的名字,task就是detect

  以及

1
2
# Run command in python
getattr(model, mode)(**overrides) # default args from model mode就是predict overrides就是命令行中你输入的其他内容

  最后getattr这一段等效于(具体的还需要再去挖):

1
model.predict(source="0", device=0)

  而检测的值就是predict的返回值,返回的结构也可以顺藤摸瓜找出来:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
Args:
orig_img (numpy.ndarray): The original image as a numpy array.
path (str): The path to the image file.
names (Dict): A dictionary of class names.
boxes (torch.Tensor | None): A 2D tensor of bounding box coordinates for each detection.
masks (torch.Tensor | None): A 3D tensor of detection masks, where each mask is a binary image.
probs (torch.Tensor | None): A 1D tensor of probabilities of each class for classification task.
keypoints (torch.Tensor | None): A 2D tensor of keypoint coordinates for each detection.
obb (torch.Tensor | None): A 2D tensor of oriented bounding box coordinates for each detection.
speed (Dict | None): A dictionary containing preprocess, inference, and postprocess speeds (ms/image).

Examples:
>>> results = model("path/to/image.jpg")
>>> result = results[0] # Get the first result
>>> boxes = result.boxes # Get the boxes for the first result
>>> masks = result.masks # Get the masks for the first result

调用YOLOv8的api

  总结一下前面所发现的,可以知道要使用YOLOv8来进行预测,只需要:

1
2
3
from ultralytics import YOLO
model = YOLO(model=yolov8n.engine, task=detect)
result = model.predict(source='0', device=0, conf=0.5, show=True)

  你就得到预测结果啦
  但是在实际的使用当中,由于我在这个板子上用了Anaconda,ROS2官方好像说了两个不能混一块,所以正常使用会有报错。。。需要我们自行解决
  那就让我开始着手去写ROS2的yolov8的包吧!
  首先新建文件夹,这是你的ROS2工作空间,名字随你意,但是不要有中文以及中文符号,然后在文件夹当中再新建一个src文件夹,这是你存放源码的地方,举例:创建新文件夹:gimbal/scr
  然后在当前文件夹下用终端输入以下命令

1
ros2 pkg create yolov8_ros2 --build-type ament_python --dependencies rclpy ultralytics cv_bridge sensor_msgs vision_msgs cv2 usb_cam  --license Apache-2.0

  这行命令的意思是 用ros2命令的pkg子命令创建一个名为yolov8_ros2的ROS2的包,使用的是python,相关的依赖包是rclp··· ···,支持的开源协议是··· ···
  运行完这行命令之后,你会发现在src文件夹下多了一个yolov8_ros2的文件夹,这就是你刚刚创建的包的雏形。接下来在src/yolov8_ros2/yolov8_ros2文件夹下新建yolov8_node.py文件,这就是接收图像数据,并用yolov8进行检测,然后发布检测数据的节点。
  这个文件的内容如下,我直接粘贴过来:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
#file name:yolov8_node
import cv2
import time
from ament_index_python.packages import get_package_share_directory
import rclpy
from rclpy.node import Node
from cv_bridge import CvBridge
from sensor_msgs.msg import Image
from vision_msgs.msg import Detection2DArray, Detection2D, ObjectHypothesisWithPose, BoundingBox2D, Pose2D
import sys
sys.path.append(f'/home/djf/anaconda3/envs/yolov8/lib/python3.10/site-packages/')
from ultralytics import YOLO


class YOLOv8Detector(Node):
def __init__(self):
super().__init__('yolov8_detector')

# 初始化模型
package_share_directory = get_package_share_directory('yolov8_ros2')
self.model_path = package_share_directory + "/config/yolov8n.engine"
self.model = YOLO(self.model_path, task="detect")
self.bridge = CvBridge()

# 订阅图像话题
self.sub = self.create_subscription(
Image,
'/image_raw', # 根据实际话题修改
self.detect_callback,
10)

# 发布检测结果
self.pub = self.create_publisher(
Detection2DArray,
'/yolo/detections',
10)

self.get_logger().info("YOLOv8节点已启动")

def detect_callback(self, msg):
try:
# 图像格式转换
cv_image = self.bridge.imgmsg_to_cv2(msg, 'rgb8')
#cv_image = cv2.imencode('.jpg',cv_image)

# YOLO推理
start_time=time.time()
results = self.model.predict(
source=cv_image,
conf=0.3,
device=0)
end_time=time.time()
total_time=(end_time-start_time)*1000
self.get_logger().info(f'完成一次检测,用时:{total_time}')
# 构建检测消息
detections = Detection2DArray()
detections.header = msg.header

for box in results[0].boxes:
detection = Detection2D()
detection.bbox = BoundingBox2D()
detection.bbox.center = Pose2D()

# 计算边界框中心点

x1, y1, x2, y2 = box.xyxyn[0].tolist()
detection.bbox.center.position.x = (x1 + x2) / 2
detection.bbox.center.position.y = (y1 + y2) / 2
detection.bbox.size_x = x2 - x1
detection.bbox.size_y = y2 - y1
self.get_logger().info(f'中心x:{detection.bbox.center.position.x }')
self.get_logger().info(f'中心y:{detection.bbox.center.position.y }')

# 添加类别和置信度
hypothesis = ObjectHypothesisWithPose()
hypothesis.hypothesis.class_id = str(int(box.cls))
hypothesis.hypothesis.score = float(box.conf)
detection.results.append(hypothesis)

detections.detections.append(detection)
for result in results:
self.get_logger().info(f'发布检测结果:{result.boxes.data}')
self.get_logger().info(f'用时:{result.speed}')
self.pub.publish(detections)

except Exception as e:
self.get_logger().error(f"处理失败: {str(e)}")


def main(args=None):
rclpy.init(args=args)
node = YOLOv8Detector()
try:
rclpy.spin(node)
except KeyboardInterrupt:
pass
finally:
node.destroy_node()
rclpy.shutdown()


if __name__ == '__main__':
main()

  然后你也需要在src/yolov8_ros2文件路径下找到setup.py这个文件,修改一下里面的内容,这样ros才能知道你写了个节点,使用的时候才知道你写的节点在哪里,
  我把setup.py的文件内容也一并粘贴过来

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
from setuptools import find_packages, setup
import os
from glob import glob
package_name = 'yolov8_ros2'

setup(
name=package_name,
version='0.0.0',
packages=find_packages(exclude=['test']),
data_files=[
('share/ament_index/resource_index/packages',
['resource/' + package_name]),
(os.path.join('share', package_name, 'config'),glob('config/**')),
('share/' + package_name, ['package.xml']),
],
install_requires=['setuptools'],
zip_safe=True,
maintainer='djf',
maintainer_email='djf@todo.todo',
description='TODO: Package description',
license='TODO: License declaration',
tests_require=['pytest'],
entry_points={
'console_scripts': [
"yolov8_node=yolov8_ros2.yolov8_node:main"
],
},
)

  entry_point那边的修改就是告诉系统,当需要使用yolov8_node这个节点的时候,去执行yolov8_ros2/yolov8_node.py当中的main函数。
  除此之外,你还需要在src/yolv8_ros2/yolov8_ros2文件夹下新创建config目录,里面放下默认运行的yolov8的模型,我这里默认使用的是yolov8.engine,默认使用的模型文件名可以在yolov8_node.py文件当中修改。