制作YOLOv5数据集

在这里插入图片描述

1 使用labelImg标注数据集

labelImg是一个可视化的图像标定工具。Faster R-CNN,YOLO,SSD等目标检测网络所需要的数据集,均需要借此工具标定图像中的目标。可以标注两种格式:VOC标签格式(.xml)和YOLO标签格式(.txt)。

1.1 下载labelImg

可通过两种方式下载:
1.前往下载链接下载labelImg
2.使用命令行安装labelImg

pip install labelimg

1.2 使用labelImg标注数据

1.2.1 数据存放

预先准备好数据,以以下格式新建目录存放数据:

├── VOC2012  根目录
│├── JPEGImages  存放需要标注的图片
│├── Annotations  存放标签
│├── classes.txt  存放类别

先将图片放入JPEGImages文件夹中,再新建classes.txt,在txt中逐行输入类别,如下:
在这里插入图片描述

1.2.2 标注数据

如果是下载的labelImg,打开labelImg.exe即可。
如果是命令行安装的labelImg,使用命令:

labelimg JPEGImages classes.txt

打开后如下:
在这里插入图片描述
主要功能介绍:
在这里插入图片描述
使用快捷键A、D可进行图片切换。使用快捷键W进行标注,输入标签即可。
在这里插入图片描述
完成标注后,可在设置的标签保存路径查看标签。

2 将voc标签转换到yolo标签格式

如果标签是voc格式的,需要转换成yolo标签格式。如果标签是yolo格式的,可直接进行数据集划分。
首先准备voc_classes.json文件:

{
    "aeroplane": 1,
    "bicycle": 2,
    "bird": 3,
    "boat": 4,
    "bottle": 5,
    "bus": 6,
    "car": 7,
    "cat": 8,
    "chair": 9,
    "cow": 10,
    "diningtable": 11,
    "dog": 12,
    "horse": 13,
    "motorbike": 14,
    "person": 15,
    "pottedplant": 16,
    "sheep": 17,
    "sofa": 18,
    "train": 19,
    "tvmonitor": 20
}

第二步准备trans_voc2yolo.py文件进行转换,更改相应目录即可:

"""
将voc数据集标注信息(.xml)转为yolo标注格式(.txt)
"""
import os
from tqdm import tqdm
from lxml import etree
import json
import shutil

# voc数据集根目录可改成自己的数据集路径
voc_root = "VOC2012"

# 转换后的文件保存目录
save_file_root =  os.path.join(voc_root, "Yolo_labels")"

# label标签对应json文件
label_json_path = os.path.join(voc_root, "voc_classes.json")

# voc的xml标签路径
voc_xml_path = os.path.join(voc_root, "Annotations")

# 检查文件/文件夹都是否存在
assert os.path.exists(voc_xml_path), "VOC xml path not exist..."
assert os.path.exists(label_json_path), "label_json_path does not exist..."
if os.path.exists(save_file_root) is False:
    os.makedirs(save_file_root)

def parse_xml_to_dict(xml):
    """
    将xml文件解析成字典形式,参考tensorflow的recursive_parse_xml_to_dict
    Args:
        xml: xml tree obtained by parsing XML file contents using lxml.etree

    Returns:
        Python dictionary holding XML contents.
    """

    if len(xml) == 0:  # 遍历到底层,直接返回tag对应的信息
        return {xml.tag: xml.text}

    result = {}
    for child in xml:
        child_result = parse_xml_to_dict(child)  # 递归遍历标签信息
        if child.tag != 'object':
            result[child.tag] = child_result[child.tag]
        else:
            if child.tag not in result:  # 因为object可能有多个,所以需要放入列表里
                result[child.tag] = []
            result[child.tag].append(child_result[child.tag])
    return {xml.tag: result}


