【Yolov5】基于Yolov5的红外电压互感器检测

基于Yolov5的红外电压互感器检测

1. 数据集标注

1. 下载 labelimg

   pip install -i pypi.tuna.tsinghua.edu.cn/simple labelImg

   只用这个命令可能会有错误，网上有很多其他教程安装 labelImg

2. 标注数据集

   • 终端输入 labelimg打开软件，如下图所示
     

   • 在左边栏中选择 yolo 格式

     在此格式下标注会生成如下所示的 txt 文件：
     

     第一个数字表示类别，后面四个数字表示标注框的信息

   • 在左边栏中的 change save dir中选择自己想要保存标签的地方，open dir 打开要标注的数据文件夹

   • 按 W 键画框，并输入相应的标签；A/D选择上一张/下一张图片

2. Yolov5源码下载与训练准备

1. 克隆 Yolov5 源码

参考网上教程

源码结构如下所示：

3. 建立 Yolov5 格式的训练文件夹

   共需三个文件夹/文件，分别是 images、labels文件夹和一个 yaml 文件，树状结构如下：

   images：

   • train
   • test
   • val

   labels：

   • train
   • test
   • val

   yaml

4. 编写 split.py 函数将数据集和标签分成 train、test、val三部分

   # 将图片和标注数据按比例切分为 训练集和测试集、验证集
   import shutil
   import random
   import os
   
   # 原始路径
   image_original_path = 'd:/yolov5/MyDetection/voltage_transformer/data/'
   label_original_path = 'd:/yolov5/MyDetection/voltage_transformer/txt/'
   # 训练集路径
   train_image_path = 'd:/yolov5/MyDetection/voltage_transformer/images/train/'
   train_label_path = 'd:/yolov5/MyDetection/voltage_transformer/labels/train/'
   # 验证集路径
   val_image_path = 'd:/yolov5/MyDetection/voltage_transformer/images/val/'
   val_label_path = 'd:/yolov5/MyDetection/voltage_transformer/labels/val/'
   # 测试集路径
   test_image_path = 'd:/yolov5/MyDetection/voltage_transformer/images/test/'
   test_label_path = 'd:/yolov5/MyDetection/voltage_transformer/labels/test/'
   
   # 数据集划分比例，训练集80%，验证集10%，测试集10%
   train_percent = 0.8
   val_percent = 0.1
   test_percent = 0.1
   
   # 检查文件夹是否存在
   def mkdir():
       if not os.path.exists(train_image_path):
           os.makedirs(train_image_path)
       if not os.path.exists(train_label_path):
           os.makedirs(train_label_path)
   
       if not os.path.exists(val_image_path):
           os.makedirs(val_image_path)
       if not os.path.exists(val_label_path):
           os.makedirs(val_label_path)
   
       if not os.path.exists(test_image_path):
           os.makedirs(test_image_path)
       if not os.path.exists(test_label_path):
           os.makedirs(test_label_path)
   
   
   def main():
       mkdir()
   
       total_txt = os.listdir(label_original_path)
       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从list_all_txt取出num_train个元素
       # 所以list_all_txt列表只剩下了这些元素：val_test
       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)
       # 检查两个列表元素是否有重合的元素
       # set_c = set(val_test) & set(val)
       # list_c = list(set_c)
       # print(list_c)
       # print(len(list_c))
   
       print("训练集数目：{}, 验证集数目：{},测试集数目：{}".format(len(train), len(val), len(val_test) - len(val)))
       for i in list_all_txt:
           name = total_txt[i][:-4]
   
           srcImage = image_original_path + name + '.jpg'
           srcLabel = label_original_path + name + '.txt'
   
           if i in train:
               dst_train_Image = train_image_path + name + '.jpg'
               dst_train_Label = train_label_path + name + '.txt'
               shutil.copyfile(srcImage, dst_train_Image)
               shutil.copyfile(srcLabel, dst_train_Label)
           elif i in val:
               dst_val_Image = val_image_path + name + '.jpg'
               dst_val_Label = val_label_path + name + '.txt'
               shutil.copyfile(srcImage, dst_val_Image)
               shutil.copyfile(srcLabel, dst_val_Label)
           else:
               dst_test_Image = test_image_path + name + '.jpg'
               dst_test_Label = test_label_path + name + '.txt'
               shutil.copyfile(srcImage, dst_test_Image)
               shutil.copyfile(srcLabel, dst_test_Label)
   
   
   if __name__ == '__main__':
       main()
   

   （参考了网上代码）

5. 编写数据的 yaml 文件

   # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
   path: ./MyDetection/voltage_transformer  # dataset root dir
   train: images/train  
   val: images/val  
   test: images/test  
   
   # Classes
   nc: 1  # number of classes
   names: ['voltage transform']  # class names
   

3. 使用Google Colab进行训练

1. 登录 Google Drive并将已完成上述步骤的 Yolov5 文件夹导入

2. 在 Google Colab 新建笔记本并导入云盘的文件

   import os
   from google.colab import drive
   drive.mount('/content/drive')
   
   path = "/content/drive/My Drive"
   
   os.chdir(path)
   os.listdir(path)
   

3. 开启 GPU 加速

4. 开始训练

   !cd yolov5/;
   python train.py --weights ./weights/yolov5s.pt --data ./MyDetection/voltage_transformer/voltage_transformer.yaml --cfg ./models/yolov5s.yaml --batch-size 64 --device 0
   

   其中，

   --weights     表示初始化的参数，在yolov5官方下载的
   --data        表示数据的 yaml 文件
   --cfg         表示模型的 yaml 文件
   --batch-size  表示使用的 batch 大小为64
   --device      表示使用编号为 0 的 GPU
   --epochs      默认为 300
   

4. 结果分析

1. 训练结果
   

   在经历了 300 个 epoch 后：

   • 准确率§达到 0.946

   • 召回率®达到 0.966

   • 在 0.5IOU 下的 mAP达到 0.981

   • 在 0.5：0.05：0.95 IOU下的平均 mAP达到 0.673

   下面是训练产生的中间文件：

   

   可以看到：

   • 大概在前 30 次 epoch 时 loss 迅速下降，之后缓慢下降且保持在一个很小的水平
   • 准确率、召回率和 0.5IOU下的mAP 也在大概前 30 次 epoch达到最高水平，接下来在器附近波动
   • 0.5:0.05:0.95 IOU 下的平均 mAP在前 30 个 epoch 增长较快，接下来处于较慢的增长；再增加 epoch ，这个值的增长也会很慢，预计之后的最大值不会高于 0.7

   从上面可以看出，在大概 50 次的 epoch 后模型就已经表现很好了，想要提高的性能，可能需要更多的数据集。

2. 测试模型

   使用 detect.py 测试 test 中的数据：

   !cd yolov5;
   python detect.py --weights ./runs/train/exp43/weights/best.pt --source ./MyDetection/voltage_transformer/images/test/ --data ./MyDetection/voltage_transformer/voltage_transformer.yaml --conf-thres 0.8 --iou-thres 0.5 --device 0
   

   其中，

   --weights     表示训练后得到的最佳权重
   --source      表示测试集的路径
   --data        表示数据的 yaml 文件
   --conf-thres  表示置信阈值
   --iou-thres   表示 IOU 阈值
   --device 0    表示使用编号为 0 的 GPU
   

   测试结果是测试：模型在测试集上表现的很好，基本上都能正确检测且置信度大部分都达到了 0.9 以上。（由于数据集的关系这里的结果不能公开）