【目标检测】YOLOv5跑通VOC2007数据集

zstar

发布于 2022-09-07 08:54:04

2.2K10

代码可运行

文章被收录于专栏：往期博文往期博文

运行总次数：0

代码可运行

前言

如果经常阅读我博客的读者，想必对YOLOv5并不陌生。在Pytorch：YOLO-v5目标检测(上)一文中，我使用了coco128数据集，非常轻松的跑通了。然而在使用VOC2007数据集时，却遇到重重阻碍。主要问题在数据标签转化这个阶段，VOC数据集标注形式是xml，需要将其转换为txt。很多博文并未把文件的放置位置交代清楚，导致走了不少弯路，本篇博文就记录如何不走弯路地跑通VOC数据集。

VOC数据集

我是用的VOC数据集是VOC2007，它的文件格式长这样：

注：官网上下载的话训练验集和测试集是分成两个压缩包VOCtrainval_06-Nov-2007.tar和VOCtest_06-Nov-2007.tar，此处我将其进行合并，如需下载我合并后的数据集，可跳转到文末进行下载。

Annotations 该文件包含数据标注的xml文件
ImageSets 该文件包含数据集的划分
JPEGImages 该文件包含图像数据集

拿到数据集，第一件需要干的事是把JPEGImages文件夹名称改成images，具体原因后面解释，改完之后如图所示：

数据集处理

从GitHub上，下载Yolov5代码项目地址： https://github.com/ultralytics/yolov5 下载完之后，在文件中新建一个文件，我起名为voc_label.py，起名，位置均随意，这个文件主要干两件事：

1、将ImageSets中的训练集、测试集、验证集提出来，拼凑成可以直接访问的绝对路径，如图所示

2、把xml格式的标注转换成txt格式，如图所示：

图中红色方框圈住的是要提取出的数据，对于每一个目标，txt格式的数据有5个，分别为(class, x , y , w, h)，即类别、中心点(x,y)，宽高(w,h)。这里的转换公式在代码中有写道，不作细述。注意class必须是以0开始的序号，否则训练时会出现类别溢出的报错。

voc_label.py

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join

sets = ['train', 'test', 'val']

Imgpath = 'D:/Desktop/Work/VOC2007/images' 
xmlfilepath = 'D:/Desktop/Work/VOC2007/Annotations/'  
ImageSets_path = 'D:/Desktop/Work/VOC2007/ImageSets/Segmentation/'
Label_path = 'D:/Desktop/Work/VOC2007/'
classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]



def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)


def convert_annotation(image_id):
    in_file = open(xmlfilepath + '%s.xml' % (image_id))
    out_file = open(Label_path + 'labels/%s.txt' % (image_id), 'w')
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)
    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
             float(xmlbox.find('ymax').text))
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')


for image_set in sets:
    if not os.path.exists(Label_path + 'labels/'):
        os.makedirs(Label_path + 'labels/')
    image_ids = open(ImageSets_path + '%s.txt' % (image_set)).read().strip().split()
    list_file = open(Label_path + '%s.txt' % (image_set), 'w')
    for image_id in image_ids:
        list_file.write(Imgpath + '/%s.jpg\n' % (image_id))
        convert_annotation(image_id)
    list_file.close()

使用时，替换路径和类别classes即可。

运行之后，会在数据集文件夹下生成四个文件，labels中的txt文件和images中的图片名称一一对应。

训练准备

YOLOv5是通过yaml格式的配置文件来找到对应的训练测试数据，因此在训练前需要对该文件进行简单配置。尽管YOLOv5已有自带的voc.yaml配置文件，但为了之后训练自己的数据，我这里还是新建一个mydata.yaml来将VOC数据进行导入。

train: D:/Desktop/Work/VOC2007/train.txt
val: D:/Desktop/Work/VOC2007/val.txt
test: D:/Desktop/Work/VOC2007/test.txt

# number of classes
nc: 20

