当前位置:   article > 正文

YOLOv5+Deepsort训练自己的数据集实现多目标跟踪

yolov5+deepsort训练自己的数据集

引言

本文利用yolov5-6.1版本完成目标检测模块,利用deepsort跟踪算法实现目标跟踪模块,将二者集成,在自己的数据集上形成一套行之有效的目标检测+跟踪模型。

yolov5-6.1版本代码下载地址:yolov5 deepsort- CSDN搜索 (github.com)https://github.com/ultralytics/yolov5deepsort代码下载地址:https://github.com/mikel-brostrom/Yolov5_DeepSort_Pytorchhttps://github.com/mikel-brostrom/Yolov5_DeepSort_Pytorch注意:root下的models和utils是使用的yolov5的v6.1版本的代码,如使用其它版本可以用相应版本的models和utils代码替换。

 1 代码整体框架

其中分为yolov5-6.1模块,deepsort模块。

2 目标检测模块

2.1 环境依赖

YOLOv5的代码是开源的,因此我们可以从github上克隆其源码。该项目利用yolov5-6.1版本来作为讲解。下载yolov5-6.1代码,其目录结构如下:

data:主要是存放一些超参数的配置文件,是用来配置训练集和测试集还有验证集的路径的。如果是训练自己的数据集的话,那么就需要修改其中的yaml文件。

models:里面主要是一些网络构建的配置文件和函数。如果训练自己的数据集的话,就需要修改这里面相对应的yaml文件来训练自己模型。

 utils:存放的是工具类的函数,里面有loss函数,metrics函数,plots函数等。

weights:放置训练好的权重参数。

detect.py:利用训练好的权重参数进行目标检测。

 train.py:训练自己的数据集的函数。

 test.py:测试训练的结果的函数。

requirements.txt:yolov5-6.1的环境依赖包。利用以下命令完成环境依赖安装。

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

