Ubuntu配置Yolov8环境并训练自己的数据集 + ROS实时运行
文章目录
前言：需要先安装CUDA和Anaconda，它们的安装参考我这篇文章：Ubuntu配置深度学习环境（TensorFlow和PyTorch）
一、环境配置与功能测试
1.1 安装
新建一个虚拟环境下安装：
# 新建虚拟环境
conda create -n yolov8 python=3.8
# 激活虚拟环境
conda activate yolov8
pip install ultralytics
# 使用清华大学的镜像源安装
pip install ultralytics -i https://pypi.tuna.tsinghua.edu.cn/simple/
源码安装：
# 激活虚拟环境
conda activate yolov8
# 需要单独安装torch
conda install pytorch==2.0.0 torchvision==0.15.0 torchaudio==2.0.0 pytorch-cuda=11.7 -c pytorch -c nvidia
git clone https://github.com/ultralytics/ultralytics.git
cd ultralytics
# 安装依赖
pip install -r requirements.txt
# 使用清华大学的镜像源安装
pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple/
1.2 目标检测
#激活虚拟环境
conda activate yolov8
#官方的测试案例进行程序的推理测试：
yolo task=detect mode=predict model=yolov8n.pt source=/home/zard/Pictures/2.jpeg device=cpu save=True show=True
# 任务模型task=detect，YOLOv8可用于检测，分割，姿态和分类
# 会自动下载权重文件https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n.pt到当前目录
# 推理的数据为source
# 这是CPU进行测试的，将device改为0用GPU
结果如下，结果保存在当前路径下的runs/detect/predict文件夹中：
Ultralytics YOLOv8.0.145 🚀 Python-3.7.16 torch-1.13.1+cu117 CPU (12th Gen Intel Core(TM) i5-12500H)
YOLOv8n summary (fused): 168 layers, 3151904 parameters, 0 gradients
image 1/1 /home/zard/Pictures/2.jpeg: 480x640 4 persons, 5 cars, 1 motorcycle, 1 suitcase, 32.9ms
Speed: 1.8ms preprocess, 32.9ms inference, 0.7ms postprocess per image at shape (1, 3, 480, 640)
Results saved to runs/detect/predict
GPU运行
yolo task=detect mode=predict model=yolov8n.pt source=/home/zard/Pictures/1.jpeg device=0 save=True show=True
Ultralytics YOLOv8.0.145 🚀 Python-3.7.16 torch-1.13.1+cu117 CUDA:0 (NVIDIA GeForce RTX 3060 Laptop GPU, 5930MiB)
YOLOv8n summary (fused): 168 layers, 3151904 parameters, 0 gradients
image 1/1 /home/zard/Pictures/1.jpeg: 448x640 7 persons, 6 cars, 2 buss, 6.8ms
Speed: 1.5ms preprocess, 6.8ms inference, 0.8ms postprocess per image at shape (1, 3, 448, 640)
Results saved to runs/detect/predict2
也可以输入文件夹，处理多张图片：
yolo task=detect mode=predict model=yolov8n.pt source=/home/zard/Pictures/ device=0 save=True show=True
Ultralytics YOLOv8.0.145 🚀 Python-3.7.16 torch-1.13.1+cu117 CUDA:0 (NVIDIA GeForce RTX 3060 Laptop GPU, 5930MiB)
YOLOv8n summary (fused): 168 layers, 3151904 parameters, 0 gradients
image 1/3 /home/zard/Pictures/1.jpeg: 448x640 7 persons, 6 cars, 2 buss, 7.0ms
image 2/3 /home/zard/Pictures/2.jpeg: 480x640 4 persons, 5 cars, 1 motorcycle, 1 suitcase, 6.9ms
image 3/3 /home/zard/Pictures/3.jpeg: 448x640 7 persons, 1 car, 2 trucks, 3.7ms
Speed: 1.8ms preprocess, 5.9ms inference, 0.6ms postprocess per image at shape (1, 3, 448, 640)
Results saved to runs/detect/predict3
1.3 实例分割
#激活虚拟环境
conda activate yolov8
yolo task=segment mode=predict model=/home/zard/yolov8s-seg.pt source=/home/zard/Pictures/2.jpeg device=0 save=True show=True
# 自动下载权重文件https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8s-seg.pt
1.4 分类
yolo task=classify mode=predict model=yolov8x-cls.pt source=/home/zard/Pictures/2.jpeg device=0 save=True show=True
1.5 姿态检测
yolo task=pose mode=predict model=/home/zard/yolov8s-pose.pt source=/home/zard/Pictures/2.jpeg device=0 save=True show=True
二、训练数据标注
安装数据标注环境：
#激活虚拟环境
conda activate yolov8
pip install labelImg
# 使用清华大学的镜像源安装
pip install labelImg -i https://pypi.tuna.tsinghua.edu.cn/simple/
labelImg
labelImg允许用户在图像中绘制边界框、多边形、线条和点等来标注不同类型的对象或特征。也可以标注标注类别，用户可以定义不同的标注类别，使其适应不同的项目需求。每个类别都可以有自己的名称和颜色。我这里从网上随便找了一些狗子的照片示例：
选择要标注的数据和输出目录，右下角会出现目录下所有图片：
更改数据的标注结果格式，点击save下面按钮选择输出为YOLO
点击Creat RectBox框选目标，填入标签：
每一帧标注完成后均要点击save，会在输出路径下生成结果，其中.txt文件包含图片中的物体信息及对应的位置之类的，class.txt列出所有的标签：
三、数据集训练方法
标注完成后，创建data文件夹，并在images中放置所有标注原图，labels保存每个图片对应的标签文件（.txt），dataSet用来保存数据集划分，例如训练集、验证集和测试集，通常以文本文件的形式列出每个数据集中的图像名称或ID，使用下面脚本生成数据集划分：
#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os
import random
import argparse
parser = argparse.ArgumentParser()
# 标注文件的地址，根据自己的数据进行修改
parser.add_argument('--label_path', default='./labels', type=str, help='input label path')
# 数据集的划分，地址选择自己数据dataSet下
parser.add_argument('--txt_path', default='./dataSet', type=str, help='output dataset path')
opt = parser.parse_args()
# 训练与测试数据集比例
trainval_percent = 1.0
train_percent = 0.9
labelfilepath = opt.label_path
txtsavepath = opt.txt_path
# 读取所有已经标注文件的名称
total_label = os.listdir(labelfilepath)
if not os.path.exists(txtsavepath):
os.makedirs(txtsavepath)
num = len(total_label)
list_index = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list_index, tv)
train = random.sample(trainval, tr)
file_trainval = open(txtsavepath + '/trainval.txt', 'w')
file_test = open(txtsavepath + '/test.txt', 'w')
file_train = open(txtsavepath + '/train.txt', 'w')
file_val = open(txtsavepath + '/val.txt', 'w')
for i in list_index:
name = total_label[i][:-4] + '\n'
# 排除掉生成的classes.txt文件
if name=='classes' + '\n':
continue
if i in trainval:
file_trainval.write(name)
if i in train:
file_train.write(name)
else:
file_val.write(name)
else:
file_test.write(name)
file_trainval.close()
file_train.close()
file_val.close()
file_test.close()
再用下面脚本将数据组织成Yolov8需要的形式（训练与验证数据集下均包含图像与对应的标签文件夹）：
#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os,shutil
rootpath="/home/zard/Pictures/data/"#待修改路径
# 输出路径
imgtrain=rootpath+"train/images/"
imgval=rootpath+"val/images/"
labeltrain=rootpath+"train/labels/"
labelval=rootpath+"val/labels/"
if not os.path.exists(imgtrain):
os.makedirs(imgtrain)
if not os.path.exists(imgval):
os.makedirs(imgval)
if not os.path.exists(labeltrain):
os.makedirs(labeltrain)
if not os.path.exists(labelval):
os.makedirs(labelval)
f = open(rootpath+"dataSet/train.txt","r")
lines = f.readlines()
for i in lines:
shutil.copy(rootpath+"images/"+str(i).replace('\n','')+".jpg",imgtrain+str(i).replace('\n','')+".jpg")
shutil.copy(rootpath + "labels/" + str(i).replace('\n', '') + ".txt", labeltrain + str(i).replace('\n', '') + ".txt")
f = open(rootpath+"dataSet/val.txt","r")
lines = f.readlines()
for i in lines:
shutil.copy(rootpath+"images/"+str(i).replace('\n','')+".jpg",imgval+str(i).replace('\n','')+".jpg")
shutil.copy(rootpath + "labels/" + str(i).replace('\n', '') + ".txt", labelval + str(i).replace('\n', '') + ".txt")
shutil.copy(rootpath+"dataSet/train.txt",rootpath+"train.txt")
shutil.copy(rootpath+"dataSet/trainval.txt",rootpath+"trainval.txt")
shutil.copy(rootpath+"dataSet/test.txt",rootpath+"test.txt")
shutil.copy(rootpath+"dataSet/val.txt",rootpath+"val.txt")
然后编写ymal文件，${Youpath}替换为你的路径：
train: ${Youpath}/data/train/images
val: ${Youpath}/data/val/images
test: ${Youpath}/data/test/images
# number of classes
nc: 1
# class names
names: ['dog']
然后去Yolo源码里找到yolov8.yaml参数文件，复制一份并修改：
nc: 1 # number of classes
最后的文件结构如下，蓝色文件夹下均为图像和对应的标注文件：
3.1 命令训练
#激活虚拟环境
conda activate yolov8
cd ${Youpath}/data
yolo task=detect mode=train model=yolov8.yaml data=dog.yaml batch=32 epochs=100 imgsz=640 workers=16 device=0
model部分可以进行模型的选择与更改，其余的参数都可以根据自己电脑的性能进行修改。训练完成后就可以进入runs文件夹下面看自己的训练成果：
结果如下，包括一些图表和验证结果：其中weights下包含两个训练的权重文件（最好的 和 上一次的）
利用训练得到的权重推理：
#激活虚拟环境
conda activate yolov8
yolo task=detect mode=predict model=best.pt source=15.jpg device=cpu save=True show=True
由于就只有十几张图片，没有训练出来效果就不展示了
best.pt还可以作为下一次训练的初值：
#激活虚拟环境
conda activate yolov8
yolo task=detect mode=predict model= ./runs/detect/train/weights/best.pt source=15.jpg device=cpu save=True show=True
3.2 代码训练
训练：
from ultralytics import YOLO
# Load a model
# model = YOLO('yolov8n.pt') # load a pretrained model (recommended for training)
# model = YOLO('./runs/detect/train/weights/best.pt') # load a pretrained model (recommended for training)
model = YOLO('yolov8.yaml') # load a pretrained model (recommended for training)
# Train the model
model.train(data='dog.yaml', epochs=100, imgsz=640，workers=16 device=0)
运行，与上面效果是一样的：
#激活虚拟环境
conda activate yolov8
python3 train.py
推理：
from ultralytics import YOLO
from PIL import Image
import cv2
model = YOLO("runs/detect/train/weights/best.pt")
# accepts all fonmats - image/dir/Path/URL/video/PIL/ndarray. 0 for webcamresults = model.predict(source="0")
# results = model.predict(source="0") # 用摄像头
# results = model.predict(source="folder",show=True)# Display preds. Accepts all YoLO predict argument
#from PIL
im1 = Image.open("15.jpg")
results = model.predict(source=im1, save=True) # save plotted images
#from ndarray
# im2 = cv2.imread("test.jpg")
# results = model.predict(source=im2,save=True,save_txt=True) # save predictions as labels
# #from list of PIL/ndancay
# results = model. predict(source=[im1, im2])
运行：
#激活虚拟环境
conda activate yolov8
python3 train.py
四、Yolov8在ROS中实时运行
4.1 yolov8_ros
mkdir yolov8_ros/src -p
cd yolov8_ros/src
git clone https://github.com/qq44642754a/Yolov8_ros.git
cd ..
catkin_make
使用TUM数据集测试，图像话题参数修改如下：
<param name="image_topic" value="/camera/rgb/image_color" />
rosbag play rgbd_dataset_freiburg3_walking_rpy.bag
roslaunch yolov8_ros yolo_v8.launch
如果报错：ImportError: No module named cv2，说明ros调用了python2，修改：
#!/usr/bin/env python3
-- coding: utf-8 --
4.2 ultralytics_ros
这个相对功能更全面一点
cd yolov8_ros/src
git clone -b noetic-devel https://github.com/Alpaca-zip/ultralytics_ros.git
git clone -b noetic-devel https://github.com/ros-perception/vision_msgs.git
python3 -m pip install -r ultralytics_ros/requirements.txt
cd yolov8_ros
rosdep install -r -y -i --from-paths .
catkin build
使用TUM数据集测试，参数修改如下：
<arg name="input_topic" default="/camera/rgb/image_color"/>
<arg name="device" default="0"/>
rosbag play rgbd_dataset_freiburg3_walking_rpy.bag
roslaunch ultralytics_ros tracker.launch
代码阅读（yolo跟踪或者分割的数据获取）
#!/usr/bin/env python3
# 指定脚本使用的Python解释器
import cv_bridge # 导入cv_bridge模块，用于ROS图像消息和OpenCV图像格式之间的转换
import numpy as np
import roslib.packages
import rospy
from sensor_msgs.msg import Image
from ultralytics import YOLO
from vision_msgs.msg import Detection2D, Detection2DArray, ObjectHypothesisWithPose
from ultralytics_ros.msg import YoloResult
class TrackerNode:
def __init__(self):
# 初始化节点的参数
# 指定接收图像数据的ROS话题
self.input_topic = rospy.get_param("~input_topic", "image_raw")
# 将检测结果发布到的ROS话题
self.result_topic = rospy.get_param("~result_topic", "yolo_result")
# 结果图像发布的ROS话题名称
self.result_image_topic = rospy.get_param("~result_image_topic", "yolo_image")
# 指定要使用的YOLO模型
yolo_model = rospy.get_param("~yolo_model", "yolov8n.pt")
# 置信度阈值，默认值为0.25。用于过滤检测结果的置信度阈值，低于该值的结果会被丢弃。
self.conf_thres = rospy.get_param("~conf_thres", 0.25)
# IoU（交并比）阈值，默认值为0.45。表示两个边界框之间的最小IoU阈值，用于非极大值抑制（NMS）。
self.iou_thres = rospy.get_param("~iou_thres", 0.45)
# 最大检测数，默认值为300。指定每张图像上允许的最大检测数量。
self.max_det = rospy.get_param("~max_det", 300)
# 需要检测的类别，默认为None。可以指定要检测的特定类别，如果为None，则检测所有类别。
self.classes = rospy.get_param("~classes", None)
# 跟踪器的配置文件路径或名称，默认值为"bytetrack.yaml"。用于指定目标跟踪器的配置。
self.tracker = rospy.get_param("~tracker", "bytetrack.yaml")
# 设备类型，默认为None。指定模型的运行显卡设备，例如"cpu"或"cuda:0"等。
self.device = rospy.get_param("~device", None)
# result_conf：是否在结果图像中显示置信度，默认为True。控制是否在输出图像中显示检测框的置信度。
# result_line_width：结果图像中检测框的线宽，默认为None。指定在结果图像中绘制检测框的线宽。
# result_font_size：结果图像中文本的字体大小，默认为None。指定在结果图像中显示的文本的字体大小。
# result_font：结果图像中文本的字体名称，默认为"Arial.ttf"。指定在结果图像中显示的文本的字体。
# result_labels：是否在结果图像中显示标签，默认为True。控制是否在输出图像中显示类别标签。
# result_boxes：是否在结果图像中显示边界框，默认为True。控制是否在输出图像中显示检测框。
self.result_conf = rospy.get_param("~result_conf", True)
self.result_line_width = rospy.get_param("~result_line_width", None)
self.result_font_size = rospy.get_param("~result_font_size", None)
self.result_font = rospy.get_param("~result_font", "Arial.ttf")
self.result_labels = rospy.get_param("~result_labels", True)
self.result_boxes = rospy.get_param("~result_boxes", True)
# 加载YOLO模型
path = roslib.packages.get_pkg_dir("ultralytics_ros") # 获取ultralytics_ros软件包的路径
self.model = YOLO(f"{path}/models/{yolo_model}") # 加载YOLO模型
self.model.fuse() # 加载模型权重
# 创建ROS订阅器和发布器
self.sub = rospy.Subscriber(
self.input_topic,
Image,
self.image_callback,
queue_size=1,
buff_size=2**24,
)
self.results_pub = rospy.Publisher(self.result_topic, YoloResult, queue_size=1)
self.result_image_pub = rospy.Publisher(self.result_image_topic, Image, queue_size=1)
self.bridge = cv_bridge.CvBridge()
# 检查是否使用分割模型
self.use_segmentation = yolo_model.endswith("-seg.pt")
# 图像回调函数，用于接收传感器图像消息
def image_callback(self, msg):
cv_image = self.bridge.imgmsg_to_cv2(msg, desired_encoding="bgr8")
# 使用YOLO模型进行目标跟踪
results = self.model.track(
source=cv_image,
conf=self.conf_thres,
iou=self.iou_thres,
max_det=self.max_det,
classes=self.classes,
tracker=self.tracker,
device=self.device,
verbose=False,
retina_masks=True,
)
# 如果有检测结果
if results is not None:
yolo_result_image_msg = Image() # 创建图像消息
yolo_result_image_msg.header = msg.header # 设置图像消息头部
yolo_result_image_msg = self.create_result_image(results) # 创建结果图像消息
self.result_image_pub.publish(yolo_result_image_msg) # 发布结果图像消息
# YoloResult消息类型包含detections和掩膜图像
# header
# detections:
# heade
# detections: "<array type: vision_msgs/Detection2D, length: 5>"
# masks: "<array type: sensor_msgs/Image, length: 5>"
yolo_result_msg = YoloResult() # 创建YoloResult消息
yolo_result_msg.header = msg.header # 设置消息头部
yolo_result_msg.detections = self.create_detections_array(results) # 创建检测结果的消息数组
if self.use_segmentation: # 如果使用分割模型
yolo_result_msg.masks = self.create_segmentation_masks(results) # 创建分割掩模消息
self.results_pub.publish(yolo_result_msg) # 发布检测结果消息
# 创建检测结果数组消息
def create_detections_array(self, results):
detections_msg = Detection2DArray()
# 获取坐标、类别id和相似度 [0]是因为结果是一个列表，在处理多张图片时用
# 边界框的坐标格式为 [x, y, w, h]，其中 x 和 y 是边界框左上角的坐标，w 和 h 分别表示边界框的宽度和高度
bounding_box = results[0].boxes.xywh
classes = results[0].boxes.cls
confidence_score = results[0].boxes.conf
for bbox, cls, conf in zip(bounding_box, classes, confidence_score):
detection = Detection2D()
detection.bbox.center.x = float(bbox[0])
detection.bbox.center.y = float(bbox[1])
detection.bbox.size_x = float(bbox[2])
detection.bbox.size_y = float(bbox[3])
hypothesis = ObjectHypothesisWithPose()
hypothesis.id = int(cls)
hypothesis.score = float(conf)
detection.results.append(hypothesis)
detections_msg.detections.append(detection)
return detections_msg
# 创建结果图像消息
def create_result_image(self, results):
# 在图像上绘制分割结果
plotted_image = results[0].plot(
conf=self.result_conf,
line_width=self.result_line_width,
font_size=self.result_font_size,
font=self.result_font,
labels=self.result_labels,
boxes=self.result_boxes,
)
result_image_msg = self.bridge.cv2_to_imgmsg(plotted_image, encoding="bgr8")
return result_image_msg
# 创建分割掩模消息
def create_segmentation_masks(self, results):
masks_msg = []
for result in results:
if hasattr(result, "masks") and result.masks is not None:
# 遍历 result.masks 中的每个分割掩模
for mask_tensor in result.masks:
# 将 PyTorch 张量（tensor）转换为 NumPy 数组（numpy array）
mask_numpy = (
np.squeeze(mask_tensor.data.to("cpu").detach().numpy()).astype(
np.uint8
)
* 255
)
# 将 NumPy 数组转换为 ROS 图像消息
mask_image_msg = self.bridge.cv2_to_imgmsg(
mask_numpy, encoding="mono8"
)
masks_msg.append(mask_image_msg)
return masks_msg
if __name__ == "__main__":
rospy.init_node("tracker_node") # 初始化ROS节点
node = TrackerNode() # 实例化TrackerNode类
rospy.spin() # 进入ROS事件循环