深度学习数据集格式转换代码

深度学习中不同的模型训练所需要的数据集格式可能不同，这时可能就需要将数据集格式转为该模型要求的格式。

yolo2voc

import xml.dom.minidom
import glob
from PIL import Image
from math import ceil
import shutil
import os
yolo_file = r'folder'  # yolo格式下的存放txt标注文件的文件夹
turn_xml_file = r'folder' # 转换后储存xml的文件夹地址
img_file = r'folder'  # 存放图片的文件夹

#class id
labels = ['label']
src_img_dir = img_file
src_txt_dir = yolo_file
src_xml_dir = turn_xml_file  # 转换后储存xml的文件夹地址

img_Lists = glob.glob(src_img_dir + '/*.jpg')
img_basenames = []
for item in img_Lists:
    img_basenames.append(os.path.basename(item))  # os.path.basename返回path最后的文件名

img_names = []
for item in img_basenames:
    temp1, temp2 = os.path.splitext(item)  # os.path.splitext(“文件路径”)    分离文件名与扩展名
    img_names.append(temp1)

total_num = len(img_names)  # 统计当前总共要转换的图片标注数量
count = 0  # 技术变量
for img in img_names:  # 这里的img是不加后缀的图片名称，如：'GF3_SAY_FSI_002732_E122.3_N29.9_20170215_L1A_HH_L10002188179__1__4320___10368'
    count += 1
    if count % 1000 == 0:
        print("当前转换进度{}/{}".format(count, total_num))
    im = Image.open((src_img_dir + '/' + img + '.jpg'))
    width, height = im.size

    # 打开yolo格式下的txt文件
    gt = open(src_txt_dir + '/' + img + '.txt').read().splitlines()
    if gt:
        # 将主干部分写入xml文件中
        xml_file = open((src_xml_dir + '/' + img + '.xml'), 'w')
        xml_file.write('<annotation>\n')
        xml_file.write('    <folder>VOC2007</folder>\n')
        xml_file.write('    <filename>' + str(img) + '.jpg' + '</filename>\n')
        xml_file.write('    <size>\n')
        xml_file.write('        <width>' + str(width) + '</width>\n')
        xml_file.write('        <height>' + str(height) + '</height>\n')
        xml_file.write('        <depth>3</depth>\n')
        xml_file.write('    </size>\n')

        # write the region of image on xml file
        for img_each_label in gt:
            spt = img_each_label.split(' ')  # 这里如果txt里面是以逗号‘，’隔开的，那么就改为spt = img_each_label.split(',')。
            xml_file.write('    <object>\n')
            xml_file.write('        <name>' + str(labels[int(spt[0])]) + '</name>\n')
            xml_file.write('        <pose>Unspecified</pose>\n')
            xml_file.write('        <truncated>0</truncated>\n')
            xml_file.write('        <difficult>0</difficult>\n')
            xml_file.write('        <bndbox>\n')

            center_x = round(float(spt[1].strip()) * width)
            center_y = round(float(spt[2].strip()) * height)
            bbox_width = round(float(spt[3].strip()) * width)
            bbox_height = round(float(spt[4].strip()) * height)
            xmin = str(int(center_x - bbox_width / 2))
            ymin = str(int(center_y - bbox_height / 2))
            xmax = str(int(center_x + bbox_width / 2))
            ymax = str(int(center_y + bbox_height / 2))

            xml_file.write('            <xmin>' + xmin + '</xmin>\n')
            xml_file.write('            <ymin>' + ymin + '</ymin>\n')
            xml_file.write('            <xmax>' + xmax + '</xmax>\n')
            xml_file.write('            <ymax>' + ymax + '</ymax>\n')
            xml_file.write('        </bndbox>\n')
            xml_file.write('    </object>\n')

        xml_file.write('</annotation>')
    else:
        # 将主干部分写入xml文件中
        xml_file = open((src_xml_dir + '/' + img + '.xml'), 'w')
        xml_file.write('<annotation>\n')
        xml_file.write('    <folder>VOC2007</folder>\n')
        xml_file.write('    <filename>' + str(img) + '.jpg' + '</filename>\n')
        xml_file.write('    <size>\n')
        xml_file.write('        <width>' + str(width) + '</width>\n')
        xml_file.write('        <height>' + str(height) + '</height>\n')
        xml_file.write('        <depth>3</depth>\n')
        xml_file.write('    </size>\n')
        xml_file.write('</annotation>')

