我是个初学者,这个周末学习了目标检测,跑别人的模型到自己训练模型,收获不少,在这里分享学习心得,希望对大家有帮助

1.效果

(1)跑别人的模型

YOLOv11目标检测|跑别人的模型

(2)自己训练模型

YOLOv11目标检测|跑别人的模型

2.部分主要代码

(1)数据集转化

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

# VOC格式数据集路径
voc_data_path = 'mht'
voc_annotations_path = os.path.join(voc_data_path, 'labels')
voc_images_path = os.path.join(voc_data_path, 'images')

# YOLO格式数据集保存路径
yolo_data_path = '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 = {
    'kiwi': 0,
    'calyx': 1,
    # 添加更多类别...
}

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("转换完成!")

(2)数据集化分

import os
import shutil
import random

def make_yolo_dataset(images_folder, labels_folder, output_folder, train_ratio=0.8):
    # 创建目标文件夹
    images_train_folder = os.path.join(output_folder, 'data/train/images')
    images_val_folder = os.path.join(output_folder, 'data/valid/images')
    labels_train_folder = os.path.join(output_folder, 'data/train/labels')
    labels_val_folder = os.path.join(output_folder, 'data/valid/labels')

    os.makedirs(images_train_folder, exist_ok=True)
    os.makedirs(images_val_folder, exist_ok=True)
    os.makedirs(labels_train_folder, exist_ok=True)
    os.makedirs(labels_val_folder, exist_ok=True)

    # 获取图片和标签的文件名(不包含扩展名)
    image_files = [f for f in os.listdir(images_folder) if f.endswith('.png')]
    label_files = [f for f in os.listdir(labels_folder) if f.endswith('.txt')]
    image_base_names = set(os.path.splitext(f)[0] for f in image_files)
    label_base_names = set(os.path.splitext(f)[0] for f in label_files)

    # 找出图片和标签都存在的文件名
    matched_files = list(image_base_names & label_base_names)

    # 打乱顺序并划分为训练集和验证集
    random.shuffle(matched_files)
    split_idx = int(len(matched_files) * train_ratio)
    train_files = matched_files[:split_idx]
    val_files = matched_files[split_idx:]

    # 移动文件到对应文件夹
    for base_name in train_files:
        img_src = os.path.join(images_folder, f"{base_name}.png")
        lbl_src = os.path.join(labels_folder, f"{base_name}.txt")

        img_dst = os.path.join(images_train_folder, f"{base_name}.png")
        lbl_dst = os.path.join(labels_train_folder, f"{base_name}.txt")

        shutil.copyfile(img_src, img_dst)
        shutil.copyfile(lbl_src, lbl_dst)

    for base_name in val_files:
        img_src = os.path.join(images_folder, f"{base_name}.png")
        lbl_src = os.path.join(labels_folder, f"{base_name}.txt")

        img_dst = os.path.join(images_val_folder, f"{base_name}.png")
        lbl_dst = os.path.join(labels_val_folder, f"{base_name}.txt")

        shutil.copyfile(img_src, img_dst)
        shutil.copyfile(lbl_src, lbl_dst)

    print("数据集划分完成!")

# 使用示例
images_folder = 'yolo/images'  # 原始图片文件夹路径
labels_folder = 'yolo/labels'  # 原始标签文件夹路径
output_folder = 'D:/kc/yolov11_project'  # 存放结果数据集的文件夹路径
make_yolo_dataset(images_folder, labels_folder, output_folder)

(3)模型训练

from ultralytics import YOLO
if __name__  == '__main__':
    model=YOLO('yolo11n.pt')
    model.train(data='data/data.yaml',epochs=200,imgsz=640,workers=4)

(4)效果测试

import cv2
from ultralytics import YOLO

if __name__ == '__main__':
    # 加载预训练模型
    model = YOLO('last.pt')  # 确保 'best.pt' 是有效的模型文件路径
    
    # 指定要处理的图像文件路径
    img_path = r"data\test\images\000151.png"
    
    # 使用 predict 方法处理图像(注意:方法名可能因库版本而异)
    results = model(img_path)  # 这将返回一个包含 Detections 对象的列表
    # cap=cv2.VideoCapture(img_path)
    # 检查是否有检测结果,并显示第一个图像的结果(如果有的话)
#     while cap.isOpened():
#         res, frame=cap.read()
#         if res:
#             results=model(frame)
#             anotated_frame=results[0].plot()
#             cv2.imshow(winname="best.pt",mat=anotated_frame)
#             if cv2.waitKey(1)==27:
#                 break
#         else:
#             break
# cap.release()
# cv2.destroyAllWindows()
    results[0].show()

3.需要源代码的同学评论区留言

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