DIOR数据集下载及预处理

数据集下载

天翼云盘18713912310账号

其中包括四个文件夹Annotations、ImageSets、JPEGImages-test和JPEGImages-trainval四个文件夹。

Annotations包括Horizontal Bounding Boxes和Oriented Bounding Boxes两个文件夹，用于存放标签。

JPEGImages-test和JPEGImages-trainval分别存放了数据集的测试集图片和训练集验证集图片。

数据预处理

我们需要将xml格式的标签转换为yolov5所使用的txt格式标签，再对数据集按照6:2:2分为训练集、验证集和测试集。

首先，我们在DIOR数据集文件夹下创建文件夹JPEGImages用于存放所有图片，就是将JPEGImages-test和JPEGImages-trainval文件夹下的图片复制到该文件夹中。

然后在DIOR文件夹下创建.py文件用于将标签转化为txt格式，以及对数据集进行划分。可命名为DIORxml2txtYOLOv5.py，代码如下。

# DIORxml2txtYOLOv5.py
# coding:utf-8

import os
import random
import argparse

import xml.etree.ElementTree as ET
from os import getcwd
from shutil import copyfile


parser = argparse.ArgumentParser()
# xml文件的地址，根据自己的数据进行修改 xml一般存放在Annotations下
parser.add_argument('--xml_path', default='./Annotations/Horizontal Bounding Boxes', type=str, help='input xml label path')
# 数据集的划分，地址选择自己数据下的ImageSets/Main

opt = parser.parse_args()

sets = ['train', 'val', 'test']
classes = ['airplane', 'airport', 'baseballfield', 'basketballcourt', 'bridge', 'chimney', 'dam',
              'Expressway-Service-area', 'Expressway-toll-station', 'golffield', 'groundtrackfield', 'harbor',
              'overpass', 'ship', 'stadium', 'storagetank', 'tenniscourt', 'trainstation', 'vehicle', 'windmill']

abs_path = os.getcwd()
print(abs_path)


if not os.path.exists('DIOR_dataset/'):
    os.makedirs('DIOR_dataset/')

if not os.path.exists('DIOR_dataset/labels/'):
    os.makedirs('DIOR_dataset/labels/')
if not os.path.exists('DIOR_dataset/labels/train'):
    os.makedirs('DIOR_dataset/labels/train')
if not os.path.exists('DIOR_dataset_yolo/labels/test'):
    os.makedirs('DIOR_dataset/labels/test')
if not os.path.exists('DIOR_dataset_yolo/labels/val'):
    os.makedirs('DIOR_dataset/labels/val')


if not os.path.exists('DIOR_dataset/images/'):
    os.makedirs('DIOR_dataset/images/')
if not os.path.exists('DIOR_dataset/images/train'):
    os.makedirs('DIOR_dataset/images/train')
if not os.path.exists('DIOR_dataset/images/test'):
    os.makedirs('DIOR_dataset/images/test')
if not os.path.exists('DIOR_dataset/images/val'):
    os.makedirs('DIOR_dataset/images/val')


def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = (box[0] + box[1]) / 2.0 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    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, path):
#输入输出文件夹，根据实际情况进行修改
    in_file = open('./Annotations/Horizontal Bounding Boxes/%s.xml' % (image_id), encoding='UTF-8')
    out_file = open('DIOR_dataset/labels/' + path + '/%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
        #difficult = obj.find('Difficult').text
        cls = obj.find('name').text
        if cls not in classes:
            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))
        b1, b2, b3, b4 = b
        # 标注越界修正
        if b2 > w:
            b2 = w
        if b4 > h:
            b4 = h
        b = (b1, b2, b3, b4)
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')



train_percent = 0.6
test_percent = 0.2
val_percent = 0.2

xmlfilepath = opt.xml_path
# txtsavepath = opt.txt_path
total_xml = os.listdir(xmlfilepath)
# if not os.path.exists(txtsavepath):
#     os.makedirs(txtsavepath)

num = len(total_xml)
list_index = range(num)
list_index = list(list_index)
random.shuffle(list_index)


train_nums = list_index[:int(num * train_percent)]
test_nums = list_index[int(num * train_percent): int(num * test_percent)+int(num * train_percent)]
val_nums = list_index[int(num * test_percent)+int(num * train_percent):]



for i in list_index:
    name = total_xml[i][:-4]
    if i in train_nums:
        convert_annotation(name, 'train')   # lables
        image_origin_path = './JPEGImages/' + name + '.jpg'
        image_target_path = 'DIOR_dataset/images/train/' + name + '.jpg'
        copyfile(image_origin_path, image_target_path)

    if i in test_nums:
        convert_annotation(name, 'test')   # lables
        image_origin_path = './JPEGImages/' + name + '.jpg'
        image_target_path = 'DIOR_dataset/images/test/' + name + '.jpg'
        copyfile(image_origin_path, image_target_path)

    if i in val_nums:
        convert_annotation(name, 'val')   # lables
        image_origin_path = './JPEGImages/' + name + '.jpg'
        image_target_path = 'DIOR_dataset/images/val/' + name + '.jpg'
        copyfile(image_origin_path, image_target_path)

最终在当前目录下生成DIOR_dataset文件夹，其中包含两个文件夹images和labels。

images和labels文件夹下分别包含了训练集、验证集和测试集的图片和标签。

但这种格式还无法直接进行训练，因为yolov5查询标签的方式是在储存图片的文件夹的上级目录下找到labels文件夹，所以我们还需对格式进行调整。

首先在DIOR_dataset文件夹下创建test、train、val文件夹。再将原images下的test、train、val下的图片剪切到对应的刚创建好的文件夹下的images文件夹下，将原labels文件夹下的test、train、val下的标签剪切到对应的刚创建好的文件夹下的labels文件夹下，即：

DIOR_dataset/images/test-----→DIOR_dataset/test/images

DIOR_dataset/images/train-----→DIOR_dataset/train/images

DIOR_dataset/images/val-----→DIOR_dataset/val/images

DIOR_dataset/labels/test-----→DIOR_dataset/test/labels

DIOR_dataset/labels/train-----→DIOR_dataset/train/labels

DIOR_dataset/labels/val-----→DIOR_dataset/val/labels

完成格式调整，最终如下所示。

test、train和val文件夹下分别是各自的images和labels文件夹。

至此，预处理部分就结束了。

编写data.yaml文件

在DIOR_dataset文件夹下创建data.yaml文件，内容如下。

# Path: ../datasets/DIOR_data/DIOR_dataset/data.yaml
# path
# ├── Yolo_v5
# └── datasets
#     └── DIOR_data
#         └── DIOR_dataset
#             ├── test
#             ├── train
#             ├── val
#             └── data.yaml   ← There

train: ../../datasets/DIOR_data/DIOR_dataset/train/images  # train images (relative to 'path') 128 images
val: ../../datasets/DIOR_data/DIOR_dataset/val/images  # val images (relative to 'path') 128 images
test: ../../datasets/DIOR_data/DIOR_dataset/test/images  # test images (optional)
nc: 20
# Classes
names: ['airplane',
        'airport',
        'baseball field',
        'basketball court',
        'bridge',
        'chimney',
        'dam',
        'expressway service area',
        'expressway toll station',
        'golf course',
        'ground track field',
        'harbor',
        'overpass',
        'ship',
        'stadium',
        'storage tank',
        'tennis court',
        'train station',
        'vehicle',
        'wind mill']