当前位置:   article > 正文

yolov4目标检测与训练_训练集数量小于500,属于较小的数据量,请注意设置较大的训练世代(epoch)以满足足够

训练集数量小于500,属于较小的数据量,请注意设置较大的训练世代(epoch)以满足足够

前言

随着深度学习技术的不断发展,目标检测作为计算机视觉领域的重要分支,已经在实际应用中取得了显著的成果。YOLO(You Only Look Once)系列模型作为目标检测领域的经典之作,以其高效、快速的特点,赢得了广大研究者和开发者的青睐。YOLOv4更是在前代模型的基础上,通过一系列改进和创新,进一步提升了模型的性能,使得其在处理复杂场景和多样化目标时,表现出色。

本次实验旨在通过实际训练和测试YOLOv4模型,深入了解其内部机制、优化策略以及在实际应用中的表现。在实验中,将重点关注模型的训练过程、参数调整、性能评估等方面。我们将使用自己的标注数据集,通过合理的数据划分和预处理,为模型的训练提供优质的数据支持。最后,我们将对训练好的模型进行测试和评估,验证其在目标检测任务中的准确性和有效性,并展示检测结果。

一、标注数据集

下载labellmg

整理好需要标注的数据集,放在一个文件夹JPEJImages

进入到labellmg

保存好的标注数据集Annotations文件夹里面的数据都是xml类型。

二、处理VOC格式数据集

处理VOC格式数据集的脚本,目的是从数据集的XML标注文件中提取出目标的位置和类别信息,并统计每个类别的目标数量。

先创建一个包含数据集标注标签的class.txt文件。

在代码中导入依赖库、设置参数、获取类别列表、初始化统计变量、函数定义convert_annotation、主程序if_name_=="_main_"

以下是代码voc_annotation部分

  1. import os
  2. import random
  3. import xml.etree.ElementTree as ET
  4. import numpy as np
  5. from utils.utils import get_classes
  6. #--------------------------------------------------------------------------------------------------------------------------------#
  7. # annotation_mode用于指定该文件运行时计算的内容
  8. # annotation_mode为0代表整个标签处理过程,包括获得VOCdevkit/VOC2008/ImageSets里面的txt以及训练用的2008_train.txt、2008_val.txt
  9. # annotation_mode为1代表获得VOCdevkit/VOC2008/ImageSets里面的txt
  10. # annotation_mode为2代表获得训练用的2008_train.txt、2008_val.txt
  11. #--------------------------------------------------------------------------------------------------------------------------------#
  12. annotation_mode = 0
  13. #-------------------------------------------------------------------#
  14. # 必须要修改,用于生成2008_train.txt、2008_val.txt的目标信息
  15. # 与训练和预测所用的classes_path一致即可
  16. # 如果生成的2008_train.txt里面没有目标信息
  17. # 那么就是因为classes没有设定正确
  18. # 仅在annotation_mode为0和2的时候有效
  19. #-------------------------------------------------------------------#
  20. classes_path = 'model_data/class.txt'
  21. #--------------------------------------------------------------------------------------------------------------------------------#
  22. # trainval_percent用于指定(训练集+验证集)与测试集的比例,默认情况下 (训练集+验证集):测试集 = 9:1
  23. # train_percent用于指定(训练集+验证集)中训练集与验证集的比例,默认情况下 训练集:验证集 = 9:1
  24. # 仅在annotation_mode为0和1的时候有效
  25. #--------------------------------------------------------------------------------------------------------------------------------#
  26. trainval_percent = 0.9
  27. train_percent = 0.9
  28. #-------------------------------------------------------#
  29. # 指向VOC数据集所在的文件夹
  30. # 默认指向根目录下的VOC数据集
  31. #-------------------------------------------------------#
  32. VOCdevkit_path = 'VOCdevkit'
  33. VOCdevkit_sets = [('2008', 'train'), ('2008', 'val')]
  34. classes, _ = get_classes(classes_path)
  35. #-------------------------------------------------------#
  36. # 统计目标数量
  37. #-------------------------------------------------------#
  38. photo_nums = np.zeros(len(VOCdevkit_sets))
  39. nums = np.zeros(len(classes))
  40. def convert_annotation(year, image_id, list_file):
  41. in_file = open(os.path.join(VOCdevkit_path, 'VOC%s/Annotations/%s.xml'%(year, image_id)), encoding='utf-8')
  42. tree=ET.parse(in_file)
  43. root = tree.getroot()
  44. for obj in root.iter('object'):
  45. difficult = 0
  46. if obj.find('difficult')!=None:
  47. difficult = obj.find('difficult').text
  48. cls = obj.find('name').text
  49. if cls not in classes or int(difficult)==1:
  50. continue
  51. cls_id = classes.index(cls)
  52. xmlbox = obj.find('bndbox')
  53. b = (int(float(xmlbox.find('xmin').text)), int(float(xmlbox.find('ymin').text)), int(float(xmlbox.find('xmax').text)), int(float(xmlbox.find('ymax').text)))
  54. list_file.write(" " + ",".join([str(a) for a in b]) + ',' + str(cls_id))
  55. nums[classes.index(cls)] = nums[classes.index(cls)] + 1
  56. if __name__ == "__main__":
  57. random.seed(0)
  58. if " " in os.path.abspath(VOCdevkit_path):
  59. raise ValueError("数据集存放的文件夹路径与图片名称中不可以存在空格,否则会影响正常的模型训练,请注意修改。")
  60. if annotation_mode == 0 or annotation_mode == 1:
  61. print("Generate txt in ImageSets.")
  62. xmlfilepath = os.path.join(VOCdevkit_path, 'VOC2008/Annotations')
  63. saveBasePath = os.path.join(VOCdevkit_path, 'VOC2008/ImageSets/Main')
  64. temp_xml = os.listdir(xmlfilepath)
  65. total_xml = []
  66. for xml in temp_xml:
  67. if xml.endswith(".xml"):
  68. total_xml.append(xml)
  69. num = len(total_xml)
  70. list = range(num)
  71. tv = int(num*trainval_percent)
  72. tr = int(tv*train_percent)
  73. trainval= random.sample(list,tv)
  74. train = random.sample(trainval,tr)
  75. print("train and val size",tv)
  76. print("train size",tr)
  77. ftrainval = open(os.path.join(saveBasePath,'trainval.txt'), 'w')
  78. ftest = open(os.path.join(saveBasePath,'test.txt'), 'w')
  79. ftrain = open(os.path.join(saveBasePath,'train.txt'), 'w')
  80. fval = open(os.path.join(saveBasePath,'val.txt'), 'w')
  81. for i in list:
  82. name=total_xml[i][:-4]+'\n'
  83. if i in trainval:
  84. ftrainval.write(name)
  85. if i in train:
  86. ftrain.write(name)
  87. else:
  88. fval.write(name)
  89. else:
  90. ftest.write(name)
  91. ftrainval.close()
  92. ftrain.close()
  93. fval.close()
  94. ftest.close()
  95. print("Generate txt in ImageSets done.")
  96. if annotation_mode == 0 or annotation_mode == 2:
  97. print("Generate 2008_train.txt and 2008_val.txt for train.")
  98. type_index = 0
  99. for year, image_set in VOCdevkit_sets:
  100. image_ids = open(os.path.join(VOCdevkit_path, 'VOC%s/ImageSets/Main/%s.txt'%(year, image_set)), encoding='utf-8').read().strip().split()
  101. list_file = open('%s_%s.txt'%(year, image_set), 'w', encoding='utf-8')
  102. for image_id in image_ids:
  103. list_file.write('%s/VOC%s/JPEGImages/%s.jpg'%(os.path.abspath(VOCdevkit_path), year, image_id))
  104. convert_annotation(year, image_id, list_file)
  105. list_file.write('\n')
  106. photo_nums[type_index] = len(image_ids)
  107. type_index += 1
  108. list_file.close()
  109. print("Generate 2008_train.txt and 2008_val.txt for train done.")
  110. def printTable(List1, List2):
  111. for i in range(len(List1[0])):
  112. print("|", end=' ')
  113. for j in range(len(List1)):
  114. print(List1[j][i].rjust(int(List2[j])), end=' ')
  115. print("|", end=' ')
  116. print()
  117. str_nums = [str(int(x)) for x in nums]
  118. tableData = [
  119. classes, str_nums
  120. ]
  121. colWidths = [0]*len(tableData)
  122. len1 = 0
  123. for i in range(len(tableData)):
  124. for j in range(len(tableData[i])):
  125. if len(tableData[i][j]) > colWidths[i]:
  126. colWidths[i] = len(tableData[i][j])
  127. printTable(tableData, colWidths)
  128. if photo_nums[0] <= 500:
  129. print("训练集数量小于500,属于较小的数据量,请注意设置较大的训练世代(Epoch)以满足足够的梯度下降次数(Step)。")
  130. if np.sum(nums) == 0:
  131. print("在数据集中并未获得任何目标,请注意修改classes_path对应自己的数据集,并且保证标签名字正确,否则训练将会没有任何效果!")
  132. print("在数据集中并未获得任何目标,请注意修改classes_path对应自己的数据集,并且保证标签名字正确,否则训练将会没有任何效果!")
  133. print("在数据集中并未获得任何目标,请注意修改classes_path对应自己的数据集,并且保证标签名字正确,否则训练将会没有任何效果!")
  134. print("(重要的事情说三遍)。")

