YOLO电力物目标检测训练

最近需要进行电力物检测相关的业务，因此制作了一个电力物数据集，使用YOLO目标检测方法进行实验，记录实验过程如下：

数据集标注

首先需要对电力物相关设备进行标注，这里我们选用labelme进行标注，使用无人机进行数据采集，得到了600余张图像，我们的数据集包含三类：防振锤、间隔棒以及压接管。

数据集转换为YOLO格式

使用Labelme标注完后，得到的是JSON文件（COCO格式），我们需要将其进行转换，同时，还需要将其按照8:2的比例划分数据集，代码如下：

import argparse
import json
import os
from tqdm import tqdm
import shutil
import random

#矩形转换
def convert_label_json(json_dir, save_dir, classes):
    json_paths = os.listdir(json_dir)
    classes = classes.split(',')

    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    for json_path in tqdm(json_paths):
        path = os.path.join(json_dir, json_path)

        # 尝试多种编码读取 JSON
        json_dict = None
        encodings = ["utf-8", "gbk"]
        for encoding in encodings:
            try:
                with open(path, "r", encoding=encoding) as load_f:
                    json_dict = json.load(load_f)
                break
            except (UnicodeDecodeError, json.JSONDecodeError):
                continue

        if json_dict is None:
            print(f"Failed to read {json_path}")
            continue

        h, w = json_dict['imageHeight'], json_dict['imageWidth']
        txt_path = os.path.join(save_dir, json_path.replace('.json', '.txt'))
        txt_file = open(txt_path, 'w')

        for shape_dict in json_dict['shapes']:
            label = shape_dict['label']
            label_index = classes.index(label)
            points = shape_dict['points']

            # 提取所有点的 x 和 y 坐标
            xs = [p[0] for p in points]
            ys = [p[1] for p in points]

            # 计算最小外接矩形
            x_min, x_max = min(xs), max(xs)
            y_min, y_max = min(ys), max(ys)

            # 计算中心点和宽高
            xc = (x_min + x_max) / 2 / w
            yc = (y_min + y_max) / 2 / h
            bw = (x_max - x_min) / w
            bh = (y_max - y_min) / h

            # 写入 YOLO 格式
            line = f"{label_index} {xc:.6f} {yc:.6f} {bw:.6f} {bh:.6f}\n"
            txt_file.write(line)

        txt_file.close()

# 检查文件夹是否存在
def mkdir(path):
    if not os.path.exists(path):
        os.makedirs(path)
 
def split_datasets(image_dir, txt_dir, save_dir):
    # 创建文件夹
    mkdir(save_dir)
    images_dir = os.path.join(save_dir, 'images')
    labels_dir = os.path.join(save_dir, 'labels')
 
    img_train_path = os.path.join(images_dir, 'train')
    img_test_path = os.path.join(images_dir, 'test')
    img_val_path = os.path.join(images_dir, 'val')
 
    label_train_path = os.path.join(labels_dir, 'train')
    label_test_path = os.path.join(labels_dir, 'test')
    label_val_path = os.path.join(labels_dir, 'val')
 
    mkdir(images_dir)
    mkdir(labels_dir)
    mkdir(img_train_path)
    mkdir(img_test_path)
    mkdir(img_val_path)
    mkdir(label_train_path)
    mkdir(label_test_path)
    mkdir(label_val_path)
 
    # 数据集划分比例，训练集75%，验证集15%，测试集15%，按需修改
    train_percent = 0.8
    val_percent = 0.2
    test_percent = 0
 
    total_txt = os.listdir(txt_dir)
    num_txt = len(total_txt)
    list_all_txt = range(num_txt)  # 范围 range(0, num)
 
    num_train = int(num_txt * train_percent)
    num_val = int(num_txt * val_percent)
    num_test = num_txt - num_train - num_val
 
    train = random.sample(list_all_txt, num_train)
    # 在全部数据集中取出train
    val_test = [i for i in list_all_txt if not i in train]
    # 再从val_test取出num_val个元素，val_test剩下的元素就是test
    val = random.sample(val_test, num_val)
 
    print("训练集数目：{}, 验证集数目：{},测试集数目：{}".format(len(train), len(val), len(val_test) - len(val)))
    for i in list_all_txt:
        name = total_txt[i][:-4]
 
        srcImage = os.path.join(image_dir, name + '.jpg')
        srcLabel = os.path.join(txt_dir, name + '.txt')
 
        if i in train:
            dst_train_Image = os.path.join(img_train_path, name + '.jpg')
            dst_train_Label = os.path.join(label_train_path, name + '.txt')
            shutil.copyfile(srcImage, dst_train_Image)
            shutil.copyfile(srcLabel, dst_train_Label)
        elif i in val:
            dst_val_Image = os.path.join(img_val_path, name + '.jpg')
            dst_val_Label = os.path.join(label_val_path, name + '.txt')
            shutil.copyfile(srcImage, dst_val_Image)
            shutil.copyfile(srcLabel, dst_val_Label)
        else:
            dst_test_Image = os.path.join(img_test_path, name + '.jpg')
            dst_test_Label = os.path.join(label_test_path, name + '.txt')
            shutil.copyfile(srcImage, dst_test_Image)
            shutil.copyfile(srcLabel, dst_test_Label)
 
  

