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允许用户在图像中绘制边界框、多边形、线条和点等来标注不同类型的对象或特征。也可以标注标注类别，用户可以定义不同的标注类别，使其适应不同的项目需求。每个类别都可以有自己的名称和颜色。我这里自己拍了大约50张锥桶的照片示例：
选择要标注的数据和输出目录，右下角会出现目录下所有图片，更改数据的标注结果格式，点击save下面按钮选择输出为YOLO：

点击Creat RectBox框选目标，填入标签：

每一帧标注完成后均要点击save（下拉菜单可以选择自动保存），会在输出路径下生成结果，其中.txt文件包含图片中的物体信息及对应的位置之类的，class.txt列出所有的标签（类别ID从上到下依次增加：0，1，2，…）：

三、数据集训练方法

标注完成后，创建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}替换为你的路径，类别与class.txt对应：

train: ${Youpath}/data/train/images
val: ${Youpath}/data/val/images
test: ${Youpath}/data/test/images

# number of classes
nc: 1

# class names
names: ['Cone_bucket']

然后去Yolo源码里找到yolov8.yaml参数文件，复制一份并修改：

nc: 1  # number of classes

或者找到yolov8n.pt文件直接用

最后的文件结构如下，蓝色文件夹下均为图像和对应的标注文件：

3.1 命令训练

#激活虚拟环境
conda activate yolov8

cd ${Youpath}/data
yolo task=detect mode=train model=yolov8.yaml data=trafficlight.yaml batch=32 epochs=100 imgsz=640 workers=16 device=0
# 或者
yolo task=detect mode=train model=yolov8n.pt data=trafficlight.yaml batch=16 epochs=300 imgsz=640 workers=16 device=0

model部分可以进行模型的选择与更改，其余的参数都可以根据自己电脑的性能进行修改。训练完成后就可以进入runs文件夹下面看自己的训练成果：

结果如下，包括一些图表和验证结果：其中weights下包含两个训练的权重文件（最好的 和 上一次的）



利用训练得到的权重（last更好一点）推理：

#激活虚拟环境
conda activate yolov8
yolo task=detect mode=predict model=last.pt source=test.jpg  device=cpu save=True show=True

best.pt还可以作为下一次（result.png中损失值一直在变小，还没有稳定，则还有优化的空间，需要再次训练）训练的初值：

#激活虚拟环境
conda activate yolov8
yolo task=detect mode=train model=./runs/detect/train/weights/best.pt  data=trafficlight.yaml batch=32 epochs=300 imgsz=640 workers=16 device=0

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='trafficlight.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("test.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事件循环