该代码的运行结果:

该代码为训练生成2008年的训练集(2008_train.txt)和验证集(2008_val.txt)的列表文件。从VOC数据集的现有标注文件中提取图像ID,并将这些ID对应的图像路径写入训练集(2008_train.txt)和验证集(2008_val.txt)的列表文件中,同时还对这些图像的标注进行了转换。这些列表文件用于模型训练时的数据加载。

三、模型训练

训练的代码main要放在VOCdevkit的同级目录。

主要负责设置训练所需的损失函数、日志记录、混合精度训练、模型配置和数据加载,迭代数据而得到最优的模型,以便后面的检测。监控训练过程中的损失变化,确保模型正常收敛。

以下是代码train部分

  1. #-------------------------------------#
  2. # 对数据集进行训练
  3. #-------------------------------------#
  4. import datetime
  5. import os
  6. os.environ['KMP_DUPLICATE_LIB_OK']='True'
  7. from functools import partial
  8. import numpy as np
  9. import torch
  10. import torch.backends.cudnn as cudnn
  11. import torch.distributed as dist
  12. import torch.nn as nn
  13. import torch.optim as optim
  14. from torch import nn
  15. from torch.utils.data import DataLoader
  16. from nets.yolo import YoloBody
  17. from nets.yolo_training import (YOLOLoss, get_lr_scheduler, set_optimizer_lr,
  18. weights_init)
  19. from utils.callbacks import EvalCallback, LossHistory
  20. from utils.dataloader import YoloDataset, yolo_dataset_collate
  21. from utils.utils import (get_anchors, get_classes, seed_everything,
  22. show_config, worker_init_fn)
  23. from utils.utils_fit import fit_one_epoch
  24. '''
  25. 训练自己的目标检测模型一定需要注意以下几点:
  26. 1、训练前仔细检查自己的格式是否满足要求,该库要求数据集格式为VOC格式,需要准备好的内容有输入图片和标签
  27. 输入图片为.jpg图片,无需固定大小,传入训练前会自动进行resize。
  28. 灰度图会自动转成RGB图片进行训练,无需自己修改。
  29. 输入图片如果后缀非jpg,需要自己批量转成jpg后再开始训练。
  30. 标签为.xml格式,文件中会有需要检测的目标信息,标签文件和输入图片文件相对应。
  31. 2、损失值的大小用于判断是否收敛,比较重要的是有收敛的趋势,即验证集损失不断下降,如果验证集损失基本上不改变的话,模型基本上就收敛了。
  32. 损失值的具体大小并没有什么意义,大和小只在于损失的计算方式,并不是接近于0才好。如果想要让损失好看点,可以直接到对应的损失函数里面除上10000。
  33. 训练过程中的损失值会保存在logs文件夹下的loss_%Y_%m_%d_%H_%M_%S文件夹中
  34. 3、训练好的权值文件保存在logs文件夹中,每个训练世代(Epoch)包含若干训练步长(Step),每个训练步长(Step)进行一次梯度下降。
  35. 如果只是训练了几个Step是不会保存的,Epoch和Step的概念要捋清楚一下。
  36. '''
  37. if __name__ == "__main__":
  38. #---------------------------------#
  39. # Cuda 是否使用Cuda
  40. # 没有GPU可以设置成False
  41. #---------------------------------#
  42. Cuda = False
  43. #----------------------------------------------#
  44. # Seed 用于固定随机种子
  45. # 使得每次独立训练都可以获得一样的结果
  46. #----------------------------------------------#
  47. seed = 11
  48. #---------------------------------------------------------------------#
  49. # distributed 用于指定是否使用单机多卡分布式运行
  50. # 终端指令仅支持Ubuntu。CUDA_VISIBLE_DEVICES用于在Ubuntu下指定显卡。
  51. # Windows系统下默认使用DP模式调用所有显卡,不支持DDP。
  52. # DP模式:
  53. # 设置 distributed = False
  54. # 在终端中输入 CUDA_VISIBLE_DEVICES=0,1 python train.py
  55. # DDP模式:
  56. # 设置 distributed = True
  57. # 在终端中输入 CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 train.py
  58. #---------------------------------------------------------------------#
  59. distributed = False
  60. #---------------------------------------------------------------------#
  61. # sync_bn 是否使用sync_bn,DDP模式多卡可用
  62. #---------------------------------------------------------------------#
  63. sync_bn = False
  64. #---------------------------------------------------------------------#
  65. # fp16 是否使用混合精度训练
  66. # 可减少约一半的显存、需要pytorch1.7.1以上
  67. #---------------------------------------------------------------------#
  68. fp16 = False
  69. #---------------------------------------------------------------------#
  70. # classes_path 指向model_data下的txt,与自己训练的数据集相关
  71. # 训练前一定要修改classes_path,使其对应自己的数据集
  72. #---------------------------------------------------------------------#
  73. classes_path = 'model_data/class.txt'
  74. #---------------------------------------------------------------------#
  75. # anchors_path 代表先验框对应的txt文件,一般不修改。
  76. # anchors_mask 用于帮助代码找到对应的先验框,一般不修改。
  77. #---------------------------------------------------------------------#
  78. anchors_path = 'model_data/yolo_anchors.txt'
  79. anchors_mask = [[6, 7, 8], [3, 4, 5], [0, 1, 2]]
  80. #----------------------------------------------------------------------------------------------------------------------------#
  81. # 权值文件的下载请看README,可以通过网盘下载。模型的 预训练权重 对不同数据集是通用的,因为特征是通用的。
  82. # 模型的 预训练权重 比较重要的部分是 主干特征提取网络的权值部分,用于进行特征提取。
  83. # 预训练权重对于99%的情况都必须要用,不用的话主干部分的权值太过随机,特征提取效果不明显,网络训练的结果也不会好
  84. #
  85. # 如果训练过程中存在中断训练的操作,可以将model_path设置成logs文件夹下的权值文件,将已经训练了一部分的权值再次载入。
  86. # 同时修改下方的 冻结阶段 或者 解冻阶段 的参数,来保证模型epoch的连续性。
  87. #
  88. # 当model_path = ''的时候不加载整个模型的权值。
  89. #
  90. # 此处使用的是整个模型的权重,因此是在train.py进行加载的,下面的pretrain不影响此处的权值加载。
  91. # 如果想要让模型从主干的预训练权值开始训练,则设置model_path = '',下面的pretrain = True,此时仅加载主干。
  92. # 如果想要让模型从0开始训练,则设置model_path = '',下面的pretrain = Fasle,Freeze_Train = Fasle,此时从0开始训练,且没有冻结主干的过程。
  93. #
  94. # 一般来讲,网络从0开始的训练效果会很差,因为权值太过随机,特征提取效果不明显,因此非常、非常、非常不建议大家从0开始训练!
  95. # 从0开始训练有两个方案:
  96. # 1、得益于Mosaic数据增强方法强大的数据增强能力,将UnFreeze_Epoch设置的较大(300及以上)、batch较大(16及以上)、数据较多(万以上)的情况下,
  97. # 可以设置mosaic=True,直接随机初始化参数开始训练,但得到的效果仍然不如有预训练的情况。(像COCO这样的大数据集可以这样做)
  98. # 2、了解imagenet数据集,首先训练分类模型,获得网络的主干部分权值,分类模型的 主干部分 和该模型通用,基于此进行训练。
  99. #----------------------------------------------------------------------------------------------------------------------------#
  100. model_path = 'model_data/yolo4_weights.pth'
  101. #------------------------------------------------------#
  102. # input_shape 输入的shape大小,一定要是32的倍数
  103. #------------------------------------------------------#
  104. input_shape = [416, 416]
  105. #----------------------------------------------------------------------------------------------------------------------------#
  106. # pretrained 是否使用主干网络的预训练权重,此处使用的是主干的权重,因此是在模型构建的时候进行加载的。
  107. # 如果设置了model_path,则主干的权值无需加载,pretrained的值无意义。
  108. # 如果不设置model_path,pretrained = True,此时仅加载主干开始训练。
  109. # 如果不设置model_path,pretrained = False,Freeze_Train = Fasle,此时从0开始训练,且没有冻结主干的过程。
  110. #----------------------------------------------------------------------------------------------------------------------------#
  111. pretrained = True
  112. #------------------------------------------------------------------#
  113. # mosaic 马赛克数据增强。
  114. # mosaic_prob 每个step有多少概率使用mosaic数据增强,默认50%。
  115. #
  116. # mixup 是否使用mixup数据增强,仅在mosaic=True时有效。
  117. # 只会对mosaic增强后的图片进行mixup的处理。
  118. # mixup_prob 有多少概率在mosaic后使用mixup数据增强,默认50%。
  119. # 总的mixup概率为mosaic_prob * mixup_prob。
  120. #
  121. # special_aug_ratio 参考YoloX,由于Mosaic生成的训练图片,远远脱离自然图片的真实分布。
  122. # 当mosaic=True时,本代码会在special_aug_ratio范围内开启mosaic。
  123. # 默认为前70%个epoch,100个世代会开启70个世代。
  124. #
  125. # 余弦退火算法的参数放到下面的lr_decay_type中设置
  126. #------------------------------------------------------------------#
  127. mosaic = True
  128. mosaic_prob = 0.5
  129. mixup = True
  130. mixup_prob = 0.5
  131. special_aug_ratio = 0.7
  132. #------------------------------------------------------------------#
  133. # label_smoothing 标签平滑。一般0.01以下。如0.01、0.005。
  134. #------------------------------------------------------------------#
  135. label_smoothing = 0
  136. #----------------------------------------------------------------------------------------------------------------------------#
  137. # 训练分为两个阶段,分别是冻结阶段和解冻阶段。设置冻结阶段是为了满足机器性能不足的同学的训练需求。
  138. # 冻结训练需要的显存较小,显卡非常差的情况下,可设置Freeze_Epoch等于UnFreeze_Epoch,此时仅仅进行冻结训练。
  139. #
  140. # 在此提供若干参数设置建议,各位训练者根据自己的需求进行灵活调整:
  141. # (一)从整个模型的预训练权重开始训练:
  142. # Adam:
  143. # Init_Epoch = 0,Freeze_Epoch = 50,UnFreeze_Epoch = 100,Freeze_Train = True,optimizer_type = 'adam',Init_lr = 1e-3,weight_decay = 0。(冻结)
  144. # Init_Epoch = 0,UnFreeze_Epoch = 100,Freeze_Train = False,optimizer_type = 'adam',Init_lr = 1e-3,weight_decay = 0。(不冻结)
  145. # SGD:
  146. # Init_Epoch = 0,Freeze_Epoch = 50,UnFreeze_Epoch = 300,Freeze_Train = True,optimizer_type = 'sgd',Init_lr = 1e-2,weight_decay = 5e-4。(冻结)
  147. # Init_Epoch = 0,UnFreeze_Epoch = 300,Freeze_Train = False,optimizer_type = 'sgd',Init_lr = 1e-2,weight_decay = 5e-4。(不冻结)
  148. # 其中:UnFreeze_Epoch可以在100-300之间调整。
  149. # (二)从主干网络的预训练权重开始训练:
  150. # Adam:
  151. # Init_Epoch = 0,Freeze_Epoch = 50,UnFreeze_Epoch = 100,Freeze_Train = True,optimizer_type = 'adam',Init_lr = 1e-3,weight_decay = 0。(冻结)
  152. # Init_Epoch = 0,UnFreeze_Epoch = 100,Freeze_Train = False,optimizer_type = 'adam',Init_lr = 1e-3,weight_decay = 0。(不冻结)
  153. # SGD:
  154. # Init_Epoch = 0,Freeze_Epoch = 50,UnFreeze_Epoch = 300,Freeze_Train = True,optimizer_type = 'sgd',Init_lr = 1e-2,weight_decay = 5e-4。(冻结)
  155. # Init_Epoch = 0,UnFreeze_Epoch = 300,Freeze_Train = False,optimizer_type = 'sgd',Init_lr = 1e-2,weight_decay = 5e-4。(不冻结)
  156. # 其中:由于从主干网络的预训练权重开始训练,主干的权值不一定适合目标检测,需要更多的训练跳出局部最优解。
  157. # UnFreeze_Epoch可以在150-300之间调整,YOLOV5和YOLOX均推荐使用300。
  158. # Adam相较于SGD收敛的快一些。因此UnFreeze_Epoch理论上可以小一点,但依然推荐更多的Epoch。
  159. # (三)从0开始训练:
  160. # Init_Epoch = 0,UnFreeze_Epoch >= 300,Unfreeze_batch_size >= 16,Freeze_Train = False(不冻结训练)
  161. # 其中:UnFreeze_Epoch尽量不小于300。optimizer_type = 'sgd',Init_lr = 1e-2,mosaic = True。
  162. # (四)batch_size的设置:
  163. # 在显卡能够接受的范围内,以大为好。显存不足与数据集大小无关,提示显存不足(OOM或者CUDA out of memory)请调小batch_size。
  164. # 受到BatchNorm层影响,batch_size最小为2,不能为1。
  165. # 正常情况下Freeze_batch_size建议为Unfreeze_batch_size的1-2倍。不建议设置的差距过大,因为关系到学习率的自动调整。
  166. #----------------------------------------------------------------------------------------------------------------------------#
  167. #------------------------------------------------------------------#
  168. # 冻结阶段训练参数
  169. # 此时模型的主干被冻结了,特征提取网络不发生改变
  170. # 占用的显存较小,仅对网络进行微调
  171. # Init_Epoch 模型当前开始的训练世代,其值可以大于Freeze_Epoch,如设置:
  172. # Init_Epoch = 60、Freeze_Epoch = 50、UnFreeze_Epoch = 100
  173. # 会跳过冻结阶段,直接从60代开始,并调整对应的学习率。
  174. # (断点续练时使用)
  175. # Freeze_Epoch 模型冻结训练的Freeze_Epoch
  176. # (当Freeze_Train=False时失效)
  177. # Freeze_batch_size 模型冻结训练的batch_size
  178. # (当Freeze_Train=False时失效)
  179. #------------------------------------------------------------------#
  180. Init_Epoch = 0
  181. Freeze_Epoch = 20
  182. Freeze_batch_size = 2
  183. #------------------------------------------------------------------#
  184. # 解冻阶段训练参数
  185. # 此时模型的主干不被冻结了,特征提取网络会发生改变
  186. # 占用的显存较大,网络所有的参数都会发生改变
  187. # UnFreeze_Epoch 模型总共训练的epoch
  188. # SGD需要更长的时间收敛,因此设置较大的UnFreeze_Epoch
  189. # Adam可以使用相对较小的UnFreeze_Epoch
  190. # Unfreeze_batch_size 模型在解冻后的batch_size
  191. #------------------------------------------------------------------#
  192. UnFreeze_Epoch = 20
  193. Unfreeze_batch_size = 2
  194. #------------------------------------------------------------------#
  195. # Freeze_Train 是否进行冻结训练
  196. # 默认先冻结主干训练后解冻训练。
  197. #------------------------------------------------------------------#
  198. Freeze_Train = True
  199. #------------------------------------------------------------------#
  200. # 其它训练参数:学习率、优化器、学习率下降有关
  201. #------------------------------------------------------------------#
  202. #------------------------------------------------------------------#
  203. # Init_lr 模型的最大学习率
  204. # Min_lr 模型的最小学习率,默认为最大学习率的0.01
  205. #------------------------------------------------------------------#
  206. Init_lr = 1e-2
  207. Min_lr = Init_lr * 0.01
  208. #------------------------------------------------------------------#
  209. # optimizer_type 使用到的优化器种类,可选的有adam、sgd
  210. # 当使用Adam优化器时建议设置 Init_lr=1e-3
  211. # 当使用SGD优化器时建议设置 Init_lr=1e-2
  212. # momentum 优化器内部使用到的momentum参数
  213. # weight_decay 权值衰减,可防止过拟合
  214. # adam会导致weight_decay错误,使用adam时建议设置为0。
  215. #------------------------------------------------------------------#
  216. optimizer_type = "sgd"
  217. momentum = 0.937
  218. weight_decay = 5e-4
  219. #------------------------------------------------------------------#
  220. # lr_decay_type 使用到的学习率下降方式,可选的有step、cos
  221. #------------------------------------------------------------------#
  222. lr_decay_type = "cos"
  223. #------------------------------------------------------------------#
  224. # focal_loss 是否使用Focal Loss平衡正负样本
  225. # focal_alpha Focal Loss的正负样本平衡参数
  226. # focal_gamma Focal Loss的难易分类样本平衡参数
  227. #------------------------------------------------------------------#
  228. focal_loss = False
  229. focal_alpha = 0.25
  230. focal_gamma = 2
  231. #------------------------------------------------------------------#
  232. # iou_type 使用什么iou损失,ciou或者siou
  233. #------------------------------------------------------------------#
  234. iou_type = 'ciou'
  235. #------------------------------------------------------------------#
  236. # save_period 多少个epoch保存一次权值
  237. #------------------------------------------------------------------#
  238. save_period = 10
  239. #------------------------------------------------------------------#
  240. # save_dir 权值与日志文件保存的文件夹
  241. #------------------------------------------------------------------#
  242. save_dir = 'logs'
  243. #------------------------------------------------------------------#
  244. # eval_flag 是否在训练时进行评估,评估对象为验证集
  245. # 安装pycocotools库后,评估体验更佳。
  246. # eval_period 代表多少个epoch评估一次,不建议频繁的评估
  247. # 评估需要消耗较多的时间,频繁评估会导致训练非常慢
  248. # 此处获得的mAP会与get_map.py获得的会有所不同,原因有二:
  249. # (一)此处获得的mAP为验证集的mAP。
  250. # (二)此处设置评估参数较为保守,目的是加快评估速度。
  251. #------------------------------------------------------------------#
  252. eval_flag = True
  253. eval_period = 10
  254. #------------------------------------------------------------------#
  255. # num_workers 用于设置是否使用多线程读取数据
  256. # 开启后会加快数据读取速度,但是会占用更多内存
  257. # 内存较小的电脑可以设置为2或者0
  258. #------------------------------------------------------------------#
  259. num_workers = 4
  260. #------------------------------------------------------#
  261. # train_annotation_path 训练图片路径和标签
  262. # val_annotation_path 验证图片路径和标签
  263. #------------------------------------------------------#
  264. train_annotation_path = '2008_train.txt'
  265. val_annotation_path = '2008_val.txt'
  266. seed_everything(seed)
  267. #------------------------------------------------------#
  268. # 设置用到的显卡
  269. #------------------------------------------------------#
  270. ngpus_per_node = torch.cuda.device_count()
  271. if distributed:
  272. dist.init_process_group(backend="nccl")
  273. local_rank = int(os.environ["LOCAL_RANK"])
  274. rank = int(os.environ["RANK"])
  275. device = torch.device("cuda", local_rank)
  276. if local_rank == 0:
  277. print(f"[{os.getpid()}] (rank = {rank}, local_rank = {local_rank}) training...")
  278. print("Gpu Device Count : ", ngpus_per_node)
  279. else:
  280. device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  281. local_rank = 0
  282. rank = 0
  283. #------------------------------------------------------#
  284. # 获取classes和anchor
  285. #------------------------------------------------------#
  286. class_names, num_classes = get_classes(classes_path)
  287. anchors, num_anchors = get_anchors(anchors_path)
  288. #------------------------------------------------------#
  289. # 创建yolo模型
  290. #------------------------------------------------------#
  291. model = YoloBody(anchors_mask, num_classes, pretrained = pretrained)
  292. if not pretrained:
  293. weights_init(model)
  294. if model_path != '':
  295. #------------------------------------------------------#
  296. # 权值文件请看README,百度网盘下载
  297. #------------------------------------------------------#
  298. if local_rank == 0:
  299. print('Load weights {}.'.format(model_path))
  300. #------------------------------------------------------#
  301. # 根据预训练权重的Key和模型的Key进行加载
  302. #------------------------------------------------------#
  303. model_dict = model.state_dict()
  304. pretrained_dict = torch.load(model_path, map_location = device)
  305. load_key, no_load_key, temp_dict = [], [], {}
  306. for k, v in pretrained_dict.items():
  307. if k in model_dict.keys() and np.shape(model_dict[k]) == np.shape(v):
  308. temp_dict[k] = v
  309. load_key.append(k)
  310. else:
  311. no_load_key.append(k)
  312. model_dict.update(temp_dict)
  313. model.load_state_dict(model_dict)
  314. #------------------------------------------------------#
  315. # 显示没有匹配上的Key
  316. #------------------------------------------------------#
  317. if local_rank == 0:
  318. print("\nSuccessful Load Key:", str(load_key)[:500], "……\nSuccessful Load Key Num:", len(load_key))
  319. print("\nFail To Load Key:", str(no_load_key)[:500], "……\nFail To Load Key num:", len(no_load_key))
  320. print("\n\033[1;33;44m温馨提示,head部分没有载入是正常现象,Backbone部分没有载入是错误的。\033[0m")
  321. #----------------------#
  322. # 获得损失函数
  323. #----------------------#
  324. yolo_loss = YOLOLoss(anchors, num_classes, input_shape, Cuda, anchors_mask, label_smoothing, focal_loss, focal_alpha, focal_gamma, iou_type)
  325. #----------------------#
  326. # 记录Loss
  327. #----------------------#
  328. if local_rank == 0:
  329. time_str = datetime.datetime.strftime(datetime.datetime.now(),'%Y_%m_%d_%H_%M_%S')
  330. log_dir = os.path.join(save_dir, "loss_" + str(time_str))
  331. loss_history = LossHistory(log_dir, model, input_shape=input_shape)
  332. else:
  333. loss_history = None
  334. #------------------------------------------------------------------#
  335. # torch 1.2不支持amp,建议使用torch 1.7.1及以上正确使用fp16
  336. # 因此torch1.2这里显示"could not be resolve"
  337. #------------------------------------------------------------------#
  338. if fp16:
  339. from torch.cuda.amp import GradScaler as GradScaler
  340. scaler = GradScaler()
  341. else:
  342. scaler = None
  343. model_train = model.train()
  344. #----------------------------#
  345. # 多卡同步Bn
  346. #----------------------------#
  347. if sync_bn and ngpus_per_node > 1 and distributed:
  348. model_train = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model_train)
  349. elif sync_bn:
  350. print("Sync_bn is not support in one gpu or not distributed.")
  351. if Cuda:
  352. if distributed:
  353. #----------------------------#
  354. # 多卡平行运行
  355. #----------------------------#
  356. model_train = model_train.cuda(local_rank)
  357. model_train = torch.nn.parallel.DistributedDataParallel(model_train, device_ids=[local_rank], find_unused_parameters=True)
  358. else:
  359. model_train = torch.nn.DataParallel(model)
  360. cudnn.benchmark = True
  361. model_train = model_train.cuda()
  362. #---------------------------#
  363. # 读取数据集对应的txt
  364. #---------------------------#
  365. with open(train_annotation_path, encoding='utf-8') as f:
  366. train_lines = f.readlines()
  367. with open(val_annotation_path, encoding='utf-8') as f:
  368. val_lines = f.readlines()
  369. num_train = len(train_lines)
  370. num_val = len(val_lines)
  371. if local_rank == 0:
  372. show_config(
  373. classes_path = classes_path, anchors_path = anchors_path, anchors_mask = anchors_mask, model_path = model_path, input_shape = input_shape, \
  374. Init_Epoch = Init_Epoch, Freeze_Epoch = Freeze_Epoch, UnFreeze_Epoch = UnFreeze_Epoch, Freeze_batch_size = Freeze_batch_size, Unfreeze_batch_size = Unfreeze_batch_size, Freeze_Train = Freeze_Train, \
  375. Init_lr = Init_lr, Min_lr = Min_lr, optimizer_type = optimizer_type, momentum = momentum, lr_decay_type = lr_decay_type, \
  376. save_period = save_period, save_dir = save_dir, num_workers = num_workers, num_train = num_train, num_val = num_val
  377. )
  378. #---------------------------------------------------------#
  379. # 总训练世代指的是遍历全部数据的总次数
  380. # 总训练步长指的是梯度下降的总次数
  381. # 每个训练世代包含若干训练步长,每个训练步长进行一次梯度下降。
  382. # 此处仅建议最低训练世代,上不封顶,计算时只考虑了解冻部分
  383. #----------------------------------------------------------#
  384. wanted_step = 5e4 if optimizer_type == "sgd" else 1.5e4
  385. total_step = num_train // Unfreeze_batch_size * UnFreeze_Epoch
  386. if total_step <= wanted_step:
  387. if num_train // Unfreeze_batch_size == 0:
  388. raise ValueError('数据集过小,无法进行训练,请扩充数据集。')
  389. wanted_epoch = wanted_step // (num_train // Unfreeze_batch_size) + 1
  390. print("\n\033[1;33;44m[Warning] 使用%s优化器时,建议将训练总步长设置到%d以上。\033[0m"%(optimizer_type, wanted_step))
  391. print("\033[1;33;44m[Warning] 本次运行的总训练数据量为%d,Unfreeze_batch_size为%d,共训练%d个Epoch,计算出总训练步长为%d。\033[0m"%(num_train, Unfreeze_batch_size, UnFreeze_Epoch, total_step))
  392. print("\033[1;33;44m[Warning] 由于总训练步长为%d,小于建议总步长%d,建议设置总世代为%d。\033[0m"%(total_step, wanted_step, wanted_epoch))
  393. #------------------------------------------------------#
  394. # 主干特征提取网络特征通用,冻结训练可以加快训练速度
  395. # 也可以在训练初期防止权值被破坏。
  396. # Init_Epoch为起始世代
  397. # Freeze_Epoch为冻结训练的世代
  398. # UnFreeze_Epoch总训练世代
  399. # 提示OOM或者显存不足请调小Batch_size
  400. #------------------------------------------------------#
  401. if True:
  402. UnFreeze_flag = False
  403. #------------------------------------#
  404. # 冻结一定部分训练
  405. #------------------------------------#
  406. if Freeze_Train:
  407. for param in model.backbone.parameters():
  408. param.requires_grad = False
  409. #-------------------------------------------------------------------#
  410. # 如果不冻结训练的话,直接设置batch_size为Unfreeze_batch_size
  411. #-------------------------------------------------------------------#
  412. batch_size = Freeze_batch_size if Freeze_Train else Unfreeze_batch_size
  413. #-------------------------------------------------------------------#
  414. # 判断当前batch_size,自适应调整学习率
  415. #-------------------------------------------------------------------#
  416. nbs = 64
  417. lr_limit_max = 1e-3 if optimizer_type in ['adam', 'adamw'] else 5e-2
  418. lr_limit_min = 3e-4 if optimizer_type in ['adam', 'adamw'] else 5e-4
  419. Init_lr_fit = min(max(batch_size / nbs * Init_lr, lr_limit_min), lr_limit_max)
  420. Min_lr_fit = min(max(batch_size / nbs * Min_lr, lr_limit_min * 1e-2), lr_limit_max * 1e-2)
  421. #---------------------------------------#
  422. # 根据optimizer_type选择优化器
  423. #---------------------------------------#
  424. pg0, pg1, pg2 = [], [], []
  425. for k, v in model.named_modules():
  426. if hasattr(v, "bias") and isinstance(v.bias, nn.Parameter):
  427. pg2.append(v.bias)
  428. if isinstance(v, nn.BatchNorm2d) or "bn" in k:
  429. pg0.append(v.weight)
  430. elif hasattr(v, "weight") and isinstance(v.weight, nn.Parameter):
  431. pg1.append(v.weight)
  432. optimizer = {
  433. 'adam' : optim.Adam(pg0, Init_lr_fit, betas = (momentum, 0.999)),
  434. 'adamw' : optim.AdamW(pg0, Init_lr_fit, betas = (momentum, 0.999)),
  435. 'sgd' : optim.SGD(pg0, Init_lr_fit, momentum = momentum, nesterov=True)
  436. }[optimizer_type]
  437. optimizer.add_param_group({"params": pg1, "weight_decay": weight_decay})
  438. optimizer.add_param_group({"params": pg2})
  439. #---------------------------------------#
  440. # 获得学习率下降的公式
  441. #---------------------------------------#
  442. lr_scheduler_func = get_lr_scheduler(lr_decay_type, Init_lr_fit, Min_lr_fit, UnFreeze_Epoch)
  443. #---------------------------------------#
  444. # 判断每一个世代的长度
  445. #---------------------------------------#
  446. epoch_step = num_train // batch_size
  447. epoch_step_val = num_val // batch_size
  448. if epoch_step == 0 or epoch_step_val == 0:
  449. raise ValueError("数据集过小,无法继续进行训练,请扩充数据集。")
  450. #---------------------------------------#
  451. # 构建数据集加载器。
  452. #---------------------------------------#
  453. train_dataset = YoloDataset(train_lines, input_shape, num_classes, epoch_length = UnFreeze_Epoch, \
  454. mosaic=mosaic, mixup=mixup, mosaic_prob=mosaic_prob, mixup_prob=mixup_prob, train=True, special_aug_ratio=special_aug_ratio)
  455. val_dataset = YoloDataset(val_lines, input_shape, num_classes, epoch_length = UnFreeze_Epoch, \
  456. mosaic=False, mixup=False, mosaic_prob=0, mixup_prob=0, train=False, special_aug_ratio=0)
  457. if distributed:
  458. train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset, shuffle=True,)
  459. val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset, shuffle=False,)
  460. batch_size = batch_size // ngpus_per_node
  461. shuffle = False
  462. else:
  463. train_sampler = None
  464. val_sampler = None
  465. shuffle = True
  466. gen = DataLoader(train_dataset, shuffle = shuffle, batch_size = batch_size, num_workers = num_workers, pin_memory=True,
  467. drop_last=True, collate_fn=yolo_dataset_collate, sampler=train_sampler,
  468. worker_init_fn=partial(worker_init_fn, rank=rank, seed=seed))
  469. gen_val = DataLoader(val_dataset , shuffle = shuffle, batch_size = batch_size, num_workers = num_workers, pin_memory=True,
  470. drop_last=True, collate_fn=yolo_dataset_collate, sampler=val_sampler,
  471. worker_init_fn=partial(worker_init_fn, rank=rank, seed=seed))
  472. #----------------------#
  473. # 记录eval的map曲线
  474. #----------------------#
  475. if local_rank == 0:
  476. eval_callback = EvalCallback(model, input_shape, anchors, anchors_mask, class_names, num_classes, val_lines, log_dir, Cuda, \
  477. eval_flag=eval_flag, period=eval_period)
  478. else:
  479. eval_callback = None
  480. #---------------------------------------#
  481. # 开始模型训练
  482. #---------------------------------------#
  483. for epoch in range(Init_Epoch, UnFreeze_Epoch):
  484. #---------------------------------------#
  485. # 如果模型有冻结学习部分
  486. # 则解冻,并设置参数
  487. #---------------------------------------#
  488. if epoch >= Freeze_Epoch and not UnFreeze_flag and Freeze_Train:
  489. batch_size = Unfreeze_batch_size
  490. #-------------------------------------------------------------------#
  491. # 判断当前batch_size,自适应调整学习率
  492. #-------------------------------------------------------------------#
  493. nbs = 64
  494. lr_limit_max = 1e-3 if optimizer_type in ['adam', 'adamw'] else 5e-2
  495. lr_limit_min = 3e-4 if optimizer_type in ['adam', 'adamw'] else 5e-4
  496. Init_lr_fit = min(max(batch_size / nbs * Init_lr, lr_limit_min), lr_limit_max)
  497. Min_lr_fit = min(max(batch_size / nbs * Min_lr, lr_limit_min * 1e-2), lr_limit_max * 1e-2)
  498. #---------------------------------------#
  499. # 获得学习率下降的公式
  500. #---------------------------------------#
  501. lr_scheduler_func = get_lr_scheduler(lr_decay_type, Init_lr_fit, Min_lr_fit, UnFreeze_Epoch)
  502. for param in model.backbone.parameters():
  503. param.requires_grad = True
  504. epoch_step = num_train // batch_size
  505. epoch_step_val = num_val // batch_size
  506. if epoch_step == 0 or epoch_step_val == 0:
  507. raise ValueError("数据集过小,无法继续进行训练,请扩充数据集。")
  508. if distributed:
  509. batch_size = batch_size // ngpus_per_node
  510. gen = DataLoader(train_dataset, shuffle = shuffle, batch_size = batch_size, num_workers = num_workers, pin_memory=True,
  511. drop_last=True, collate_fn=yolo_dataset_collate, sampler=train_sampler,
  512. worker_init_fn=partial(worker_init_fn, rank=rank, seed=seed))
  513. gen_val = DataLoader(val_dataset , shuffle = shuffle, batch_size = batch_size, num_workers = num_workers, pin_memory=True,
  514. drop_last=True, collate_fn=yolo_dataset_collate, sampler=val_sampler,
  515. worker_init_fn=partial(worker_init_fn, rank=rank, seed=seed))
  516. UnFreeze_flag = True
  517. gen.dataset.epoch_now = epoch
  518. gen_val.dataset.epoch_now = epoch
  519. if distributed:
  520. train_sampler.set_epoch(epoch)
  521. set_optimizer_lr(optimizer, lr_scheduler_func, epoch)
  522. fit_one_epoch(model_train, model, yolo_loss, loss_history, eval_callback, optimizer, epoch, epoch_step, epoch_step_val, gen, gen_val, UnFreeze_Epoch, Cuda, fp16, scaler, save_period, save_dir, local_rank)
  523. if distributed:
  524. dist.barrier()
  525. if local_rank == 0:
  526. loss_history.writer.close()