def translate_info(file_names: list ,class_dict: dict):
    """
    将对应xml文件信息转为yolo中使用的txt文件信息
    :param file_names:
    :param save_root:
    :param class_dict:
    :return:
    """

    for file in tqdm(file_names, desc="translate xml file..."):
        # 检查xml文件是否存在
        xml_path = os.path.join(voc_xml_path, file)
        assert os.path.exists(xml_path), "file:{} not exist...".format(xml_path)

        # read xml
        with open(xml_path) as fid:
            xml_str = fid.read()
        xml = etree.fromstring(xml_str)
        data = parse_xml_to_dict(xml)["annotation"]
        img_height = int(data["size"]["height"])
        img_width = int(data["size"]["width"])

        # write object info into txt
        assert "object" in data.keys(), "file: '{}' lack of object key.".format(xml_path)
        if len(data["object"]) == 0:
            # 如果xml文件中没有目标就直接忽略该样本
            print("Warning: in '{}' xml, there are no objects.".format(xml_path))
            continue

        with open(os.path.join(save_file_root, file[:-4] + ".txt"), "w") as f:
            for index, obj in enumerate(data["object"]):
                # 获取每个object的box信息
                xmin = float(obj["bndbox"]["xmin"])
                xmax = float(obj["bndbox"]["xmax"])
                ymin = float(obj["bndbox"]["ymin"])
                ymax = float(obj["bndbox"]["ymax"])
                class_name = obj["name"]
                class_index = class_dict[class_name] - 1  # 目标id从0开始

                # 进一步检查数据,有的标注信息中可能有w或h为0的情况,这样的数据会导致计算回归loss为nan
                if xmax <= xmin or ymax <= ymin:
                    print("Warning: in '{}' xml, there are some bbox w/h <=0".format(xml_path))
                    continue

                # 将box信息转换到yolo格式
                xcenter = xmin + (xmax - xmin) / 2
                ycenter = ymin + (ymax - ymin) / 2
                w = xmax - xmin
                h = ymax - ymin

                # 绝对坐标转相对坐标,保存6位小数
                xcenter = round(xcenter / img_width, 6)
                ycenter = round(ycenter / img_height, 6)
                w = round(w / img_width, 6)
                h = round(h / img_height, 6)

                info = [str(i) for i in [class_index, xcenter, ycenter, w, h]]

                if index == 0:
                    f.write(" ".join(info))
                else:
                    f.write("\n" + " ".join(info))


def main():
    # read class_indict
    json_file = open(label_json_path, 'r')
    class_dict = json.load(json_file)

    # voc信息转yolo,并将图像文件复制到相应文件夹
    file_names = os.listdir(voc_xml_path)
    translate_info(file_names, class_dict)

if __name__ == "__main__":
    main()

转换后的标签格式如下:

14 0.538066 0.452 0.360082 0.5

文件目录结构:

├── VOC2012  根目录
│├── JPEGImages  存放需要标注的图片(.jpg,.bmp...)
│├── Annotations  存放voc标签(.xml)
│├── Yolo_labels  存放yolo标签(.txt)
│├── voc_classes.json 存放分类json文件

3 数据集划分训练集、验证集和测试集

import os
import random
import shutil

# 改成自己的路径
root_path = 'VOC2012'
txtfilepath = root_path + '\Yolo_labels\\'
imgfilepath = root_path + '\JPEGImages\\'
txtsavepath = root_path + '\ImageSets\MainTest'
yolo_root_path = root_path + '\YOLO'
yolo_images_train_path = yolo_root_path + '\images\\train\\'
yolo_images_val_path = yolo_root_path + '\images\\val\\'
yolo_images_test_path = yolo_root_path + '\images\\test\\'

yolo_labels_train_path = yolo_root_path + '\labels\\train\\'
yolo_labels_val_path = yolo_root_path + '\labels\\val\\'
yolo_labels_test_path = yolo_root_path + '\labels\\test\\'

if not os.path.exists(txtsavepath):
    os.makedirs(txtsavepath)

if not os.path.exists(yolo_root_path):
    os.makedirs(yolo_root_path)

if not os.path.exists(yolo_images_train_path):
    os.makedirs(yolo_images_train_path)

