[人工智能]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))
            return 2


# 进行归一化操作
def convert(size, box):  # size:(原图w,原图h) , box:(xmin,xmax,ymin,ymax)
    dw = 1. / size[0]  # 1/w
    dh = 1. / size[1]  # 1/h
    x = (box[0] + box[1]) / 2.0  # 物体在图中的中心点x坐标
    y = (box[2] + box[3]) / 2.0  # 物体在图中的中心点y坐标
    w = box[1] - box[0]  # 物体实际像素宽度
    h = box[3] - box[2]  # 物体实际像素高度
    x = x * dw  # 物体中心点x的坐标比(相当于 x/原图w)
    w = w * dw  # 物体宽度的宽度比(相当于 w/原图w)
    y = y * dh  # 物体中心点y的坐标比(相当于 y/原图h)
    h = h * dh  # 物体宽度的宽度比(相当于 h/原图h)
    return (x, y, w, h)  # 返回 相对于原图的物体中心点的x坐标比,y坐标比,宽度比,高度比,取值范围[0-1]


# year ='2012', 对应图片的id（文件名）
def convert_annotation(image_id):
    '''
    将对应文件名的xml文件转化为label文件，xml文件包含了对应的bunding框以及图片长款大小等信息，
    通过对其解析，然后进行归一化最终读到label文件中去，也就是说
    一张图片文件对应一个xml文件，然后通过解析和归一化，能够将对应的信息保存到唯一一个label文件中去
    labal文件中的格式：calss x y w h　　同时，一张图片对应的类别有多个，所以对应的ｂｕｎｄｉｎｇ的信息也有多个
    '''
    # 对应的通过year 找到相应的文件夹，并且打开相应image_id的xml文件，其对应bund文件
    in_file = open('./data/Annotations/%s.xml' % (image_id), encoding='utf-8')
    # 准备在对应的image_id 中写入对应的label，分别为
    # <object-class> <x> <y> <width> <height>
    out_file = open('./data/JPEGimages/%s.txt' % (image_id), 'w', encoding='utf-8')
    # 解析xml文件
    tree = ET.parse(in_file)
    # 获得对应的键值对
    root = tree.getroot()
    # 获得图片的尺寸大小
    size = root.find('size')
    # 如果xml内的标记为空，增加判断条件
    if size != None:
        # 获得宽
        w = int(size.find('width').text)
        # 获得高
        h = int(size.find('height').text)
        # 遍历目标obj
        for obj in root.iter('object'):
            # 获得difficult ？？
            difficult= 0
            if obj.find('difficult')!=None :
                difficult = obj.find('difficult').text
            # 获得类别 =string 类型
            cls = obj.find('name').text
            #print("display",image_id,cls,difficult)
            # 如果类别不是对应在我们预定好的class文件中，或difficult==1则跳过
            if cls not in classes or int(difficult) == 1:
                continue
            # 通过类别名称找到id
            cls_id = classes.index(cls)
            # 找到bndbox 对象
            xmlbox = obj.find('bndbox')
            # 获取对应的bndbox的数组 = ['xmin','xmax','ymin','ymax']
            b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
                 float(xmlbox.find('ymax').text))
            #print(image_id, cls, b)
            # 带入进行归一化操作
            # w = 宽, h = 高， b= bndbox的数组 = ['xmin','xmax','ymin','ymax']
            bb = convert((w, h), b)
            # bb 对应的是归一化后的(x,y,w,h)
            # 生成 calss x y w h 在label文件中
            out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')


# 返回当前工作目录
wd = getcwd()
#print(wd)

# 先找labels文件夹如果不存在则创建
#labels = './data/labels'
#if os.path.exists(labels):
#    shutil.rmtree(labels)  # delete output folder
#os.makedirs(labels)  # make new output folder


for image_set in sets:
    '''
    对所有的文件数据集进行遍历
    做了两个工作：
　　　　１．将所有图片文件都遍历一遍，并且将其所有的全路径都写在对应的txt文件中去，方便定位
　　　　２．同时对所有的图片文件进行解析和转化，将其对应的bundingbox 以及类别的信息全部解析写到label 文件中去
    　　　　　最后再通过直接读取文件，就能找到对应的label 信息
    '''
    # 读取在ImageSets/Main 中的train、test..等文件的内容
    # 包含对应的文件名称
    image_ids = open('./data/ImageSets/%s.txt' % (image_set)).read().strip().split()
    # 打开对应的2012_train.txt 文件对其进行写入准备

    txt_name = './data/%s.txt' % (image_set)
    if os.path.exists(txt_name):
        os.remove(txt_name)
    else:
        open(txt_name, 'w')

    list_file = open(txt_name, 'w')
    # 将对应的文件_id以及全路径写进去并换行
    for image_id in image_ids:
        list_file.write('./data/JPEGimages/%s%s\n' % (image_id, style))
        # 调用  year = 年份  image_id = 对应的文件名_id
        #print(image_id)
        pathxml='./data/Annotations/%s.xml' % (image_id)
        xmlresult = ifnone(pathxml)#判断xml是否有空文件
        if xmlresult == 1 :
            convert_annotation(image_id)
        elif xmlresult == 2 :
            path2 = './data/JPEGImages/%s.txt' % (image_id)
            file = open(path2,'w')
    # 关闭文件
    print("end")
    list_file.close()

