yolo11 训练coco128数据集，抽取部分类型，并整和自己的训练集

笼小粥

已于 2025-02-28 14:46:48 修改

阅读量732

点赞数 4

文章标签： python YOLO 图像处理视觉检测

于 2025-02-28 14:19:54 首次发布

本文链接：https://blog.youkuaiyun.com/mlzRuby/article/details/145919345

版权

一、训练coco128数据集

1. 环境准备

首先，你需要确保已经安装了 Python 环境（建议 Python 3.7 及以上版本），并安装 YOLOv8 相关依赖。可以使用以下命令来安装：

pip install ultralytics

2. 准备数据集

COCO128 数据集是 COCO 数据集的一个小型子集，包含 128 张图像。不过，我们需要对其进行筛选，只保留人和车的标注信息。

2.1 下载 COCO128 数据集

可以使用 YOLOv8 提供的下载功能来获取 COCO128 数据集：

from ultralytics import YOLO

# 下载COCO128数据集
model = YOLO('yolo11n.pt')
model.predict(task='detect', source='coco128', save=True)

2.2 筛选标注信息（**如果需要全部数据集类型，则无需筛选 **）
COCO 数据集的类别 ID 中，人对应的 ID 是 0，车（这里主要考虑汽车）对应的 ID 是 2。你需要编写一个脚本，遍历 COCO128 数据集中的标注文件（通常是.txt格式），只保留类别 ID 为 0 和 2 的标注信息。以下是一个示例脚本：

import os
import shutil
from pathlib import Path

# 定义COCO128数据集路径
coco128_path = Path('/home/lemon/opencv/dahua_rtsp/coco128')
images_path = coco128_path / 'images/train2017'
labels_path = coco128_path / 'labels/train2017'

# 创建新的数据集目录
new_dataset_path = Path('/home/lemon/opencv/dahua_rtsp/coco128_person_car')
new_images_path = new_dataset_path / 'images/train2017'
new_labels_path = new_dataset_path / 'labels/train2017'

# 创建新目录
new_images_path.mkdir(parents=True, exist_ok=True)
new_labels_path.mkdir(parents=True, exist_ok=True)

# 筛选只包含人和车的标注文件，并修改车的类别ID
for label_file in labels_path.rglob('*.txt'):
    with open(label_file, 'r') as f:
        lines = f.readlines()
        valid_lines = []
        for line in lines:
            parts = line.split()
            class_id = int(parts[0])
            if class_id == 0:  # 人，ID保持为0
                valid_lines.append(line)
            elif class_id == 2:  # 车，将ID从2改为1
                parts[0] = '1'
                new_line = ' '.join(parts) + '\n'
                valid_lines.append(new_line)

        if valid_lines:
            # 复制图像和标注文件到新目录
            image_file = images_path / label_file.with_suffix('.jpg').name
            new_image_file = new_images_path / label_file.with_suffix('.jpg').name
            new_label_file = new_labels_path / label_file.name
            
            # 检查源图像文件是否存在
            if not image_file.exists():
                print(f"Image file {image_file} not found. Skipping...")
                continue

            shutil.copyfile(image_file, new_image_file)
            with open(new_label_file, 'w') as new_f:
                new_f.writelines(valid_lines)

下一步train和val都是用同一组数据，为了方便所以创建train和valid文件夹，内部images和labels内容一样

3. 创建数据集配置文件

在 coco128_person_car 目录下创建一个 YAML 文件（例如coco128_person_car.yaml），用于指定数据集的路径和类别信息：

train: /home/lemon/opencv/dahua_rtsp/coco128_person_car/train/images/train2017  # 训练集图像绝对路径
val: /home/lemon/opencv/dahua_rtsp/coco128_person_car/valid/images/train2017    # 验证集图像绝对路径，这里简单使用训练集作为验证集

# 类别数量
nc: 2

# 类别名称
names: ['person', 'car']

4. 训练模型

运行cmd命令

yolo detect train data=/home/lemon/opencv/dahua_rtsp/coco128_person_car/data.yaml model=yolo11n.pt epochs=20 batch=5 imgsz=640 device=0    project=/home/lemon/opencv/dahua_rtsp/coco128_person_car/runs

或者使用以下代码来训练 YOLOv8 模型：

from ultralytics import YOLO

# 加载预训练模型
model = YOLO('yolo11n.pt')

# 训练模型（参数根据自己需要添加）
results = model.train(
    data='/home/lemon/opencv/dahua_rtsp/coco128_person_car/data.yaml',  # 数据集配置文件路径
    epochs=20,                        # 训练轮数
    imgsz=640                         # 输入图像尺寸
)

5.验证和推理

import cv2
import torch
import os
from ultralytics import YOLO

# 明确指定第一个 CUDA 设备（如果有可用的 GPU）
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")

# RTSP 视频流地址
rtsp_url = "rtsp://admin:abc123456@169.254.182.3:554/Streaming/Channels/1"
cap = cv2.VideoCapture(rtsp_url)

# 检查视频流是否成功打开
if not cap.isOpened():
    print("Error: Could not open video stream.")
    exit()

# 加载自定义的 YOLOv8 模型到指定设备
model_path = '/home/lemon/opencv/dahua_rtsp/coco128_person_car/runs/train/weights/best.pt'  # 根据实际路径修改
model = YOLO(model_path)
model.to(device)  # 将模型移动到指定设备

# 设置目标画布大小
target_width = 800
target_height = 600

