快速使用Yolov4-tiny训练VOC格式的数据集

jupiter
2021-11-24 / 0 评论 / 845 阅读 / 正在检测是否收录...
温馨提示:
本文最后更新于2022年11月22日,已超过514天没有更新,若内容或图片失效,请留言反馈。

1.训练步骤

STEP1:下载代码并配置环境

git clone https://github.com/bubbliiiing/yolov4-tiny-pytorch.git
cd yolov4-tiny-pytorch
pip install -r requirements.txt

STEP2:根据文件结构填充VOC格式的数据集

  • 数据放置格式(只需完成#TODO部分即可)
├──VOCdevkit/VOC2007/
    ├── Annotations #可以采用软连接的方式避免对大量数据进行复制
        ├──放置xml文件 #TODO
    ├── JPEGImages
        ├──放置img文件 #TODO
    ├──ImageSets/Main #可以采用软连接的方式避免对大量数据进行复制
        ├──放置训练索引文件 (无需手动完成,自动生成)
    ├── voc2yolo.py #数据分割脚本,用于生成训练索引文件
  • 编辑voc2yolo.py。设置tarin\val\test数据分割比例
#----------------------------------------------------------------------#
#   验证集的划分在train.py代码里面进行
#   test.txt和val.txt里面没有内容是正常的。训练不会使用到。
#----------------------------------------------------------------------#
'''
#--------------------------------注意----------------------------------#
如果在pycharm中运行时提示:
FileNotFoundError: [WinError 3] 系统找不到指定的路径。: './VOCdevkit/VOC2007/Annotations'
这是pycharm运行目录的问题,最简单的方法是将该文件复制到根目录后运行。
可以查询一下相对目录和根目录的概念。在VSCODE中没有这个问题。
#--------------------------------注意----------------------------------#
'''
import os
import random 
random.seed(0)

xmlfilepath=r'./Annotations'
saveBasePath=r"./ImageSets/Main/"
 
#----------------------------------------------------------------------#
#   想要增加测试集修改trainval_percent
#   train_percent不需要修改
#----------------------------------------------------------------------#
trainval_percent=0.9
train_percent=1

temp_xml = os.listdir(xmlfilepath)
total_xml = []
for xml in temp_xml:
    if xml.endswith(".xml"):
        total_xml.append(xml)

num=len(total_xml)  
list=range(num)  
tv=int(num*trainval_percent)  
tr=int(tv*train_percent)  
trainval= random.sample(list,tv)  
train=random.sample(trainval,tr)  
 
print("train and val size",tv)
print("traub suze",tr)
ftrainval = open(os.path.join(saveBasePath,'trainval.txt'), 'w')  
ftest = open(os.path.join(saveBasePath,'test.txt'), 'w')  
ftrain = open(os.path.join(saveBasePath,'train.txt'), 'w')  
fval = open(os.path.join(saveBasePath,'val.txt'), 'w')  
 
for i  in list:  
    name=total_xml[i][:-4]+'\n'  
    if i in trainval:  
        ftrainval.write(name)  
        if i in train:  
            ftrain.write(name)  
        else:  
            fval.write(name)  
    else:  
        ftest.write(name)  
  
ftrainval.close()  
ftrain.close()  
fval.close()  
ftest .close()
  • 生成训练索引文件
python voc2yolo.py

STEP3:对VOC数据集将进行类别统计

├──VOCdevkit/VOC2007/
    ├── get_all_class.py #VOC类别统计脚本

在其中写入以下内容

import xmltodict
import os

# VOC xml文件所在文件夹
annotation_dir="./Annotations/"

label_list = list()
# 逐一处理xml文件
for file in os.listdir(annotation_dir):
    annotation_path = os.path.join(annotation_dir,file)
    
    # 读取xml文件
    with open(annotation_path,'r') as f:
        xml_str = f.read()
    
    #转为字典
    xml_dic = xmltodict.parse(xml_str)
    
    # 获取label并去重加入到label_list
    objects = xml_dic["annotation"]["object"]
    if isinstance(objects,list): # xml文件中包含多个object
        for obj in objects:
            label = obj['name']
            if label not in label_list:
                label_list.append(label)
    else:# xml文件中只包含1个object
        obj = objects
        label = obj['name']
        if label not in label_list:
            label_list.append(label)
            
print(label_list)
  • 然后运行get_all_class.py获取所有的类别信息
['Person', 'BridgeVehicle', 'LuggageVehicle', 'Plane', 'RefuelVehicle', 'FoodVehicle', 'RubbishVehicle', 'WaterVehicle', 'PlatformVehicle', 'TractorVehicle', 'BridgeConnector']

STEP4:编辑model_data/voc_classes.txt 和下载权重文件

  • 将其中的类别数改为自己的,文件内容为
Person
BridgeVehicle
LuggageVehicle
Plane
RefuelVehicle
FoodVehicle
RubbishVehicle
WaterVehicle
PlatformVehicle
TractorVehicle
BridgeConnector
  • 下载权重文件到model_data文件夹下
wget https://github.com/bubbliiiing/yolov4-tiny-pytorch/releases/download/v1.0/yolov4_tiny_weights_coco.pth

STEP4:生成最终训练所需的txt文件

  • 然后运行voc_annotation.py
python voc_annotation.py
  • 此时会生成对应的2007_train.txt,每一行对应其图片位置及其真实框的位置

STEP5:修改train.py的NUM_CLASSSES

  • 将train.py的NUM_CLASSSES修改成所需要分的类的个数(不需要+1)

STEP6:开始训练

python train.py

2.模型预测

  1. 在yolo.py文件里面,在如下部分修改model_path和classes_path使其对应训练好的文件;model_path对应logs文件夹下面的权值文件,classes_path是model_path对应分的类

    _defaults = {
        #--------------------------------------------------------------------------#
        #   使用自己训练好的模型进行预测一定要修改model_path和classes_path!
        #   model_path指向logs文件夹下的权值文件,classes_path指向model_data下的txt
        #   如果出现shape不匹配,同时要注意训练时的model_path和classes_path参数的修改
        #--------------------------------------------------------------------------#
        "model_path"        : 'model_data/yolov4_tiny_weights_coco.pth',
        "classes_path"      : 'model_data/coco_classes.txt',
        #---------------------------------------------------------------------#
        #   anchors_path代表先验框对应的txt文件,一般不修改。
        #   anchors_mask用于帮助代码找到对应的先验框,一般不修改。
        #---------------------------------------------------------------------#
        "anchors_path"      : 'model_data/yolo_anchors.txt',
        "anchors_mask"      : [[3,4,5], [1,2,3]],
        #-------------------------------#
        #   所使用的注意力机制的类型
        #   phi = 0为不使用注意力机制
        #   phi = 1为SE
        #   phi = 2为CBAM
        #   phi = 3为ECA
        #-------------------------------#
        "phi"               : 0,  
        #---------------------------------------------------------------------#
        #   输入图片的大小,必须为32的倍数。
        #---------------------------------------------------------------------#
        "input_shape"       : [416, 416],
        #---------------------------------------------------------------------#
        #   只有得分大于置信度的预测框会被保留下来
        #---------------------------------------------------------------------#
        "confidence"        : 0.5,
        #---------------------------------------------------------------------#
        #   非极大抑制所用到的nms_iou大小
        #---------------------------------------------------------------------#
        "nms_iou"           : 0.3,
        #---------------------------------------------------------------------#
        #   该变量用于控制是否使用letterbox_image对输入图像进行不失真的resize,
        #   在多次测试后,发现关闭letterbox_image直接resize的效果更好
        #---------------------------------------------------------------------#
        "letterbox_image"   : False,
        #-------------------------------#
        #   是否使用Cuda
        #   没有GPU可以设置成False
        #-------------------------------#
        "cuda"              : True,
    }
  2. 运行predict.py,输入待测试的图片路径

    img/street.jpg
  3. 在predict.py里面进行设置可以进行fps测试和video视频检测。

3.评估自己的数据集

  1. 前往get_map.py文件修改classes_path,classes_path用于指向检测类别所对应的txt,这个txt和训练时的txt一样。评估自己的数据集必须要修改。
  2. 在yolo.py里面修改model_path以及classes_path。model_path指向训练好的权值文件,在logs文件夹里。classes_path指向检测类别所对应的txt。
  3. 运行get_map.py即可获得评估结果,评估结果会保存在map_out文件夹中。

参考资料

0

评论 (0)

打卡
取消