# 将转换后的xml文件按train 和test 归类到train.txt和test.txt中
path = r'folder'
xml_Lists = glob.glob(src_xml_dir + '/*.xml')

xml_basenames = []
for item in xml_Lists:
    xml_basenames.append(os.path.basename(item))

xml_names = []  # 这里是将xml文件去掉.xml后缀储存到的列表中
for item in xml_basenames:
    temp1, temp2 = os.path.splitext(item)  # os.path.splitext(“文件路径”)    分离文件名与扩展名
    xml_names.append(temp1)

txt_file = open((path + '/val.txt'), 'w')
for item in xml_names:
    txt_file.write(str(item) + '\n')

yolo2coco

import os
import cv2
import json
from tqdm import tqdm
from sklearn.model_selection import train_test_split
import argparse

# visdrone2019
classes = ['']

parser = argparse.ArgumentParser()
parser.add_argument('--image_path', default=r'H:\dataset\Dataset4\images\val', type=str, help="path of images")
parser.add_argument('--label_path', default=r'H:\dataset\Dataset4\labels\val', type=str, help="path of labels .txt")
parser.add_argument('--save_path', default=r'H:\dataset\pest4\annotations\instances_val2017.json', type=str,
                    help="if not split the dataset, give a path to a json file")
arg = parser.parse_args()


def yolo2coco(arg):
    print("Loading data from ", arg.image_path, arg.label_path)

    assert os.path.exists(arg.image_path)
    assert os.path.exists(arg.label_path)

    originImagesDir = arg.image_path
    originLabelsDir = arg.label_path
    # images dir name
    indexes = os.listdir(originImagesDir)

    dataset = {'categories': [], 'annotations': [], 'images': []}
    for i, cls in enumerate(classes, 0):
        dataset['categories'].append({'id': i, 'name': cls, 'supercategory': 'mark'})

    # 标注的id
    ann_id_cnt = 0
    for k, index in enumerate(tqdm(indexes)):
        # 支持 png jpg 格式的图片.
        txtFile = f'{index[:index.rfind(".")]}.txt'
        stem = index[:index.rfind(".")]
        # 读取图像的宽和高
        try:
            im = cv2.imread(os.path.join(originImagesDir, index))
            height, width, _ = im.shape
        except Exception as e:
            print(f'{os.path.join(originImagesDir, index)} read error.\nerror:{e}')
        # 添加图像的信息
        if not os.path.exists(os.path.join(originLabelsDir, txtFile)):
            # 如没标签，跳过，只保留图片信息.
            continue
        dataset['images'].append({'file_name': index,
                                  'id': stem,
                                  'width': width,
                                  'height': height})
        with open(os.path.join(originLabelsDir, txtFile), 'r') as fr:
            labelList = fr.readlines()
            for label in labelList:
                label = label.strip().split()
                x = float(label[1])
                y = float(label[2])
                w = float(label[3])
                h = float(label[4])

                # convert x,y,w,h to x1,y1,x2,y2
                H, W, _ = im.shape
                x1 = (x - w / 2) * W
                y1 = (y - h / 2) * H
                x2 = (x + w / 2) * W
                y2 = (y + h / 2) * H
                # 标签序号从0开始计算, coco2017数据集标号混乱，不管它了。
                cls_id = int(label[0])
                width = max(0, x2 - x1)
                height = max(0, y2 - y1)
                dataset['annotations'].append({
                    'area': width * height,
                    'bbox': [x1, y1, width, height],
                    'category_id': cls_id,
                    'id': ann_id_cnt,
                    'image_id': stem,
                    'iscrowd': 0,
                    # mask, 矩形是从左上角点按顺时针的四个顶点
                    'segmentation': [[x1, y1, x2, y1, x2, y2, x1, y2]]
                })
                ann_id_cnt += 1

    # 保存结果
    with open(arg.save_path, 'w') as f:
        json.dump(dataset, f)
        print('Save annotation to {}'.format(arg.save_path))


