当前位置：首页 > news >正文

YOLO电力物目标检测训练

news 来源：原创 2025/6/12 19:44:13

最近需要进行电力物检测相关的业务，因此制作了一个电力物数据集，使用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)# 尝试多种编码读取 JSONjson_dict = Noneencodings = ["utf-8", "gbk"]for encoding in encodings:try:with open(path, "r", encoding=encoding) as load_f:json_dict = json.load(load_f)breakexcept (UnicodeDecodeError, json.JSONDecodeError):continueif json_dict is None:print(f"Failed to read {json_path}")continueh, 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 / wyc = (y_min + y_max) / 2 / hbw = (x_max - x_min) / wbh = (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.8val_percent = 0.2test_percent = 0total_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_valtrain = random.sample(list_all_txt, num_train)# 在全部数据集中取出trainval_test = [i for i in list_all_txt if not i in train]# 再从val_test取出num_val个元素，val_test剩下的元素就是testval = 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_dirsave_dir = args.save_dirclasses = args.classesconvert_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_dirtxt_dir = args_split.txt_dirsave_dir_split = args_split.save_dirsplit_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: fangzhenchui1: jiangebang2: 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 IDoptimizer="AdamW",lr0=1e-4,warmup_epochs=4,label_smoothing=0.1,amp=True)