if not os.path.exists(yolo_images_val_path):
    os.makedirs(yolo_images_val_path)

if not os.path.exists(yolo_images_test_path):
    os.makedirs(yolo_images_test_path)

if not os.path.exists(yolo_labels_train_path):
    os.makedirs(yolo_labels_train_path)

if not os.path.exists(yolo_labels_val_path):
    os.makedirs(yolo_labels_val_path)

if not os.path.exists(yolo_labels_test_path):
    os.makedirs(yolo_labels_test_path)

def copyfile(imgname,name,imgpath,labelpath):
    image_copy_to = os.path.join(imgpath,imgname[:-1])
    if os.path.exists(image_copy_to) is False:
        shutil.copyfile(imgfilepath +imgname[:-1] , image_copy_to)
    label_copy_to = os.path.join(labelpath,name[:-1])
    if os.path.exists(label_copy_to) is False:
        shutil.copyfile(txtfilepath + name[:-1] , label_copy_to)

def main():
    train_test_percent = 0.9  # (训练集+验证集)/(训练集+验证集+测试集)
    train_valid_percent = 0.9  # 训练集/(训练集+验证集)

    total_xml = os.listdir(txtfilepath)
    num = len(total_xml)
    list = range(num)
    tv = int(num * train_test_percent)   # 训练集+验证集数量
    ts = int(num-tv)   # 测试集数量
    tr = int(tv * train_valid_percent)  # 验证集数量
    tz = int(tv-tr)   # 训练集数量
    trainval = random.sample(list, tv)
    train = random.sample(trainval, tr)

    print("train and valid size:", tv)
    print("train size:", tz)
    print("test size:", ts)
    print("valid size:", tr)
    ftest = open(txtsavepath + '/test.txt', 'w')
    ftrain = open(txtsavepath + '/train.txt', 'w')
    fvalid = open(txtsavepath + '/valid.txt', 'w')

    ftestimg = open(txtsavepath + '/img_test.txt', 'w')
    ftrainimg = open(txtsavepath + '/img_train.txt', 'w')
    fvalidimg = open(txtsavepath + '/img_valid.txt', 'w')

    for i in list:
        name = total_xml[i][:-4] + '.txt' + '\n'
        imgname = total_xml[i][:-4] + '.jpg' + '\n'
        if i in trainval:
            if i in train:
                ftrain.write(name)
                ftrainimg.write(imgname)
                copyfile(imgname,name,yolo_images_train_path,yolo_labels_train_path)
            else:
                fvalid.write(name)
                fvalidimg.write(imgname)
                copyfile(imgname,name,yolo_images_val_path,yolo_labels_val_path)
        else:
            ftest.write(name)
            ftestimg.write(imgname)
            copyfile(imgname,name,yolo_images_test_path,yolo_labels_test_path)

    ftrain.close()
    fvalid.close()
    ftest.close()

    ftrainimg.close()
    fvalidimg.close()
    ftestimg.close()

    print("finished!")

if __name__ == "__main__":
    main()

文件目录结构:

├── VOC2012  根目录
│├── JPEGImages  存放需要标注的图片(.jpg,.bmp...)
│├── Annotations  存放voc标签(.xml)
│├── Yolo_labels  存放yolo标签(.txt)
│├── ImageSet
││├──Main
│││├──test.txt  存放测试集yolo标签
│││├──train.txt  存放训练集yolo标签
│││├──valid.txt  存放验证集集yolo标签
│││├──img_test.txt  存放测试集图片
│││├──img_train.txt  存放训练集集图片
│││├──img_valid.txt  存放验证集图片
│├── YOLO
││├──images
│││├──test  测试集图片
│││├──train  训练集图片
│││├──val  验证集图片
││├──labels
│││├──test  测试集标签
│││├──train  训练集标签
│││├──val  验证集标签
│├── voc_classes.json 存放分类json文件

后面会继续介绍如何在YOLOv5中使用自己制作的数据集。
已更新博客:
深度学习笔记:使用YOLOv5训练自己的数据集模型

评论 9
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值