while True:
    ret, frame = cap.read()
    if ret:
        # 使用 YOLO 模型进行检测
        results = model(frame)

        # 取出列表中的第一个元素
        result = results[0]

        # 检查是否有检测结果
        if len(result.boxes) > 0:
            # 获取带有检测结果的图像
            result_img = result.plot()
        else:
            result_img = frame  # 如果没有检测结果，显示原始帧

        # 缩放图像到目标大小
        resized_img = cv2.resize(result_img, (target_width, target_height))

        # 显示带有检测结果的图像
        cv2.imshow('person_car Detection', resized_img)

        # 按 'q' 键退出循环
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    else:
        print("Error: Could not read frame.")
        break

# 释放资源
cap.release()
cv2.destroyAllWindows()

二、加入自己的数据集

1.准备自己的数据集，更改类型id，不能与已有的冲突

import os

# 定义多个标注文件所在的文件夹路径
label_folders = [
    '/home/lemon/opencv/dahua_rtsp/yolo/yolo11-test/train/labels', # 训练集
    '/home/lemon/opencv/dahua_rtsp/yolo/yolo11-test/valid/labels'  # 验证集
]

# 定义旧的类别 ID 到新的类别 ID 的映射
id_mapping = {
    0: 2,  # 例如将原来 ID 为 0 的类别改为 2
    1: 3   # 将原来 ID 为 1 的类别改为 3
}

# 遍历每个标注文件夹
for label_folder in label_folders:
    # 遍历标注文件夹中的所有标注文件
    for filename in os.listdir(label_folder):
        if filename.endswith('.txt'):
            file_path = os.path.join(label_folder, filename)
            new_lines = []
            with open(file_path, 'r') as f:
                lines = f.readlines()
                for line in lines:
                    parts = line.strip().split()
                    old_id = int(parts[0])
                    if old_id in id_mapping:
                        new_id = id_mapping[old_id]
                        parts[0] = str(new_id)
                    new_line = ' '.join(parts) + '\n'
                    new_lines.append(new_line)
            # 将修改后的内容写回文件
            with open(file_path, 'w') as f:
                f.writelines(new_lines)

2. 合并数据集

2.1 合并图像文件夹
将自己数据集的 images 文件夹中的图像复制到 COCO128 的 images 文件夹中。确保文件名不冲突（可使用脚本为文件添加唯一前缀）。

import shutil
import os

# 定义多个源图像文件夹和对应的目标图像文件夹
source_target_pairs = [
    (
        '/home/lemon/opencv/dahua_rtsp/yolo/yolo11-test/train/images', #训练集目录
        '/home/lemon/opencv/dahua_rtsp/coco128_person_car/train/images/train2017'
    ),
    (
        '/home/lemon/opencv/dahua_rtsp/yolo/yolo11-test/valid/images',  # 验证集目录
        '/home/lemon/opencv/dahua_rtsp/coco128_person_car/valid/images/train2017' 
    )
    # 可以根据需要添加更多的源 - 目标文件夹对
]

# 遍历每个源 - 目标文件夹对
for source_images_folder, target_images_folder in source_target_pairs:
    # 检查源文件夹是否存在
    if not os.path.exists(source_images_folder):
        print(f"Source folder {source_images_folder} does not exist. Skipping...")
        continue

    # 检查目标文件夹是否存在，若不存在则创建
    if not os.path.exists(target_images_folder):
        os.makedirs(target_images_folder)

    # 遍历源文件夹中的所有文件
    for filename in os.listdir(source_images_folder):
        source_path = os.path.join(source_images_folder, filename)
        target_path = os.path.join(target_images_folder, filename)
        try:
            # 复制文件
            shutil.copyfile(source_path, target_path)
            print(f"Copied {source_path} to {target_path}")
        except Exception as e:
            print(f"Error copying {source_path} to {target_path}: {e}")

2.2 合并标注文件夹
根据标注格式的不同，有不同的合并方式。

import shutil
import os

# 定义多个源标注文件夹和对应的目标标注文件夹对
source_target_pairs = [
    (
        '/home/lemon/opencv/dahua_rtsp/yolo/yolo11-test/train/labels',#训练集
        '/home/lemon/opencv/dahua_rtsp/coco128_person_car/train/labels/train2017'
    ),
    (
        '/home/lemon/opencv/dahua_rtsp/yolo/yolo11-test/valid/labels',  # 验证集
        '/home/lemon/opencv/dahua_rtsp/coco128_person_car/valid/labels/train2017' 
    )
    # 可以根据需要添加更多的源 - 目标标注文件夹对
]

# 遍历每一对源和目标标注文件夹
for source_labels_folder, target_labels_folder in source_target_pairs:
    # 检查源标注文件夹是否存在
    if not os.path.exists(source_labels_folder):
        print(f"源标注文件夹 {source_labels_folder} 不存在，跳过...")
        continue

    # 检查目标标注文件夹是否存在，若不存在则创建
    if not os.path.exists(target_labels_folder):
        os.makedirs(target_labels_folder)

    # 遍历源标注文件夹中的所有文件
    for filename in os.listdir(source_labels_folder):
        source_path = os.path.join(source_labels_folder, filename)
        target_path = os.path.join(target_labels_folder, filename)
        try:
            # 复制标注文件
            shutil.copyfile(source_path, target_path)
            print(f"已将 {source_path} 复制到 {target_path}")
        except Exception as e:
            print(f"复制 {source_path} 到 {target_path} 时出错: {e}")

3. 更新数据集配置文件

train: /home/lemon/opencv/dahua_rtsp/coco128_person_car/train/images/train2017 
val: /home/lemon/opencv/dahua_rtsp/coco128_person_car/valid/images/train2017 
nc: 原始类别数 + 新增动物类别数
names: ['原始类别1', '原始类别2', ..., '新增动物类别1', '新增动物类别2', ...]