数据集转换xml2txt 、xml2json、json2coco

数据集转换

1、xml2txt

import xml.etree.ElementTree as ET
import os
from os import listdir, getcwd
from os.path import join
import random
from shutil import copyfile
# VOC类别名称映射到数字0到19
# class_names = {
#     'aeroplane': 0, 'bicycle': 1, 'bird': 2, 'boat': 3, 'bottle': 4,
#     'bus': 5, 'car': 6, 'cat': 7, 'chair': 8, 'cow': 9,
#     'diningtable': 10, 'dog': 11, 'horse': 12, 'motorbike': 13,
#     'person': 14, 'pottedplant': 15, 'sheep': 16, 'sofa': 17,
#     'train': 18, 'tvmonitor': 19
# }
class_names={'hat':0,'person':1}
TRAIN_RATIO = 80  # 训练集和测试集的比例为 8:2

def clear_hidden_files(path):
    """
    清除隐藏文件
    :param 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)

def convert(size, box):
    """
    将标注框从 VOC 格式转换为 YOLO 格式
    :param size: 图像的宽和高
    :param box: 标注框的位置 (xmin, xmax, ymin, ymax)
    :return: 转换后的标注框 (x, y, w, h)
    """
    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):
    """
    将单个图像的 XML 标注文件转换为 YOLO 格式的 txt 文件
    :param image_id: 图像ID（文件名）
    """
    in_file = open('VOCdevkit/VOC2007/Annotations/%s.xml' % image_id)
    out_file = open('VOCdevkit/VOC2007/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 class_names or int(difficult) == 1:
            continue
        cls_id = class_names[cls]  # 使用字典查找类别的 ID
        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_space_dir = os.path.join(data_base_dir, "VOC2007/")
if not os.path.isdir(work_space_dir):
    os.mkdir(work_space_dir)
annotation_dir = os.path.join(work_space_dir, "Annotations/")
if not os.path.isdir(annotation_dir):
    os.mkdir(annotation_dir)
clear_hidden_files(annotation_dir)
image_dir = os.path.join(work_space_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_space_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)

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)  # 列出所有图像文件
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))
        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):  # 训练集
            if os.path.exists(annotation_path):
                train_file.write(image_path + '\n')
                convert_annotation(nameWithoutExtention)  # 转换标签
                copyfile(image_path, yolov5_images_train_dir + voc_path)
                copyfile(label_path, yolov5_labels_train_dir + label_name)
        else:  # 测试集
            if os.path.exists(annotation_path):
                test_file.write(image_path + '\n')
                convert_annotation(nameWithoutExtention)  # 转换标签
                copyfile(image_path, yolov5_images_test_dir + voc_path)
                copyfile(label_path, yolov5_labels_test_dir + label_name)
train_file.close()
test_file.close()

2、xml2json

"""
    Author:DamonZheng
    Function:xml2json(for labelme)
    Edition:1.0
    Date:2022.2.21