这里迭代20次,得到一个最优的模型并保存在best_epoch_weights.pth文件。

该代码的运行结果:

四、检测图片

用于执行不同对象检测任务的脚本,它允许用户根据所选模式执行不同的操作。

使用训练好的模型对测试图片进行对象检测,将检测结果可视化,并展示在测试图片上。

以下是代码predict部分

  1. #-----------------------------------------------------------------------#
  2. # predict.py将单张图片预测、摄像头检测、FPS测试和目录遍历检测等功能
  3. # 整合到了一个py文件中,通过指定mode进行模式的修改。
  4. #-----------------------------------------------------------------------#
  5. import time
  6. import cv2
  7. import numpy as np
  8. from PIL import Image
  9. from yolo import YOLO, YOLO_ONNX
  10. if __name__ == "__main__":
  11. #----------------------------------------------------------------------------------------------------------#
  12. # mode用于指定测试的模式:
  13. # 'predict' 表示单张图片预测,如果想对预测过程进行修改,如保存图片,截取对象等,可以先看下方详细的注释
  14. # 'video' 表示视频检测,可调用摄像头或者视频进行检测,详情查看下方注释。
  15. # 'fps' 表示测试fps,使用的图片是img里面的street.jpg,详情查看下方注释。
  16. # 'dir_predict' 表示遍历文件夹进行检测并保存。默认遍历img文件夹,保存img_out文件夹,详情查看下方注释。
  17. # 'heatmap' 表示进行预测结果的热力图可视化,详情查看下方注释。
  18. # 'export_onnx' 表示将模型导出为onnx,需要pytorch1.7.1以上。
  19. # 'predict_onnx' 表示利用导出的onnx模型进行预测,相关参数的修改在yolo.py_416行左右处的YOLO_ONNX
  20. #----------------------------------------------------------------------------------------------------------#
  21. mode = "predict"
  22. #-------------------------------------------------------------------------#
  23. # crop 指定了是否在单张图片预测后对目标进行截取
  24. # count 指定了是否进行目标的计数
  25. # crop、count仅在mode='predict'时有效
  26. #-------------------------------------------------------------------------#
  27. crop = False
  28. count = False
  29. #----------------------------------------------------------------------------------------------------------#
  30. # video_path 用于指定视频的路径,当video_path=0时表示检测摄像头
  31. # 想要检测视频,则设置如video_path = "xxx.mp4"即可,代表读取出根目录下的xxx.mp4文件。
  32. # video_save_path 表示视频保存的路径,当video_save_path=""时表示不保存
  33. # 想要保存视频,则设置如video_save_path = "yyy.mp4"即可,代表保存为根目录下的yyy.mp4文件。
  34. # video_fps 用于保存的视频的fps
  35. #
  36. # video_path、video_save_path和video_fps仅在mode='video'时有效
  37. # 保存视频时需要ctrl+c退出或者运行到最后一帧才会完成完整的保存步骤。
  38. #----------------------------------------------------------------------------------------------------------#
  39. video_path = 0
  40. video_save_path = ""
  41. video_fps = 25.0
  42. #----------------------------------------------------------------------------------------------------------#
  43. # test_interval 用于指定测量fps的时候,图片检测的次数。理论上test_interval越大,fps越准确。
  44. # fps_image_path 用于指定测试的fps图片
  45. #
  46. # test_interval和fps_image_path仅在mode='fps'有效
  47. #----------------------------------------------------------------------------------------------------------#
  48. test_interval = 100
  49. fps_image_path = "img/street.jpg"
  50. #-------------------------------------------------------------------------#
  51. # dir_origin_path 指定了用于检测的图片的文件夹路径
  52. # dir_save_path 指定了检测完图片的保存路径
  53. #
  54. # dir_origin_path和dir_save_path仅在mode='dir_predict'时有效
  55. #-------------------------------------------------------------------------#
  56. dir_origin_path = "img/"
  57. dir_save_path = "img_out/"
  58. #-------------------------------------------------------------------------#
  59. # heatmap_save_path 热力图的保存路径,默认保存在model_data下
  60. #
  61. # heatmap_save_path仅在mode='heatmap'有效
  62. #-------------------------------------------------------------------------#
  63. heatmap_save_path = "model_data/heatmap_vision.png"
  64. #-------------------------------------------------------------------------#
  65. # simplify 使用Simplify onnx
  66. # onnx_save_path 指定了onnx的保存路径
  67. #-------------------------------------------------------------------------#
  68. simplify = True
  69. onnx_save_path = "model_data/models.onnx"
  70. if mode != "predict_onnx":
  71. yolo = YOLO()
  72. else:
  73. yolo = YOLO_ONNX()
  74. if mode == "predict":
  75. '''
  76. 1、如果想要进行检测完的图片的保存,利用r_image.save("img.jpg")即可保存,直接在predict.py里进行修改即可。
  77. 2、如果想要获得预测框的坐标,可以进入yolo.detect_image函数,在绘图部分读取top,left,bottom,right这四个值。
  78. 3、如果想要利用预测框截取下目标,可以进入yolo.detect_image函数,在绘图部分利用获取到的top,left,bottom,right这四个值
  79. 在原图上利用矩阵的方式进行截取。
  80. 4、如果想要在预测图上写额外的字,比如检测到的特定目标的数量,可以进入yolo.detect_image函数,在绘图部分对predicted_class进行判断,
  81. 比如判断if predicted_class == 'car': 即可判断当前目标是否为车,然后记录数量即可。利用draw.text即可写字。
  82. '''
  83. while True:
  84. img = input('Input image filename:')
  85. try:
  86. image = Image.open(img)
  87. except:
  88. print('Open Error! Try again!')
  89. continue
  90. else:
  91. r_image = yolo.detect_image(image, crop = crop, count=count)
  92. r_image.show()
  93. elif mode == "video":
  94. capture = cv2.VideoCapture(video_path)
  95. if video_save_path!="":
  96. fourcc = cv2.VideoWriter_fourcc(*'XVID')
  97. size = (int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)), int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)))
  98. out = cv2.VideoWriter(video_save_path, fourcc, video_fps, size)
  99. ref, frame = capture.read()
  100. if not ref:
  101. raise ValueError("未能正确读取摄像头(视频),请注意是否正确安装摄像头(是否正确填写视频路径)。")
  102. fps = 0.0
  103. while(True):
  104. t1 = time.time()
  105. # 读取某一帧
  106. ref, frame = capture.read()
  107. if not ref:
  108. break
  109. # 格式转变,BGRtoRGB
  110. frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)
  111. # 转变成Image
  112. frame = Image.fromarray(np.uint8(frame))
  113. # 进行检测
  114. frame = np.array(yolo.detect_image(frame))
  115. # RGBtoBGR满足opencv显示格式
  116. frame = cv2.cvtColor(frame,cv2.COLOR_RGB2BGR)
  117. fps = ( fps + (1./(time.time()-t1)) ) / 2
  118. print("fps= %.2f"%(fps))
  119. frame = cv2.putText(frame, "fps= %.2f"%(fps), (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
  120. cv2.imshow("video",frame)
  121. c= cv2.waitKey(1) & 0xff
  122. if video_save_path!="":
  123. out.write(frame)
  124. if c==27:
  125. capture.release()
  126. break
  127. print("Video Detection Done!")
  128. capture.release()
  129. if video_save_path!="":
  130. print("Save processed video to the path :" + video_save_path)
  131. out.release()
  132. cv2.destroyAllWindows()
  133. elif mode == "fps":
  134. img = Image.open(fps_image_path)
  135. tact_time = yolo.get_FPS(img, test_interval)
  136. print(str(tact_time) + ' seconds, ' + str(1/tact_time) + 'FPS, @batch_size 1')
  137. elif mode == "dir_predict":
  138. import os
  139. from tqdm import tqdm
  140. img_names = os.listdir(dir_origin_path)
  141. for img_name in tqdm(img_names):
  142. if img_name.lower().endswith(('.bmp', '.dib', '.png', '.jpg', '.jpeg', '.pbm', '.pgm', '.ppm', '.tif', '.tiff')):
  143. image_path = os.path.join(dir_origin_path, img_name)
  144. image = Image.open(image_path)
  145. r_image = yolo.detect_image(image)
  146. if not os.path.exists(dir_save_path):
  147. os.makedirs(dir_save_path)
  148. r_image.save(os.path.join(dir_save_path, img_name.replace(".jpg", ".png")), quality=95, subsampling=0)
  149. elif mode == "heatmap":
  150. while True:
  151. img = input('Input image filename:')
  152. try:
  153. image = Image.open(img)
  154. except:
  155. print('Open Error! Try again!')
  156. continue
  157. else:
  158. yolo.detect_heatmap(image, heatmap_save_path)
  159. elif mode == "export_onnx":
  160. yolo.convert_to_onnx(simplify, onnx_save_path)
  161. elif mode == "predict_onnx":
  162. while True:
  163. img = input('Input image filename:')
  164. try:
  165. image = Image.open(img)
  166. except:
  167. print('Open Error! Try again!')
  168. continue
  169. else:
  170. r_image = yolo.detect_image(image)
  171. r_image.show()
  172. else:
  173. raise AssertionError("Please specify the correct mode: 'predict', 'video', 'fps', 'heatmap', 'export_onnx', 'dir_predict'.")

根据调节confidence的参数来显示正确框出之前标注的数据。

该代码的运行结果:

总结

本实验的主要目标是使用YOLOv4算法,通过自己标注的数据集进行训练,得到一个训练好的模型。随后,我们将使用这个训练好的模型对测试图片进行对象检测,并展示检测结果。

经过训练,我们得到了一个针对特定对象的YOLOv4模型。在测试阶段,该模型能够准确地检测出测试图片中的对象,并给出相应的类别和位置信息。通过可视化展示,我们可以直观地看到模型的检测效果。

本实验仍存在一定的局限性。例如,在处理复杂背景和遮挡目标时,模型的检测性能可能会受到一定影响。未来,将进一步探索如何提升模型在复杂场景下的检测能力,并尝试将模型应用于更多的实际场景中。

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

闽ICP备14008679号