Yolo-FastestV2 训练自己的数据集

根据原作者源码及说明实践后总结

操作环境

Ubuntu 18.02 +cuda 11.1 +pytorch 1.9.0
作者源码下载：https://github.com/dog-qiuqiu/Yolo-FastestV2

配置环境

新建：

conda create -n yolov2 python=3.8

安装pytorch：

conda install pytorch torchvision torchaudio cudatoolkit=11.1 -c pytorch -c nvidia

下载源码：

git clone https://github.com/dog-qiuqiu/Yolo-FastestV2

安装依赖包：

pip install -r requirements.txt

测试图片：

python test.py --data data/coco.data --weights modelzoo/coco2017-0.241078ap-model.pth --img img/000139.jpg

测试成功!没有问题，就可以开始训练自己的数据集啦!

构建数据集：

此方法是将 voc 格式数据集转换为符合此算法的数据集格式
以下三个文件放在新建的data数据集里

使用maketxt.py划分测试集验证集

import os
import random

trainval_percent = 0.1
train_percent = 1
xmlfilepath = './data/Annotations'
total_xml = os.listdir(xmlfilepath)
num = len(total_xml)
list = range(num)
tv = int(num * trainval_percent)
trainval = random.sample(list, tv)

txt_train='./data/ImageSets/train.txt'
if os.path.exists(txt_train):
    os.remove(txt_train)
else:
    open(txt_train,'w')

txt_val='./data/ImageSets/val.txt'
if os.path.exists(txt_val):
    os.remove(txt_val)
else:
    open(txt_val,'w')

ftrain = open(txt_train, 'w')
fval = open(txt_val, 'w')

for i in list:
    name = total_xml[i][:-4] + '\n'
    ftrain.write(name)
    if i in trainval:
        fval.write(name)

ftrain.close()
fval.close()

划分好验证集测试集后，会在ImageSets文件中生成train.txt和val.txt
使用voc_label.py生成.txt标签文件
voc_label.py

# -*- coding: utf-8 -*-
# xml解析包
import xml.etree.ElementTree as ET
import os
from os import getcwd
import shutil

sets = ['train', 'val']
classes = ['apple','cup','house']

style = '.jpg'


#判断文件是否为空
def ifnone(path) :
    if os.path.exists(path):
        #print("文件存在")
        if os.path.getsize(path):
            #print("文件存在且不为空")
            #print(os.path.getsize(FILE))
            Size=os.path.getsize(path)

            os.system('ls -lh %s' %(path))
            return 1
        else:
            #print("文件存在但为空...")   
            os.system('ls -lh %s' %(path)