"""

import argparse
import glob
import os
import xml.etree.ElementTree as ET
import json
from tqdm import tqdm

def parse_args():
    """
        参数配置
    """
    parser = argparse.ArgumentParser(description='xml2json')
    # parser.add_argument('--raw_label_dir', help='the path of raw label', default='Annotations')
    # parser.add_argument('--pic_dir', help='the path of picture', default='JPEGImages')
    # parser.add_argument('--save_dir', help='the path of new label', default='instances_train2017')
    # xml标注文件夹
    parser.add_argument('--raw_label_dir', help='the path of raw label', default='data/LabelmeData/VOC2007/Annotations')
    # xml标注的图片文件夹
    parser.add_argument('--pic_dir', help='the path of picture', default='data/LabelmeData/VOC2007/JPEGImages')
    # xml转成json输出的文件夹
    parser.add_argument('--save_dir', help='the path of new label', default='data/coco/annotations')
    args = parser.parse_args()
    return args

def read_xml_gtbox_and_label(xml_path):
    """
        读取xml内容
    """
    tree = ET.parse(xml_path)
    root = tree.getroot()
    size = root.find('size')
    width = int(size.find('width').text)
    height = int(size.find('height').text)
    depth = int(size.find('depth').text)
    points = []
    for obj in root.iter('object'):
        cls = obj.find('name').text
        pose = obj.find('pose').text
        xmlbox = obj.find('bndbox')
        xmin = float(xmlbox.find('xmin').text)
        xmax = float(xmlbox.find('xmax').text)
        ymin = float(xmlbox.find('ymin').text)
        ymax = float(xmlbox.find('ymax').text)
        box = [xmin, ymin, xmax, ymax]
        point = [cls, box]
        points.append(point)
    return points, width, height

def main():
    """
        主函数
    """
    args = parse_args()
    labels = glob.glob(args.raw_label_dir + '/*.xml')
    for i, label_abs in tqdm(enumerate(labels), total=len(labels)):
        _, label = os.path.split(label_abs)
        label_name = label.rstrip('.xml')
        img_path = os.path.join(args.pic_dir, label_name + '.jpg')
        points, width, height = read_xml_gtbox_and_label(label_abs)
        json_str = {}
        json_str['version'] = '4.5.6'
        json_str['flags'] = {}
        shapes = []
        for i in range(len(points)):
            shape = {}
            shape['label'] = points[i][0]
            shape['points'] = [[points[i][1][0], points[i][1][1]],
                                [points[i][1][0], points[i][1][3]],
                                [points[i][1][2], points[i][1][3]],
                                [points[i][1][2], points[i][1][1]]]
            shape['group_id'] = None
            shape['shape_type'] = 'polygon'
            shape['flags'] = {}
            shapes.append(shape)
        json_str['shapes'] = shapes
        json_str['imagePath'] = img_path
        json_str['imageData'] = None
        json_str['imageHeight'] = height
        json_str['imageWidth'] = width
        with open(os.path.join(args.save_dir, label_name + '.json'), 'w') as f:
            json.dump(json_str, f, indent=2)

if __name__ == '__main__':
    main()

3、json2coco

import os
import json
import glob


def polygon_to_bbox(polygon):
    # 提取多边形所有点的 x 和 y 坐标
    x_coords = [point[0] for point in polygon]
    y_coords = [point[1] for point in polygon]

    # 计算最小和最大 x、y 坐标
    x_min = min(x_coords)
    x_max = max(x_coords)
    y_min = min(y_coords)
    y_max = max(y_coords)

    # 计算边界框的宽度和高度
    width = x_max - x_min
    height = y_max - y_min

    # 返回边界框 [x_min, y_min, width, height]
    return [x_min, y_min, width, height]


def convert_to_coco(input_folder, output_file):
    coco_data = {
        "images": [],
        "annotations": [],
        "categories": [],
    }

    category_map = {}
    annotation_id = 1
    image_id = 1

    categories = ['hat','person']  # 假设只有一个类别，在json文件里定义的类别名字
    for i, category in enumerate(categories):
        coco_data['categories'].append({
            "id": i + 1,
            "name": category,
            "supercategory": "none"
        })
        category_map[category] = i + 1

    # 获取所有json文件
    json_files = glob.glob(os.path.join(input_folder, "*.json"))

    # 处理每个json文件
    for json_file in json_files:
        with open(json_file, 'r') as f:
            content = f.read().replace('\\', '\\\\')  # 修复反斜杠问题
            try:
                data = json.loads(content)  # 使用 json.loads 来解析修复后的内容
                print(data.keys())  # 打印出所有的键，查看是否正常加载
            except json.JSONDecodeError as e:
                print(f"Error parsing {json_file}: {e}")
                continue  # 如果有解析错误，跳过当前文件

        # 假设图像路径字段是 'imagePath'，如果实际字段不同，修改成正确的字段名
        image_info = {
            "id": image_id,
            "file_name": data["imagePath"],  # 使用 imagePath 字段
            "width": data["imageWidth"],
            "height": data["imageHeight"]
        }
        coco_data["images"].append(image_info)

        for ann in data["shapes"]:
            bbox = polygon_to_bbox(ann["points"])
            annotation_info = {
                "id": annotation_id,
                "image_id": image_id,
                "category_id": category_map[ann["label"]],
                "bbox": bbox,
                "area": bbox[2] * bbox[3],  # 宽 * 高
                "iscrowd": 0
            }
            coco_data["annotations"].append(annotation_info)
            annotation_id += 1

        image_id += 1

    # 确保输出文件名带有 .json 后缀
    if not output_file.endswith(".json"):
        output_file += ".json"

    # 输出到文件
    with open(output_file, 'w') as f:
        json.dump(coco_data, f, indent=4)


# 使用示例
input_folder = 'data/LabelmeData/annotations/val'  # 这里是你的json文件夹路径
output_file = 'data/coco/annotations'  # 这里是输出文件路径（没有 .json 后缀）
convert_to_coco(input_folder, output_file)