YOLOv10训练自己的数据集

最新推荐文章于 2025-03-29 18:51:36 发布

_Mamba24

最新推荐文章于 2025-03-29 18:51:36 发布

阅读量7.4k

点赞数 27

分类专栏： YOLOv10 文章标签： YOLO 深度学习人工智能

本文链接：https://blog.csdn.net/qq_54134410/article/details/139243957

版权

YOLOv10 专栏收录该内容

2 篇文章

订阅专栏

3.1、设置官方配置文件：default.yaml，可自行修改。

0、引言

本文是使用YOLOv10训练自己的数据集，数据集包含COCO数据集的人猫狗数据以及自己制作的人猫狗数据集，类别为0:person、1:cat、2:dog三类，大家可根据自己的数据集类别进行调整。

1、环境配置

打开Anaconda3终端，进入base环境，创建新环境

conda create -n yolov10 python=3.9
conda activate yolov10
#cd到yolov10的目录下
pip install -r requirements.txt -i  https://pypi.tuna.tsinghua.edu.cn/simple/
pip install -e .

2、数据集准备

YOLOv10的训练数据集格式与YOLOv8相同

mydata
______images
____________train
_________________001.jpg
____________val
_________________002.jpg
______labels
____________train
_________________001.txt
____________val
_________________002.txt

参照这篇博客的数据集准备即可：

YOLOv8-Detect训练CoCo数据集+自己的数据集_yolov8训练coco-CSDN博客

3、创建配置文件

3.1、设置官方配置文件：default.yaml，可自行修改。

3.2、设置data.yaml

根据自己的数据集位置进行修改和配置。

path: D:\Yolov8\ultralytics-main\datasets\mydata  # dataset root dir
train: images/train  # train images (relative to 'path') 118287 images
val: images/val  # val images (relative to 'path') 5000 images
#test: test-dev2017.txt  # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794
 
# Classes
names:
  0: person
  1: cat
  2: dog
nc: 3

4、进行训练

上述步骤完成后，即可开始训练。

4.1、方法一

通过命令直接进行训练在其中指定参数，命令如下:

data.yaml根据实际路径而来

yolo detect train model=yolov10s.yaml data=data.yaml batch=16 epochs=100 imgsz=640

yolo detect train model=yolov10s.pt data=data.yaml batch=16 epochs=100 imgsz=640

4.2、方法二

通过创建py文件来进行训练

from ultralytics import YOLOv10

#数据集配置文件
data_yaml_path = r'D:\Yolov10\yolov10-main\yolov10-detect\data.yaml'
#预训练模型
pre_model_name = r'D:\Yolov10\yolov10-main\yolov10-detect\yolov10s.pt'

if __name__ == '__main__':
    #加载预训练模型
    model = YOLOv10(pre_model_name)
    #训练生成的文件保存路径名
    savename = 'train_yolov10s'
    #训练模型
    results = model.train(data=data_yaml_path,
                          epochs=10,
                          name=savename)

也可以使用yaml文件

from ultralytics import YOLOv10

#数据集配置文件
data_yaml_path = r'D:\Yolov10\yolov10-main\yolov10-detect\data.yaml'
#预训练模型
pre_model_name = r'D:\Yolov10\yolov10-main\ultralytics\cfg\models\v10\yolov10s.yaml'

if __name__ == '__main__':
    #加载预训练模型
    model = YOLOv10(pre_model_name)
    #训练生成的文件保存路径名
    savename = 'train_yolov10s'
    #训练模型
    results = model.train(data=data_yaml_path,
                          epochs=10,
                          name=savename)

注意修改类

训练过程（我这里后面多加了一类所以是4）：

训练过程中会保存以下内容，最后得到两个模型分别是：best.pt、last.pt

5、验证模型

训练进程完毕以后可使用一些验证数据进行模型验证，查看模型的识别效果。

5.1、命令行输入

yolo predict model=best.pt source='D:\Yolov10\yolov10-main\yolov10-detect\Testsets\test1'

5.2、脚本运行

from ultralytics import YOLOv10
import glob
import os
import numpy as np
import cv2

classes = {
    0: 'person', 1: 'cat', 2: 'dog', 3: 'backpack'
}
class Colors:
    """Ultralytics color palette https://ultralytics.com/."""

    def __init__(self):
        """Initialize colors as hex = matplotlib.colors.TABLEAU_COLORS.values()."""
        hexs = ('FF3838', 'FF9D97', 'FF701F', 'FFB21D', 'CFD231', '48F90A', '92CC17', '3DDB86', '1A9334', '00D4BB',
                '2C99A8', '00C2FF', '344593', '6473FF', '0018EC', '8438FF', '520085', 'CB38FF', 'FF95C8', 'FF37C7')
        self.palette = [self.hex2rgb(f'#{c}') for c in hexs]
        # print(self.palette)
        self.n = len(self.palette)

    def __call__(self, i, bgr=False):
        """Converts hex color codes to rgb values."""
        c = self.palette[int(i) % self.n]
        return (c[2], c[1], c[0]) if bgr else c

    @staticmethod
    def hex2rgb(h):  # rgb order (PIL)
        return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4))


colors = Colors()  # create instance for 'from utils.plots import colors'

#预测的图片路径
imgpath = r'D:\Yolov10\yolov10-main\yolov10-detect\Testsets\test1'
#模型路径
modelpath = r'D:\Yolov10\yolov10-main\yolov10-detect\runs\detect\weights\best.pt'
#保存结果的路径
save_dir = imgpath + '_Rst'
os.makedirs(save_dir,exist_ok=True)
model = YOLOv10(modelpath)

imgs = glob.glob(os.path.join(imgpath,'*.jpg'))
for img in imgs:
    imgname = img.split('\\')[-1]
    frame = cv2.imread(img)
    results = model.predict(img)[0]
    # results = model(img)

    for box in results.boxes:
        # print(box)
        xyxy = box.xyxy.squeeze().tolist()
        x1, y1, x2, y2 = int(xyxy[0]), int(xyxy[1]), int(xyxy[2]), int(xyxy[3])
        c, conf = int(box.cls), float(box.conf)
        name = classes[c]
        color = colors(c, True)
        cv2.rectangle(frame, (int(xyxy[0]), int(xyxy[1])), (int(xyxy[2]), int(xyxy[3])), color, thickness=2, lineType=cv2.LINE_AA)
        cv2.putText(frame, f"{name}: {conf:.2f}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, color,
                    2)
    # cv2.imshow('image', frame)
    # cv2.waitKey(0)
    cv2.imwrite(save_dir+'\\'+imgname,frame)