if __name__=="__main__":

        # labelme生成的json格式标注转为yolov8支持的txt格式
        """
        python json2txt_nomalize.py --json-dir my_datasets/color_rings/jsons --save-dir my_datasets/color_rings/txts --classes "cat,dogs"
        """
        parser = argparse.ArgumentParser(description='json convert to txt params')
        parser.add_argument('--json-dir', type=str,default='D:\project_mine\detection\datasets/fangxiandata\data/labels', help='json path dir')
        parser.add_argument('--save-dir', type=str,default='D:\project_mine\detection\datasets/fangxiandata\data/outputs' ,help='txt save dir')
        parser.add_argument('--classes', type=str, default='fangzhenchui,jiangebang,yajieguan',help='classes')#,道路,房屋,水渠,桥
        args = parser.parse_args()
        json_dir = args.json_dir
        save_dir = args.save_dir
        classes = args.classes
        convert_label_json(json_dir, save_dir, classes)

        # 将图片和标注数据按比例切分为 训练集和测试集
 
        # """
        # python split_datasets.py --image-dir my_datasets/color_rings/imgs --txt-dir my_datasets/color_rings/txts --save-dir my_datasets/color_rings/train_data
        # """
        parser = argparse.ArgumentParser(description='split datasets to train,val,test params')
        parser.add_argument('--image-dir', type=str,default='D:\project_mine\detection\datasets/fangxiandata\data/images', help='image path dir')
        parser.add_argument('--txt-dir', type=str,default='D:\project_mine\detection\datasets/fangxiandata\data/outputs' , help='txt path dir')
        parser.add_argument('--save-dir', default='D:\project_mine\detection\datasets/fangxian',type=str, help='save dir')
        args_split = parser.parse_args()
        image_dir = args_split.image_dir
        txt_dir = args_split.txt_dir
        save_dir_split = args_split.save_dir

        split_datasets(image_dir, txt_dir, save_dir_split)

数据集配置

我们的数据集放到了fangxian文件夹中，需要设置对应的数据集配置文件：

path: ../datasets/fangxian # 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: fangzhenchui
  1: jiangebang
  2: yajieguan

开启训练

我们使用YOLO11网络，设置batch为8，epoch为100

from ultralytics import YOLO
model=YOLO("yolo11.yaml")
# Train the model
results = model.train(data="fangxian.yaml",
                      epochs=100,
                      batch=8,       # 根据GPU显存调整（T4建议batch=8）
                      imgsz=640,
                      device="0",     # 指定GPU ID
                      optimizer="AdamW",
                      lr0=1e-4,
                      warmup_epochs=4,
                      label_smoothing=0.1,
                      amp=True)

训练过程

从训练过程来看，这个数据集可能较为复杂，且博主没有使用任何预训练模型，因此其拟合的较慢，前十几个epoch都均为0，但从第20个epoch开始，其AP值逐渐有起色

随后AP便逐渐正常了。