yolov10＋strongsort实现多目标跟踪

v_15717690526

于 2025-02-13 15:08:49 发布

阅读量733

点赞数 15

文章标签：机器学习深度学习人工智能

本文链接：https://blog.youkuaiyun.com/l15717690526/article/details/145612865

版权

一、准备数据集

1.1 数据集标注

这时候使用labelimg工具标注自己的数据集，一定要分清楚自己的类别信息

安装完anconda之后打开命令框输入以下命令进入标注页面

conda activate labelimg

labelimg

然后标注自己的数据集就行。

1.2 数据集划分

使用以下代码可以按照比例划分自己的数据集

import os
import random
import shutil
 
# 设置路径
image_folder = ''  # 图像文件夹路径
label_folder = ''  # 标签文件夹路径
train_image_folder = ''  # 训练集图像文件夹路径
train_label_folder = ''  # 训练集标签文件夹路径
val_image_folder = ''  # 验证集图像文件夹路径
val_label_folder = ''  # 验证集标签文件夹路径
 
# 创建训练集和验证集文件夹
os.makedirs(train_image_folder, exist_ok=True)
os.makedirs(train_label_folder, exist_ok=True)
os.makedirs(val_image_folder, exist_ok=True)
os.makedirs(val_label_folder, exist_ok=True)
 
# 获取所有的图像文件（假设图像是.jpg 格式，标签是 .txt 格式）
image_files = [f for f in os.listdir(image_folder) if f.endswith('.jpg')]
label_files = [f for f in os.listdir(label_folder) if f.endswith('.txt')]
 
# 确保每个图像都有一个对应的标签
assert len(image_files) == len(label_files), "图像和标签数量不匹配"
 
# 打乱文件顺序
data = list(zip(image_files, label_files))
random.shuffle(data)
 
# 按照 9:1 划分数据集
train_size = int(0.9 * len(data))
train_data = data[:train_size]
val_data = data[train_size:]
 
# 复制文件到训练集和验证集目录
for image, label in train_data:
    shutil.copy(os.path.join(image_folder, image), os.path.join(train_image_folder, image))
    shutil.copy(os.path.join(label_folder, label), os.path.join(train_label_folder, label))
 
for image, label in val_data:
    shutil.copy(os.path.join(image_folder, image), os.path.join(val_image_folder, image))
    shutil.copy(os.path.join(label_folder, label), os.path.join(val_label_folder, label))
 
print(f"数据集划分完成：{len(train_data)} 张图像用于训练，{len(val_data)} 张图像用于验证。")

只需要输入好自己数据集的路径就可以开始划分数据集了，到此数据集准备完毕。

二、训练yolov10的权重

2.1 修改训练用的yaml文件

在训练yolov10权重的时候一定要注意将类别信息改成自己的信息

将代码里面的nc改成自己的真实类别

# Ultralytics YOLO 🚀, AGPL-3.0 license
# YOLOv8 object detection model with P3-P5 outputs. For Usage examples see https://docs.ultralytics.com/tasks/detect
 
# Parameters
nc: 80 # number of classes
scales: # model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8.yaml with scale 'n'
  # [depth, width, max_channels]
  n: [0.33, 0.25, 1024] # YOLOv8n summary: 225 layers,  3157200 parameters,  3157184 gradients,   8.9 GFLOPs
 
# YOLOv8.0n backbone
backbone:
  # [from, repeats, module, args]
  - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2
  - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4
  - [-1, 3, C2f, [128, True]]
  - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8
  - [-1, 6, C2f, [256, True]]
  - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16
  - [-1, 6, C2f, [512, True]]
  - [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32
  - [-1, 3, C2f, [1024, True]]
  - [-1, 1, SPPF, [1024, 5]] # 9
 
# YOLOv8.0n head
head:
  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 6], 1, Concat, [1]] # cat backbone P4
  - [-1, 3, C2f, [512]] # 12
 
  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 4], 1, Concat, [1]] # cat backbone P3
  - [-1, 3, C2f, [256]] # 15 (P3/8-small)
 
  - [-1, 1, Conv, [256, 3, 2]]
  - [[-1, 12], 1, Concat, [1]] # cat head P4
  - [-1, 3, C2f, [512]] # 18 (P4/16-medium)
 
  - [-1, 1, Conv, [512, 3, 2]]
  - [[-1, 9], 1, Concat, [1]] # cat head P5
  - [-1, 3, C2f, [1024]] # 21 (P5/32-large)
 
  - [[15, 18, 21], 1, Detect, [nc]] # Detect(P3, P4, P5)

2.2 创建yaml文件

自己创建一个.yaml文件用于自己数据集的路径选择

# Ultralytics YOLO 🚀, AGPL-3.0 license
# COCO8 dataset (first 8 images from COCO train2017) by Ultralytics
# Documentation: https://docs.ultralytics.com/datasets/detect/coco8/
# Example usage: yolo train data=coco8.yaml
# parent
# ├── ultralytics
# └── datasets
#     └── coco8  ← downloads here (1 MB)
 
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ../datasets/coco8 # dataset root dir
train: images/train # train images (relative to 'path') 4 images
val: images/val # val images (relative to 'path') 4 images
test: # test images (optional)
 
# Classes
names:
  0: person
  1: bicycle
  2: car
  3: motorcycle
  4: airplane
  5: bus
  6: train
  7: truck
  8: boat
  9: traffic light
  10: fire hydrant
  11: stop sign
  12: parking meter
  13: bench
  14: bird
  15: cat
  16: dog
  17: horse
  18: sheep
  19: cow
  20: elephant
  21: bear
  22: zebra
  23: giraffe
  24: backpack
  25: umbrella
  26: handbag
  27: tie
  28: suitcase
  29: frisbee
  30: skis
  31: snowboard
  32: sports ball
  33: kite
  34: baseball bat
  35: baseball glove
  36: skateboard
  37: surfboard
  38: tennis racket
  39: bottle
  40: wine glass
  41: cup
  42: fork
  43: knife
  44: spoon
  45: bowl
  46: banana
  47: apple
  48: sandwich
  49: orange
  50: broccoli
  51: carrot
  52: hot dog
  53: pizza
  54: donut
  55: cake
  56: chair
  57: couch
  58: potted plant
  59: bed
  60: dining table
  61: toilet
  62: tv
  63: laptop
  64: mouse
  65: remote
  66: keyboard
  67: cell phone
  68: microwave
  69: oven
  70: toaster
  71: sink
  72: refrigerator
  73: book
  74: clock
  75: vase
  76: scissors
  77: teddy bear
  78: hair drier
  79: toothbrush
 
# Download script/URL (optional)
download: https://github.com/ultralytics/assets/releases/download/v0.0.0/coco8.zip

把类比什么的改成自己的数据集，数据集的路径也要替换自己的路径。

2.3 创建一个train.py文件

import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLOv8
 
if __name__ == '__main__':
    model = YOLOv8('yolov8n.yaml')
   #  model.load('yolov8n.pt') # loading pretrain weights
    model.train(data='创建的.yaml',
                cache=False,
                imgsz=640,
                epochs=100,
                batch=32,
                close_mosaic=0,
                workers=8,
                device='0',
                optimizer='SGD', # using SGD
                project='runs/train',
                name='exp',
                )