[DL series] VOC data set to YOLO data set (making YOLOv5/v7 custom data set)



foreword

Usually, when we use the official YOLOv5 and YOLOv7 models for magical training, we generally do not use large data sets such as COCO2017 , and usually debug on our own custom data sets or some small open source data sets. At this time, the problem of data sets will be involved.

Here I will discuss and share how to convert a dataset in VOC format into a dataset in YOLO format , and how to write datasets.yaml files required for training. Take the details of my personal actual operation as an example. If there are some details and data file locations, please discuss and exchange in the comment area.

In addition, the codes of YOLOv5 and YOLOv7 are very similar (you know), so many places in the two networks can communicate with each other, so the operation of making data sets in the two networks can be used in common, specific details such as folder addresses Just pay attention and it's okk.

[Note]: The data set division and operation problems in the article are based on the direction of <Deep Learning Vision-Target Detection> , and the data set is used for YOLOv5/v7 for image target detection.


1. Dataset preparation

The dataset I divided here is the public dataset, VOC2007trainval + VOC2012trainval . Since the data volume of the two datasets is very small, the two datasets are merged into one dataset for operation (it can be regarded as a custom VOC format dataset), which is the VOC07+ that will be encountered in the previous paper. 12 datasets.

First prepare two datasets of VOC07+12, and merge the two datasets (directly merge the contents of the Annotations and JPEGImages folders of the two datasets), among which we only need the Annotations tag folder, JPEGImages picture folder two folders are enough. The dataset folder division in VOC format is shown in the figure below.

insert image description here

Make your own dataset VOC0712 (images-22136, labels-22136) based on the code of the divided dataset, and put it in the root directory of the yolo code (that is, put it in the yolo project). This is the dataset file to be divided ( custom dataset).

[Note]: The production process of the VOC07+12 dataset is to merge the data in the Annotations and JPEGImages in the above VOC2007 + VOC2012 respectively (ctrl+c, ctrl+v)

insert image description here


2. Divide the dataset

Create the voc_to_yolo.py script in the root directory of yolo , which is the script for dividing the data set.

The VOC dataset has 20 classes. If it is a custom dataset, the classes need to be modified. The division ratio is determined by yourself, here 80 means train: val = 8: 2

  • voc_to_yolo.py
"""
# VOC数据集 转 YOLO数据集 格式
"""


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

classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog",
           "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
# classes=["ball"]

# 划分比率
TRAIN_RATIO = 80


def clear_hidden_files(path):
    dir_list = os.listdir(path)
    for i in dir_list:
        abspath = os.path.join(os.path.abspath(path), i)
        if os.path.isfile(abspath):
            if i.startswith("._"):
                os.remove(abspath)
        else:
            clear_hidden_files(abspath)


# size是原图的宽和高
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('VOCdevkit/VOC0712/Annotations/%s.xml' % image_id, 'rb')
    out_file = open('VOCdevkit/VOC0712/YOLOLabels/%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')
    in_file.close()
    out_file.close()


wd = os.getcwd()
data_base_dir = os.path.join(wd, "VOCdevkit/")
if not os.path.isdir(data_base_dir):
    os.mkdir(data_base_dir)
work_sapce_dir = os.path.join(data_base_dir, "VOC0712/")
if not os.path.isdir(work_sapce_dir):
    os.mkdir(work_sapce_dir)
annotation_dir = os.path.join(work_sapce_dir, "Annotations/")
if not os.path.isdir(annotation_dir):
    os.mkdir(annotation_dir)
clear_hidden_files(annotation_dir)
image_dir = os.path.join(work_sapce_dir, "JPEGImages/")
if not os.path.isdir(image_dir):
    os.mkdir(image_dir)
clear_hidden_files(image_dir)

# 这个部分可以不要
yolo_labels_dir = os.path.join(work_sapce_dir, "YOLOLabels/")
if not os.path.isdir(yolo_labels_dir):
    os.mkdir(yolo_labels_dir)
clear_hidden_files(yolo_labels_dir)

yolov5_images_dir = os.path.join(data_base_dir, "images/")
if not os.path.isdir(yolov5_images_dir):
    os.mkdir(yolov5_images_dir)
clear_hidden_files(yolov5_images_dir)
yolov5_labels_dir = os.path.join(data_base_dir, "labels/")
if not os.path.isdir(yolov5_labels_dir):
    os.mkdir(yolov5_labels_dir)
clear_hidden_files(yolov5_labels_dir)
yolov5_images_train_dir = os.path.join(yolov5_images_dir, "train/")
if not os.path.isdir(yolov5_images_train_dir):
    os.mkdir(yolov5_images_train_dir)
clear_hidden_files(yolov5_images_train_dir)
yolov5_images_test_dir = os.path.join(yolov5_images_dir, "val/")
if not os.path.isdir(yolov5_images_test_dir):
    os.mkdir(yolov5_images_test_dir)
clear_hidden_files(yolov5_images_test_dir)
yolov5_labels_train_dir = os.path.join(yolov5_labels_dir, "train/")
if not os.path.isdir(yolov5_labels_train_dir):
    os.mkdir(yolov5_labels_train_dir)
clear_hidden_files(yolov5_labels_train_dir)
yolov5_labels_test_dir = os.path.join(yolov5_labels_dir, "val/")
if not os.path.isdir(yolov5_labels_test_dir):
    os.mkdir(yolov5_labels_test_dir)
clear_hidden_files(yolov5_labels_test_dir)

