​​【深度学习】手把手教你玩转YOLOv8：训练自己的目标检测模型​​

import os
import shutil
import xml.etree.ElementTree as ET

# VOC格式数据集路径
voc_data_path = 'E:\\DataSet\\helmet-VOC'
voc_annotations_path = os.path.join(voc_data_path, 'Annotations')
voc_images_path = os.path.join(voc_data_path, 'JPEGImages')

# YOLO格式数据集保存路径
yolo_data_path = 'E:\\DataSet\\helmet-YOLO'
yolo_images_path = os.path.join(yolo_data_path, 'images')
yolo_labels_path = os.path.join(yolo_data_path, 'labels')

# 创建YOLO格式数据集目录
os.makedirs(yolo_images_path, exist_ok=True)
os.makedirs(yolo_labels_path, exist_ok=True)

# 类别映射 (可以根据自己的数据集进行调整)
class_mapping = {
    'head': 0,
    'helmet': 1,
    'person': 2,
    # 添加更多类别...
}

def convert_voc_to_yolo(voc_annotation_file, yolo_label_file):
    tree = ET.parse(voc_annotation_file)
    root = tree.getroot()

    size = root.find('size')
    width = float(size.find('width').text)
    height = float(size.find('height').text)

    with open(yolo_label_file, 'w') as f:
        for obj in root.findall('object'):
            cls = obj.find('name').text
            if cls not in class_mapping:
                continue
            cls_id = class_mapping[cls]
            xmlbox = obj.find('bndbox')
            xmin = float(xmlbox.find('xmin').text)
            ymin = float(xmlbox.find('ymin').text)
            xmax = float(xmlbox.find('xmax').text)
            ymax = float(xmlbox.find('ymax').text)

            x_center = (xmin + xmax) / 2.0 / width
            y_center = (ymin + ymax) / 2.0 / height
            w = (xmax - xmin) / width
            h = (ymax - ymin) / height

            f.write(f"{cls_id} {x_center} {y_center} {w} {h}\n")

# 遍历VOC数据集的Annotations目录，进行转换
for voc_annotation in os.listdir(voc_annotations_path):
    if voc_annotation.endswith('.xml'):
        voc_annotation_file = os.path.join(voc_annotations_path, voc_annotation)
        image_id = os.path.splitext(voc_annotation)[0]
        voc_image_file = os.path.join(voc_images_path, f"{image_id}.jpg")
        yolo_label_file = os.path.join(yolo_labels_path, f"{image_id}.txt")
        yolo_image_file = os.path.join(yolo_images_path, f"{image_id}.jpg")

        convert_voc_to_yolo(voc_annotation_file, yolo_label_file)
        if os.path.exists(voc_image_file):
            shutil.copy(voc_image_file, yolo_image_file)

print("转换完成！")

import os
import shutil
import random
 
# random.seed(0)  #随机种子，可自选开启
def split_data(file_path, label_path, new_file_path, train_rate, val_rate, test_rate):
	images = os.listdir(file_path)
	labels = os.listdir(label_path)
	images_no_ext = {os.path.splitext(image)[0]: image for image in images}
	labels_no_ext = {os.path.splitext(label)[0]: label for label in labels}
	matched_data = [(img, images_no_ext[img], labels_no_ext[img]) for img in images_no_ext if img in labels_no_ext]
 
	unmatched_images = [img for img in images_no_ext if img not in labels_no_ext]
	unmatched_labels = [label for label in labels_no_ext if label not in images_no_ext]
	if unmatched_images:
		print("未匹配的图片文件:")
		for img in unmatched_images:
			print(images_no_ext[img])
	if unmatched_labels:
		print("未匹配的标签文件:")
		for label in unmatched_labels:
			print(labels_no_ext[label])
 
	random.shuffle(matched_data)
	total = len(matched_data)
	train_data = matched_data[:int(train_rate * total)]
	val_data = matched_data[int(train_rate * total):int((train_rate + val_rate) * total)]
	test_data = matched_data[int((train_rate + val_rate) * total):]
 
	# 处理训练集
	for img_name, img_file, label_file in train_data:
		old_img_path = os.path.join(file_path, img_file)
		old_label_path = os.path.join(label_path, label_file)
		new_img_dir = os.path.join(new_file_path, 'train', 'images')
		new_label_dir = os.path.join(new_file_path, 'train', 'labels')
		os.makedirs(new_img_dir, exist_ok=True)
		os.makedirs(new_label_dir, exist_ok=True)
		shutil.copy(old_img_path, os.path.join(new_img_dir, img_file))
		shutil.copy(old_label_path, os.path.join(new_label_dir, label_file))
	# 处理验证集
	for img_name, img_file, label_file in val_data:
		old_img_path = os.path.join(file_path, img_file)
		old_label_path = os.path.join(label_path, label_file)
		new_img_dir = os.path.join(new_file_path, 'val', 'images')
		new_label_dir = os.path.join(new_file_path, 'val', 'labels')
		os.makedirs(new_img_dir, exist_ok=True)
		os.makedirs(new_label_dir, exist_ok=True)
		shutil.copy(old_img_path, os.path.join(new_img_dir, img_file))
		shutil.copy(old_label_path, os.path.join(new_label_dir, label_file))
	# 处理测试集
	for img_name, img_file, label_file in test_data:
		old_img_path = os.path.join(file_path, img_file)
		old_label_path = os.path.join(label_path, label_file)
		new_img_dir = os.path.join(new_file_path, 'test', 'images')
		new_label_dir = os.path.join(new_file_path, 'test', 'labels')
		os.makedirs(new_img_dir, exist_ok=True)
		os.makedirs(new_label_dir, exist_ok=True)
		shutil.copy(old_img_path, os.path.join(new_img_dir, img_file))
		shutil.copy(old_label_path, os.path.join(new_label_dir, label_file))
	print("数据集已划分完成")
 