if __name__ == "__main__":
    yolo2coco(arg)

voc2yolo

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


# classes = ['hard_hat', 'other', 'regular', 'long_hair', 'braid', 'bald', 'beard']


def convert(size, box):
    # size=(width, height)  b=(xmin, xmax, ymin, ymax)
    # x_center = (xmax+xmin)/2        y_center = (ymax+ymin)/2
    # x = x_center / width            y = y_center / height
    # w = (xmax-xmin) / width         h = (ymax-ymin) / height

    x_center = (box[0] + box[1]) / 2.0
    y_center = (box[2] + box[3]) / 2.0
    x = x_center / size[0]
    y = y_center / size[1]

    w = (box[1] - box[0]) / size[0]
    h = (box[3] - box[2]) / size[1]

    # print(x, y, w, h)
    return (x, y, w, h)


def convert_annotation(xml_files_path, save_txt_files_path, classes):
    xml_files = os.listdir(xml_files_path)
    print(xml_files)
    for xml_name in xml_files:
        print(xml_name)
        xml_file = os.path.join(xml_files_path, xml_name)
        out_txt_path = os.path.join(save_txt_files_path, xml_name.split('.')[0] + '.txt')
        out_txt_f = open(out_txt_path, 'w')
        tree = ET.parse(xml_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))
            # b=(xmin, xmax, ymin, ymax)
            print(w, h, b)
            bb = convert((w, h), b)
            out_txt_f.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')


if __name__ == "__main__":
    # 测试程序
    # classes = ['hard_hat', 'other', 'regular', 'long_hair', 'braid', 'bald', 'beard']
    # xml_files = r'D:\ZF\1_ZF_proj\3_脚本程序\2_voc格式转yolo格式\voc_labels'
    # save_txt_files = r'D:\ZF\1_ZF_proj\3_脚本程序\2_voc格式转yolo格式\yolo_labels'
    # convert_annotation(xml_files, save_txt_files, classes)

    # ====================================================================================================
    # 把帽子头发胡子的voc的xml标签文件转化为yolo的txt标签文件
    # 1、帽子头发胡子的类别
    classes1 = ['class']
    # 2、voc格式的xml标签文件路径
    xml_files1 = r'xml'
    # 3、转化为yolo格式的txt标签文件存储路径
    save_txt_files1 = r'label'

    convert_annotation(xml_files1, save_txt_files1, classes1)

voc2coco

import sys
import os
import json
import xml.etree.ElementTree as ET

START_BOUNDING_BOX_ID = 0

PRE_DEFINE_CATEGORIES = {"LicensePlate": 1}  # 修改的地方，修改为自己的类别


# If necessary, pre-define category and its id
#  PRE_DEFINE_CATEGORIES = {"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}


def get(root, name):
    vars = root.findall(name)
    return vars


def get_and_check(root, name, length):
    vars = root.findall(name)
    if len(vars) == 0:
        raise NotImplementedError('Can not find %s in %s.' % (name, root.tag))
    if length > 0 and len(vars) != length:
        raise NotImplementedError('The size of %s is supposed to be %d, but is %d.' % (name, length, len(vars)))
    if length == 1:
        vars = vars[0]
    return vars


def get_filename_as_int(filename):
    try:
        filename = os.path.splitext(filename)[0]
        return filename
    except:
        raise NotImplementedError('Filename %s is supposed to be an integer.' % (filename))


