devbox
AllinToyou
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

【矩池云】YOLOv3~YOLOv5训练红外小目标数据集

作者：AllinToyou | 2024-04-04 21:53:34

赞

踩

红外小目标数据集

一、数据集准备

数据集下载地址：https://github.com/YimianDai/sirst

1. 需要将数据集转换为YOLO所需要的txt格式

参考链接：https://github.com/pprp/voc2007_for_yolo_torch

1.1 检测图片及其xml文件

  1. import os, shutil
  2.  
  3. def checkPngXml(dir1, dir2, dir3, is_move=True):
  4.     """
  5.     dir1 是图片所在文件夹
  6.     dir2 是标注文件所在文件夹
  7.     dir3 是创建的，如果图片没有对应的xml文件，那就将图片放入dir3
  8.     is_move 是确认是否进行移动，否则只进行打印
  9.     """
  10.     if not os.path.exists(dir3):
  11.         os.mkdir(dir3)
  12.     cnt = 0
  13.     for file in os.listdir(dir1):
  14.         f_name,f_ext = file.split(".")
  15.         if not os.path.exists(os.path.join(dir2, f_name+".xml")):
  16.             print(f_name)
  17.             if is_move:
  18.                 cnt += 1
  19.                 shutil.move(os.path.join(dir1,file), os.path.join(dir3, file))
  20.     if cnt > 0:
  21.         print("有%d个文件不符合要求，已打印。"%(cnt))
  22.     else:
  23.         print("所有图片和对应的xml文件都是一一对应的。")
  24.  
  25. if __name__ == "__main__":
  26.     dir1 = r"dataset/images/images"   # 修改为自己的图片路径
  27.     dir2 = r"dataset/masks/masks"     # 修改为自己的图片路径
  28.     dir3 = r"dataset/Allempty"        # 修改为自己的图片路径
  29.     checkPngXml(dir1, dir2, dir3, False)

1.2 划分训练集

  1. import os
  2. import random
  3. import os, fnmatch
  4.  
  5.  
  6. trainval_percent = 0.8
  7. train_percent = 0.8
  8.  
  9. xmlfilepath = r"dataset/masks/masks"
  10. txtsavepath = r"dataset"
  11. total_xml = fnmatch.filter(os.listdir(xmlfilepath), '*.xml')
  12. print(total_xml)
  13.  
  14. num=len(total_xml)
  15. list=range(num)
  16. tv=int(num*trainval_percent)
  17. tr=int(tv*train_percent)
  18. trainval= random.sample(list,tv)
  19. train=random.sample(trainval,tr)
  20.  
  21. ftrainval = open('dataset/trainval.txt', 'w')
  22. ftest = open('dataset/test.txt', 'w')
  23. ftrain = open('dataset/train.txt', 'w')
  24. fval = open('dataset/val.txt', 'w')
  25.  
  26. for i  in list:
  27.     name=total_xml[i][:-4]+'\n'
  28.     if i in trainval:
  29.         ftrainval.write(name)
  30.         if i in train:
  31.             ftrain.write(name)
  32.         else:
  33.             fval.write(name)
  34.     else:
  35.         ftest.write(name)
  36.  
  37. ftrainval.close()
  38. ftrain.close()
  39. fval.close()
  40. ftest.close()

1.3 转为txt标签

  1. # -*- coding: utf-8 -*-
  2. """
  3. 需要修改的地方：
  4. 1. sets中替换为自己的数据集
  5. 2. classes中替换为自己的类别
  6. 3. 将本文件放到VOC2007目录下
  7. 4. 直接开始运行
  8. """
  9.  
  10. import xml.etree.ElementTree as ET
  11. import pickle
  12. import os
  13. from os import listdir, getcwd
  14. from os.path import join
  15. sets=[('2007', 'train'), ('2007', 'val'), ('2007', 'test')]  #替换为自己的数据集
  16. classes = ["Target"]     #修改为自己的类别
  17.  
  18. def convert(size, box):
  19.     dw = 1./(size[0])
  20.     dh = 1./(size[1])
  21.     x = (box[0] + box[1])/2.0
  22.     y = (box[2] + box[3])/2.0
  23.     w = box[1] - box[0]
  24.     h = box[3] - box[2]
  25.     x = x*dw
  26.     w = w*dw
  27.     y = y*dh
  28.     h = h*dh
  29.     return (x,y,w,h)
  30.  
  31. def convert_annotation(year, image_id):
  32.     in_file = open('dataset/masks/masks/%s.xml'%(image_id))  #将数据集放于当前目录下
  33.     out_file = open('dataset/labels/%s.txt'%(image_id), 'w')
  34.     tree=ET.parse(in_file)
  35.     root = tree.getroot()
  36.     size = root.find('size')
  37.     w = int(size.find('width').text)
  38.     h = int(size.find('height').text)
  39.     print(w,h)
  40.     for obj in root.iter('object'):
  41.         difficult = obj.find('difficult').text
  42.         cls = obj.find('name').text
  43.         if cls not in classes or int(difficult)==1:
  44.             continue
  45.         cls_id = classes.index(cls)
  46.         print(cls_id)
  47.         xmlbox = obj.find('bndbox')
  48.         b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
  49.         bb = convert((w,h), b)
  50.         print(bb)
  51.         out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
  52. wd = getcwd()
  53. for year, image_set in sets:
  54.     if not os.path.exists('dataset/labels/'):
  55.         os.makedirs('dataset/labels/')
  56.     image_ids = open('dataset/%s.txt'%(image_set)).read().strip().split()
  57.     list_file = open('%s_%s.txt'%(year, image_set), 'w')
  58.     for image_id in image_ids:
  59.         list_file.write('dataset/images/images/%s.png\n'%(image_id))
  60.         convert_annotation(year, image_id)
  61.     list_file.close()
  62. # os.system("cat 2007_train.txt 2007_val.txt > train.txt")     #修改为自己的数据集用作训练