# class names
names: [ 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor' ]

文件内容很简单，需要指定三个数据集的路径，类别个数和类别名称。

开始训练

打开train.py文件主要需要调整三个参数我使用yolov5s作为预训练模型，如下图所示进行修改。

运行之后，程序会自动从官网下载yolov5s.pt模型文件，放在根目录下面。如果前面的操作和我一样，到此时应该能顺利跑通训练程序。程序默认的epochs为300，为了尽快看到效果，可以改小些，训练完后，会在runs -> train -> exp 中看到训练和验证的文件。

如果再次运行，会发现生成了一个新的文件夹exp2…这是由于yolov5默认将每次运行结果单独放一个文件夹存储。具体代码可参见train.py的518行

opt.save_dir = increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok | opt.evolve)  # increment run

这里的increment_path函数就是每次累加exp序号，如果需要将结果保存到指定路径，可以直接将路径替换到等号右边。

当然，Yolov5也提供了一种便捷的方式放置不断生成exp文件。在超参数设定中，有个参数为exist-ok，其默认为False，将其设置为True，则程序再次运行会在最近的exp生成文件。

另外，还有各种不同的参数，参数含义补充如下，可供研究参考。

    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default='yolov5s.pt', help='initial weights path')
    parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml', help='model.yaml path  模型配置和网络结构的yaml文件路径')
    parser.add_argument('--data', type=str, default='data/mydata.yaml', help='data.yaml path 数据集配置的yaml文件路径，包括数据集的路径，类名等关键信息')
    parser.add_argument('--hyp', type=str, default='data/hyp.scratch.yaml', help='超参数配置的yaml文件路径')
    parser.add_argument('--epochs', type=int, default=2)
    parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs')
    parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[train, test] image sizes  输入图像的分辨率大小')
    parser.add_argument('--rect', action='store_true', help='rectangular training  是否采用矩形图像训练，默认是False')
    parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training 从给定的path/last.pt恢复训练，如果为空，则从最近保存的path/last.pt恢复训练')
    parser.add_argument('--nosave', action='store_true', help='only save final checkpoint 保存模型，只保存最后的检查点，默认为False')
    parser.add_argument('--notest', action='store_true', help='only test final epoch 训练中不进行test，只测试最后一个轮次，默认为False')
    parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check 不自动调整anchor，默认为False')
    parser.add_argument('--evolve', action='store_true', help='evolve hyperparameters 是否进行超参数进化，默认为False')
    parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
    parser.add_argument('--cache-images', action='store_true', help='cache images for faster training 是否提前缓存图片到内存，以加快训练速度，默认为False')
    parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%% 是否进行多尺度训练，默认为False')
    parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class 将所有数据按照一个类别进行训')
    parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer 是否使用 torch.optim.Adam() 优化器，默认是False')
    parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode 是否使用跨卡同步BN')
    parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify')
    parser.add_argument('--workers', type=int, default=0, help='maximum number of dataloader workers')
    parser.add_argument('--project', default='runs/train', help='save to project/name')
    parser.add_argument('--entity', default=None, help='W&B entity')
    parser.add_argument('--name', default='exp', help='save to project/name')
    parser.add_argument('--exist-ok', default=True, action='store_true', help='existing project/name ok, do not increment')
    parser.add_argument('--quad', action='store_true', help='quad dataloader')
    parser.add_argument('--linear-lr', action='store_true', help='linear LR')
    parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')
    parser.add_argument('--upload_dataset', action='store_true', help='Upload dataset as W&B artifact table')
    parser.add_argument('--bbox_interval', type=int, default=-1, help='Set bounding-box image logging interval for W&B')
    parser.add_argument('--save_period', type=int, default=-1, help='Log model after every "save_period" epoch')
    parser.add_argument('--artifact_alias', type=str, default="latest", help='version of dataset artifact to be used')
    opt = parser.parse_args()

开始测试

上面的过程只运用到了模型的训练和验证，打开test.py，可使用测试集对模型进行测试。需要修改的参数和训练程序差不多，值得注意的是，yolov5在每次训练之后会输出两个模型文件，一个是best.pt，即验证集上表现最好(mAP最高)的模型，另一个是last.pt，即最后一轮训练的模型。通常来说，输入best.pt进行测试即可。测试运行之后，会在runs中生成一个test文件，其中可进行训练结果和标注结果的对比。