# xml_list为xml文件存放的txt文件名    xml_dir为真实xml的存放路径    json_file为存放的json路径
def convert(xml_list, xml_dir, json_file):
    list_fp = open(xml_list, 'r')
    json_dict = {"images": [], "type": "instances", "annotations": [],
                 "categories": []}
    categories = PRE_DEFINE_CATEGORIES
    bnd_id = START_BOUNDING_BOX_ID
    for line in list_fp:
        line = line.strip()
        line = line + ".xml"
        print("Processing %s" % (line))
        xml_f = os.path.join(xml_dir, line)
        tree = ET.parse(xml_f)
        root = tree.getroot()
        path = get(root, 'path')
        if len(path) == 1:
            filename = os.path.basename(path[0].text)
        elif len(path) == 0:
            filename = get_and_check(root, 'filename', 1).text
        else:
            raise NotImplementedError('%d paths found in %s' % (len(path), line))
        ## The filename must be a number
        image_id = get_filename_as_int(filename)
        size = get_and_check(root, 'size', 1)
        width = int(get_and_check(size, 'width', 1).text)
        height = int(get_and_check(size, 'height', 1).text)
        image = {'file_name': filename, 'height': height, 'width': width,
                 'id': image_id}
        json_dict['images'].append(image)
        ## Cruuently we do not support segmentation
        #  segmented = get_and_check(root, 'segmented', 1).text
        #  assert segmented == '0'
        for obj in get(root, 'object'):
            category = get_and_check(obj, 'name', 1).text
            if category not in categories:
                new_id = len(categories)
                categories[category] = new_id
            category_id = categories[category]
            bndbox = get_and_check(obj, 'bndbox', 1)
            xmin = int(get_and_check(bndbox, 'xmin', 1).text) - 1
            ymin = int(get_and_check(bndbox, 'ymin', 1).text) - 1
            xmax = int(get_and_check(bndbox, 'xmax', 1).text)
            ymax = int(get_and_check(bndbox, 'ymax', 1).text)
            assert (xmax > xmin)
            assert (ymax > ymin)
            o_width = abs(xmax - xmin)
            o_height = abs(ymax - ymin)
            ann = {'area': o_width * o_height, 'iscrowd': 0, 'image_id':
                image_id, 'bbox': [xmin, ymin, o_width, o_height],
                   'category_id': category_id, 'id': bnd_id, 'ignore': 0,
                   'segmentation': []}
            json_dict['annotations'].append(ann)
            bnd_id = bnd_id + 1

    for cate, cid in categories.items():
        cat = {'supercategory': 'none', 'id': cid, 'name': cate}
        json_dict['categories'].append(cat)
    json_fp = open(json_file, 'w')
    json_str = json.dumps(json_dict)
    json_fp.write(json_str)
    json_fp.close()
    list_fp.close()


if __name__ == '__main__':
    # xml_list为xml文件存放的txt文件名    xml_dir为真实xml的存放路径    json_file为存放的json路径
    # xml_list = './data/VOCdevkit/ImageSets/Main/test.txt'
    # xml_list = './data/VOCdevkit/ImageSets/Main/train.txt'
    xml_list = './VOC_plate/VOC2007/ImageSets/Main/test.txt'
    xml_dir = './VOC_plate/VOC2007/Annotations'
    # json_dir = './data/COCO/annotations/test.json'  # 注意！！！这里test.json先要自己创建，不然
    # json_dir = './data/COCO/annotations/train.json'  # 注意！！！这里test.json先要自己创建，不然
    json_dir = 'RT-DETR/rtdetr_pytorch/COCO_plate/annotations/test.json'  # 注意！！！这里test.json先要自己创建，不然																  #程序回报权限不足
    convert(xml_list, xml_dir, json_dir)

注意，实际使用时要根据自己的实际情况修改代码中的路径和参数。