1.4 构造数据集

  1. import os, shutil
  2.  
  3. """
  4. 需要满足以下条件：
  5. 1. 在JPEGImages中准备好图片
  6. 2. 在labels中准备好labels
  7. 3. 创建好如下所示的文件目录：
  8.     - images
  9.         - train2014
  10.         - val2014
  11.     - labels(由于voc格式中有labels文件夹，所以重命名为label)
  12.         - train2014
  13.         - val2014
  14. """
  15.  
  16.  
  17. def make_for_torch_yolov3(dir_image,
  18.                                  dir_label,
  19.                                  dir1_train,
  20.                                  dir1_val,
  21.                                  dir2_train,
  22.                                  dir2_val,
  23.                                  main_trainval,
  24.                                  main_test):
  25.     if not os.path.exists(dir1_train):
  26.         os.mkdir(dir1_train)
  27.     if not os.path.exists(dir1_val):
  28.         os.mkdir(dir1_val)
  29.     if not os.path.exists(dir2_train):
  30.         os.mkdir(dir2_train)
  31.     if not os.path.exists(dir2_val):
  32.         os.mkdir(dir2_val)
  33.  
  34.     with open(main_trainval, "r") as f1:
  35.         for line in f1:
  36.             print(line[:-1])
  37.             # print(os.path.join(dir_image, line[:-1]+".png"), os.path.join(dir1_train, line[:-1]+".png"))
  38.             shutil.copy(os.path.join(dir_image, line[:-1]+".png"),
  39.                         os.path.join(dir1_train, line[:-1]+".png"))
  40.             shutil.copy(os.path.join(dir_label, line[:-1]+".txt"),
  41.                         os.path.join(dir2_train, line[:-1]+".txt"))
  42.  
  43.  
  44.     with open(main_test, "r") as f2:
  45.         for line in f2:
  46.             print(line[:-1])
  47.             shutil.copy(os.path.join(dir_image, line[:-1]+".png"),
  48.                         os.path.join(dir1_val, line[:-1]+".png"))
  49.             shutil.copy(os.path.join(dir_label, line[:-1]+".txt"),
  50.                         os.path.join(dir2_val, line[:-1]+".txt"))
  51.  
  52. if __name__ == "__main__":
  53.     '''
  54.     https://github.com/ultralytics/yolov3
  55.     这个pytorch版本的数据集组织
  56.     - images
  57.         - train2014 # dir1_train
  58.         - val2014 # dir1_val
  59.     - labels
  60.         - train2014 # dir2_train
  61.         - val2014 # dir2_val
  62.     trainval.txt, test.txt 是由create_main.py构建的
  63.     '''
  64.  
  65.     dir_image = r"dataset/images/images"
  66.     dir_label = r"dataset/labels"
  67.  
  68.     dir1_train = r"dataset/image/train2014"
  69.     dir1_val = r"dataset/image/val2014"
  70.  
  71.     dir2_train = r"dataset/label/train2014"
  72.     dir2_val = r"dataset/label/val2014"
  73.  
  74.     main_trainval = r"dataset/trainval.txt"
  75.     main_test = r"dataset/test.txt"
  76.  
  77.     make_for_torch_yolov3(dir_image,
  78.                             dir_label,
  79.                             dir1_train,
  80.                             dir1_val,
  81.                             dir2_train,
  82.                             dir2_val,
  83.                             main_trainval,
  84.                             main_test)

最终数据集格式如下：

037b886f0f004da1b6e0ed1ae8cd45c6.png

 2. 构造训练所需要的数据集

根据以上数据集 需要单独构建一个datasets文件夹，存放标签和图像，具体格式如下：

36edb51c6b4c4003ac8723e253c6b8dd.png

可以参考该链接：https://github.com/ultralytics/yolov5/issues/7389

3. 构建数据配置文档，需要注意 YOLOv5目录需要和datasets目录同级。

命名为hongwai.yaml

  1. # YOLOv3 
    声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/AllinToyou/article/detail/361349
    推荐阅读
    相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。