计算锚点偏差

python3 genanchors.py --traintxt data/train.txt

计算以后会生成anchor6.txt，把第一行复制到data的coco.data里anchors里

coco.data和coco.names解析见其他文章

训练自己的数据集

python train.py --data data/coco.data

测试

python test.py --data data/coco.data --weights modelzoo/coco-140-epoch-0.872988ap-model.pth --img img/000005.jpg

测试多张图片用如下代码

import os
import cv2
import time
import argparse
import torch.nn as nn
import torch
import model.detector
import utils.utils
from thop import profile

if __name__ == '__main__':
    # 指定训练配置文件
    parser = argparse.ArgumentParser()
    parser.add_argument('--data', type=str, default='',
                        help='Specify training profile *.data')
    parser.add_argument('--weights', type=str, default='',
                        help='The path of the .pth model to be transformed')
    parser.add_argument('--img', type=str, default='',
                        help='The path of test image')

    opt = parser.parse_args()
    cfg = utils.utils.load_datafile(opt.data)
    # assert os.path.exists(opt.weights), "请指定正确的模型路径"
    # assert os.path.exists(opt.img), "请指定正确的测试图像路径"
    assert os.path.exists(opt.weights), "./weights/best-model.pth"
    assert os.path.exists(opt.img), "./data/voc/JPEGImages/000009.jpg"

    #指定后端设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


    #初始化模型
    model = model.detector.Detector(cfg["classes"], cfg["anchor_num"], True).to(device) 
    model = nn.DataParallel(model)  # 在加载时需要把网络也转成DataParallel的
    model.to(device)  
    model.load_state_dict(torch.load(opt.weights, map_location=device))
    # sets the module in eval node
    model.eval()

    # 数据预处理
    for jpg in os.listdir(opt.img):
        ori_img = os.path.join(opt.img, jpg)
        print(ori_img)
        ori_img = cv2.imread(ori_img)
        res_img = cv2.resize(ori_img, (cfg["width"], cfg["height"]), interpolation=cv2.INTER_LINEAR)
        img = res_img.reshape(1, cfg["height"], cfg["width"], 3)
        img = torch.from_numpy(img.transpose(0, 3, 1, 2))
        img = img.to(device).float() / 255.0

        # 模型推理
        computertime = 0
        start = time.perf_counter()
        preds = model(img)
        end = time.perf_counter()
        computertime = (end - start) * 1000.
        print("forward time:%fms" % computertime)

        # 特征图后处理
        output = utils.utils.handel_preds(preds, cfg, device)
        output_boxes = utils.utils.non_max_suppression(output, conf_thres=0.5, iou_thres=0.4)

        # 加载label names
        LABEL_NAMES = []
        with open(cfg["names"], 'r') as f:
            for line in f.readlines():
                LABEL_NAMES.append(line.strip())

        h, w, _ = ori_img.shape
        scale_h, scale_w = h / cfg["height"], w / cfg["width"]

        # 绘制预测框
        for box in output_boxes[0]:
            box = box.tolist()

            obj_score = box[4]
            category = LABEL_NAMES[int(box[5])]

            x1, y1 = int(box[0] * scale_w), int(box[1] * scale_h)
            x2, y2 = int(box[2] * scale_w), int(box[3] * scale_h)

            cv2.rectangle(ori_img, (x1, y1), (x2, y2), (255, 255, 0), 2)
            cv2.putText(ori_img, '%.2f' % obj_score, (x1, y1 - 5), 0, 0.7, (0, 255, 0), 2)
            cv2.putText(ori_img, category, (x1, y1 - 25), 0, 0.7, (0, 255, 0), 2)

        cv2.imwrite("datatest/testresult/{}.jpg".format(jpg.split(".")[0]), ori_img)

评估数据集

python3 evaluation.py --data data/coco.data --weights weights/best-model.pth