# 这两个部分yolov5_train.txt, yolov5_val.txt可以不要
train_file = open(os.path.join(wd, "yolov5_train.txt"), 'w')
test_file = open(os.path.join(wd, "yolov5_val.txt"), 'w')
train_file.close()
test_file.close()
train_file = open(os.path.join(wd, "yolov5_train.txt"), 'a')
test_file = open(os.path.join(wd, "yolov5_val.txt"), 'a')
list_imgs = os.listdir(image_dir)  # list image_one files
prob = random.randint(1, 100)
print("Probability: %d" % prob)
for i in range(0, len(list_imgs)):
    path = os.path.join(image_dir, list_imgs[i])
    if os.path.isfile(path):
        image_path = image_dir + list_imgs[i]
        voc_path = list_imgs[i]
        (nameWithoutExtention, extention) = os.path.splitext(os.path.basename(image_path))
        (voc_nameWithoutExtention, voc_extention) = os.path.splitext(os.path.basename(voc_path))
        annotation_name = nameWithoutExtention + '.xml'
        annotation_path = os.path.join(annotation_dir, annotation_name)
        label_name = nameWithoutExtention + '.txt'
        label_path = os.path.join(yolo_labels_dir, label_name)
    prob = random.randint(1, 100)
    print("Probability: %d" % prob)
    if (prob < TRAIN_RATIO):  # train dataset
        if os.path.exists(annotation_path):
            train_file.write(image_path + '\n')
            convert_annotation(nameWithoutExtention)  # convert label
            copyfile(image_path, yolov5_images_train_dir + voc_path)
            copyfile(label_path, yolov5_labels_train_dir + label_name)
    else:  # test dataset
        if os.path.exists(annotation_path):
            test_file.write(image_path + '\n')
            convert_annotation(nameWithoutExtention)  # convert label
            copyfile(image_path, yolov5_images_test_dir + voc_path)
            copyfile(label_path, yolov5_labels_test_dir + label_name)
train_file.close()
test_file.close()

Run the script to generate the following files in the original dataset:

insert image description here

Among them , the images and labels folders are the divided data sets, which can be directly used for YOLOv5 training, without subsequent generation of data address.txt files. YOLOv7 does not work, YOLOv7 must use .txt files for training.

If you directly use the original picture for training (YOLOv5 supports direct reference to picture data for training; YOLOv7 does not support it, you need to refer to the absolute path .txt file generated by the picture data for training). What is mentioned here is the special case in YOLOv5. The yaml file of the custom data set can be written in the following format like coco.yaml.

  • mydataset.yaml
# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: ./VOCdevkit/images/train  # 17420 images
val: ./VOCdevkit/images/val # 4716 images

# number of classes
nc: 20

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

[Note]: In addition, several useless files generated by the script can be deleted directly. These include the YOLOLabels folder generated under the original data set in the above figure, and the yolov5_train.txt and yolov5_val.txt files generated in the root directory in the figure below can be deleted directly.

(This is because the source code is down on the Internet, it does not affect the use, so I am too lazy to modify it)

insert image description here


3. Generate datasets.txt file

If you have seen the COCO.yaml file in the YOLOv5/v7 code, it is clear that the data path in the file uses a txt file . In particular, v7 must use txt address files for training (I don’t know whether v7 can directly use pictures, but I think the code needs to be modified in many places to operate, which is prone to bugs. It is not as convenient as directly generating the address txt file of the data set)

Create the get_data_img_dir.py script in the root directory to generate the absolute path of the images in train and val under images after the data set is divided.

  • get_data_img_dir.py
"""
# 制作数据集的步骤:
1. voc_to_yolo.py划分数据集,留下images、labels
2. get_data_img_dir.py 分别生成images下train、val图片的绝对路径(yolov7训练数据集需要.txt文件)
------------------------------------------------------------------------------------------------------------------------
# 生成图片数据的绝对路径
"""

import os


def listdir(path, list_name):  # 传入存储的list
    for file in os.listdir(path):
        file_path = os.path.join(path, file)
        if os.path.isdir(file_path):
            listdir(file_path, list_name)
        else:
            list_name.append(file_path)


list_name = []
path = '/root/Desktop/code/base-yolov7/VOCdevkit/images/val'  # 文件夹路径
listdir(path, list_name)
print(list_name)

with open('./val0712.txt', 'w') as f:  # 要存入的txt
    write = ''
    for i in list_name:
        write = write + str(i) + '\n'
    f.write(write)

Run the script to generate two txt files, train0712.txt and val0712.txt , in the root directory . These two files are the absolute addresses of the image data in the training set. The files are shown in the figure below.

I am here because I want to train the network on the server, so the generated address is the address under the linux system.

insert image description here


4. Create a training datasets.yaml file based on the dataset

During the yolov5/v7 training process, move the txt file generated in the above picture to the created datasets-VOC0712 folder, so that the dataset is created, and the next step is to make the yaml file needed for model training, modeled after coco.yaml File to make your own yaml file, the code is shown in the figure below.

  • mydatasets.yaml
# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: ./datasets/VOC0712/train0712.txt  # 17420 images
val: ./datasets/VOC0712/val0712.txt  # 4716 images

# number of classes
nc: 20

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

OK, at this point the custom data set has been created as a whole. Next, you can carry out the next step of training according to the points that need to be paid attention to in model training!



[Note]: Due to my limited level, if there are any mistakes, please correct me! ! !

Guess you like

Origin blog.csdn.net/qq_39770163/article/details/127715144