if __name__ == '__main__':
	file_path = r"f:\data\JPEGImages"  # 图片文件夹
	label_path = r'f:\data\labels'  # 标签文件夹
	new_file_path = r"f:\VOCdevkit"  # 新数据存放位置
	split_data(file_path, label_path, new_file_path, train_rate=0.8, val_rate=0.1, test_rate=0.1)

from setuptools import setup, find_packages

setup(
    name="ultralytics",
    version="8.2.23",
    packages=find_packages(),
    install_requires=[
        "torch>=1.7.0",
        "opencv-python",
        "numpy",
        "matplotlib",
        "pillow",
        "pyyaml",
        "requests",
        "scipy",
        "tqdm",
    ],
    python_requires=">=3.7",
)

pip install -e .

import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO
if __name__ == '__main__':
  model = YOLO('ultralytics/cfg/models/v8/yolov8n.yaml')
  model.load('yolov8n.pt')  #注释则不加载
  results = model.train(
    data='data.yaml',  #数据集配置文件的路径
    epochs=200,  #训练轮次总数
    batch=16,  #批量大小，即单次输入多少图片训练
    imgsz=640,  #训练图像尺寸
    workers=8,  #加载数据的工作线程数
    device= 0,  #指定训练的计算设备，无nvidia显卡则改为 'cpu'
    optimizer='SGD',  #训练使用优化器，可选 auto,SGD,Adam,AdamW 等
    amp= True,  #True 或者 False, 解释为：自动混合精度(AMP) 训练
    cache=False  # True 在内存中缓存数据集图像，服务器推荐开启
)

# 找到这行代码：
ckpt = torch.load(file, map_location="cpu")

# 修改为：
ckpt = torch.load(file, map_location="cpu", weights_only=False)

import warnings
import torch

warnings.filterwarnings('ignore')
from ultralytics import YOLO

if __name__ == '__main__':
    # 使用上下文管理器允许不安全加载
    with torch.serialization.safe_globals([torch.nn.modules.container.Sequential]):
        model = YOLO('ultralytics/cfg/models/v8/yolov8n.yaml')
        model.load('yolov8n.pt')  #注释则不加载
        results = model.train(
            data='data.yaml',  #数据集配置文件的路径
            epochs=200,  #训练轮次总数
            batch=16,  #批量大小，即单次输入多少图片训练
            imgsz=640,  #训练图像尺寸
            workers=0,  #加载数据的工作线程数
            device= 0,  #指定训练的计算设备，无nvidia显卡则改为 'cpu'
            optimizer='SGD',  #训练使用优化器，可选 auto,SGD,Adam,AdamW 等
            amp= True,  #True 或者 False, 解释为：自动混合精度(AMP) 训练
            cache=False  # True 在内存中缓存数据集图像，服务器推荐开启
        )

import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO
if __name__ == '__main__':
    model = YOLO('runs/detect/train/weights/best.pt')  #修改为自己训练的模型路径
    model.val(data='hat.yaml',  #修改为自己的数据集yaml文件
              split='test',
              imgsz=640,
              batch=16,
              iou=0.6,  #阈值可以改，mAP50为0.5的情况下
              conf=0.001,
              workers=8,
              )

from ultralytics import YOLO
# 加载训练好的模型，改为自己的路径
model = YOLO('runs/detect/train/weights/best.pt')
# 修改为自己的图像或者文件夹的路径
source = 'test1.jpg' #修改为自己的图片路径及文件名
# 运行推理，并附加参数
model.predict(source, save=True)