2.2 数据集准备

       这里需要将VOC(xml)格式的数据集转换成yolo所需的txt格式,需要对xml格式的标签文件转换为txt文件。将标注好的数据集按照以下形式存放。

 运行代码:

  1. import xml.etree.ElementTree as ET
  2. import pickle
  3. import os
  4. from os import listdir, getcwd
  5. from os.path import join
  6. import random
  7. from shutil import copyfile
  8. classes = ["hat", "person"]
  9. #classes=["ball"]
  10. TRAIN_RATIO = 80
  11. def clear_hidden_files(path):
  12. dir_list = os.listdir(path)
  13. for i in dir_list:
  14. abspath = os.path.join(os.path.abspath(path), i)
  15. if os.path.isfile(abspath):
  16. if i.startswith("._"):
  17. os.remove(abspath)
  18. else:
  19. clear_hidden_files(abspath)
  20. def convert(size, box):
  21. dw = 1./size[0]
  22. dh = 1./size[1]
  23. x = (box[0] + box[1])/2.0
  24. y = (box[2] + box[3])/2.0
  25. w = box[1] - box[0]
  26. h = box[3] - box[2]
  27. x = x*dw
  28. w = w*dw
  29. y = y*dh
  30. h = h*dh
  31. return (x,y,w,h)
  32. def convert_annotation(image_id):
  33. in_file = open('VOCdevkit/VOC2007/Annotations/%s.xml' %image_id)
  34. out_file = open('VOCdevkit/VOC2007/YOLOLabels/%s.txt' %image_id, 'w')
  35. tree=ET.parse(in_file)
  36. root = tree.getroot()
  37. size = root.find('size')
  38. w = int(size.find('width').text)
  39. h = int(size.find('height').text)
  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. xmlbox = obj.find('bndbox')
  47. b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
  48. bb = convert((w,h), b)
  49. out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
  50. in_file.close()
  51. out_file.close()
  52. wd = os.getcwd()
  53. wd = os.getcwd()
  54. data_base_dir = os.path.join(wd, "VOCdevkit/")
  55. if not os.path.isdir(data_base_dir):
  56. os.mkdir(data_base_dir)
  57. work_sapce_dir = os.path.join(data_base_dir, "VOC2007/")
  58. if not os.path.isdir(work_sapce_dir):
  59. os.mkdir(work_sapce_dir)
  60. annotation_dir = os.path.join(work_sapce_dir, "Annotations/")
  61. if not os.path.isdir(annotation_dir):
  62. os.mkdir(annotation_dir)
  63. clear_hidden_files(annotation_dir)
  64. image_dir = os.path.join(work_sapce_dir, "JPEGImages/")
  65. if not os.path.isdir(image_dir):
  66. os.mkdir(image_dir)
  67. clear_hidden_files(image_dir)
  68. yolo_labels_dir = os.path.join(work_sapce_dir, "YOLOLabels/")
  69. if not os.path.isdir(yolo_labels_dir):
  70. os.mkdir(yolo_labels_dir)
  71. clear_hidden_files(yolo_labels_dir)
  72. yolov5_images_dir = os.path.join(data_base_dir, "images/")
  73. if not os.path.isdir(yolov5_images_dir):
  74. os.mkdir(yolov5_images_dir)
  75. clear_hidden_files(yolov5_images_dir)
  76. yolov5_labels_dir = os.path.join(data_base_dir, "labels/")
  77. if not os.path.isdir(yolov5_labels_dir):
  78. os.mkdir(yolov5_labels_dir)
  79. clear_hidden_files(yolov5_labels_dir)
  80. yolov5_images_train_dir = os.path.join(yolov5_images_dir, "train/")
  81. if not os.path.isdir(yolov5_images_train_dir):
  82. os.mkdir(yolov5_images_train_dir)
  83. clear_hidden_files(yolov5_images_train_dir)
  84. yolov5_images_test_dir = os.path.join(yolov5_images_dir, "val/")
  85. if not os.path.isdir(yolov5_images_test_dir):
  86. os.mkdir(yolov5_images_test_dir)
  87. clear_hidden_files(yolov5_images_test_dir)
  88. yolov5_labels_train_dir = os.path.join(yolov5_labels_dir, "train/")
  89. if not os.path.isdir(yolov5_labels_train_dir):
  90. os.mkdir(yolov5_labels_train_dir)
  91. clear_hidden_files(yolov5_labels_train_dir)
  92. yolov5_labels_test_dir = os.path.join(yolov5_labels_dir, "val/")
  93. if not os.path.isdir(yolov5_labels_test_dir):
  94. os.mkdir(yolov5_labels_test_dir)
  95. clear_hidden_files(yolov5_labels_test_dir)
  96. train_file = open(os.path.join(wd, "yolov5_train.txt"), 'w')
  97. test_file = open(os.path.join(wd, "yolov5_val.txt"), 'w')
  98. train_file.close()
  99. test_file.close()
  100. train_file = open(os.path.join(wd, "yolov5_train.txt"), 'a')
  101. test_file = open(os.path.join(wd, "yolov5_val.txt"), 'a')
  102. list_imgs = os.listdir(image_dir) # list image files
  103. prob = random.randint(1, 100)
  104. print("Probability: %d" % prob)
  105. for i in range(0,len(list_imgs)):
  106. path = os.path.join(image_dir,list_imgs[i])
  107. if os.path.isfile(path):
  108. image_path = image_dir + list_imgs[i]
  109. voc_path = list_imgs[i]
  110. (nameWithoutExtention, extention) = os.path.splitext(os.path.basename(image_path))
  111. (voc_nameWithoutExtention, voc_extention) = os.path.splitext(os.path.basename(voc_path))
  112. annotation_name = nameWithoutExtention + '.xml'
  113. annotation_path = os.path.join(annotation_dir, annotation_name)
  114. label_name = nameWithoutExtention + '.txt'
  115. label_path = os.path.join(yolo_labels_dir, label_name)
  116. prob = random.randint(1, 100)
  117. print("Probability: %d" % prob)
  118. if(prob < TRAIN_RATIO): # train dataset
  119. if os.path.exists(annotation_path):
  120. train_file.write(image_path + '\n')
  121. convert_annotation(nameWithoutExtention) # convert label
  122. copyfile(image_path, yolov5_images_train_dir + voc_path)
  123. copyfile(label_path, yolov5_labels_train_dir + label_name)
  124. else: # test dataset
  125. if os.path.exists(annotation_path):
  126. test_file.write(image_path + '\n')
  127. convert_annotation(nameWithoutExtention) # convert label
  128. copyfile(image_path, yolov5_images_test_dir + voc_path)
  129. copyfile(label_path, yolov5_labels_test_dir + label_name)
  130. train_file.close()
  131. test_file.close()

       在VOCdevkit目录下生成images和labels文件夹,文件夹下分别生成了train文件夹和val文件夹,里面分别保存着训练集的照片和txt格式的标签,还有验证集的照片和txt格式的标签。images文件夹和labels文件夹就是训练yolov5模型所需的训练集和验证集。

2.3 预训练权重

 为了缩短训练时间,可以从下面网址中获取预训练权重。https://www.wpsshop.cn/w/知新_RL/article/detail/732775

推荐阅读
相关标签
  

闽ICP备14008679号