目标检测实例分割数据集转换:从XML和JSON到YOLOv8(txt)

2023-12-29 06:44:07

yolov8导航

????????如果大家想要了解关于yolov8的其他任务和相关内容可以点击这个链接,我这边整理了许多其他任务的说明博文,后续也会持续更新,包括yolov8模型优化、sam等等的相关内容。

YOLOv8(附带各种任务详细说明链接)

源码下载地址:

XML&JSON 目标检测、实例分割标签转换给yolo用脚本

引言

????????在计算机视觉领域,目标检测是一个重要而复杂的任务。随着深度学习技术的发展,各种高效的算法和模型如YOLOv8不断涌现。然而,在使用这些先进模型之前,必须确保我们的数据集是正确格式化的。今天,我将分享一段Python代码,它能够将XML和JSON格式的标注文件转换为适合YOLOv8模型训练的格式。

数据集转换的重要性

????????在目标检测任务中,我们通常有两种常见的标注格式:XML和JSON。而YOLOv8需要的是一种特定格式的文本文件,其中包含了目标的类别和位置信息。因此,将现有的标注文件转换为YOLO格式是实施有效训练的第一步。

源码

? ? ? ? 这段代码大体包含两个功能,可以转换xml、json格式的目标检测以及实例分割任务的标签,填写的时候需要根据要求填写标签所在的路径:

import json
import pandas as pd
import xml.etree.ElementTree as ET
import os, cv2
import numpy as np
import glob

classes = []

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(xmlpath, xmlname):
    with open(xmlpath, "r", encoding='utf-8') as in_file:
        txtname = xmlname[:-4] + '.txt'
        txtfile = os.path.join(txtpath, txtname)
        tree = ET.parse(in_file)
        root = tree.getroot()
        filename = root.find('filename')

        img = cv2.imdecode(np.fromfile('{}/{}.{}'.format(imgpath, xmlname[:-4], postfix), np.uint8), cv2.IMREAD_COLOR)
        h, w = img.shape[:2]
        res = []
        for obj in root.iter('object'):
            cls = obj.find('name').text
            if cls not in classes:
                classes.append(cls)
            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)
            res.append(str(cls_id) + " " + " ".join([str(a) for a in bb]))
        if len(res) != 0:
            with open(txtfile, 'w+') as f:
                f.write('\n'.join(res))


def query_json_to_txt(dir_json, dir_txt):
    classes2id = {}
    num = 0
    jsons = os.listdir(dir_json)
    for i in jsons:
        json_path = os.path.join(dir_json, i)
        with open(json_path, 'r', encoding="utf-8") as f:
            json_data = json.load(f)
            # print(json_data['shapes'])
            for j in json_data['shapes']:
                if j['label'] not in classes2id:
                    classes2id[j['label']] = num
                    num += 1

    def json2txt(path_json, path_txt):  # 可修改生成格式
        with open(path_json, 'r', encoding='utf-8') as path_json:
            jsonx = json.load(path_json)
            with open(path_txt, 'w+') as ftxt:
                shapes = jsonx['shapes']
                # 获取图片长和宽
                width = jsonx['imageWidth']
                height = jsonx['imageHeight']
                # print(shapes)
                cat = shapes[0]['label']
                cat = classes2id[cat]

                for shape in shapes:
                    # 获取矩形框两个角点坐标
                    x1 = shape['points'][0][0]
                    y1 = shape['points'][0][1]
                    x2 = shape['points'][1][0]
                    y2 = shape['points'][1][1]

                    dw = 1. / width
                    dh = 1. / height
                    x = dw * (x1 + x2) / 2
                    y = dh * (y1 + y2) / 2
                    w = dw * abs(x2 - x1)
                    h = dh * abs(y2 - y1)
                    yolo = f"{cat} {x} {y} {w} {h} \n"
                    ftxt.writelines(yolo)


    list_json = os.listdir(dir_json)
    for cnt, json_name in enumerate(list_json):
        if os.path.splitext(json_name)[-1] == ".json":
            path_json = dir_json + json_name
            path_txt = dir_txt + json_name.replace('.json', '.txt')
            json2txt(path_json, path_txt)

    pd.DataFrame([{"原始类别": k, "编码": v} for k,v in classes2id.items()]).to_excel("label_codes.xlsx", index=None)

###################### 实例分割处理 #######################################
def parse_json_for_instance_segmentation(json_path, label_dict={}):
    # 打开并读取JSON文件
    with open(json_path, 'r', encoding='utf-8') as file:
        json_info = json.load(file)

    annotations = []  # 初始化一个用于存储注释的列表
    # 遍历JSON文件中的每个形状(shape)
    for shape in json_info["shapes"]:
        label = shape["label"]  # 获取实例分割类别的标签
        # 如果标签在标签字典中 则直接编码
        if label in label_dict:
            label_dict[label] = label_dict[label]  # 为该标签分配一个编码
        # 如果不在标签中
        else:
            next_label_code = max(label_dict.values(), default=-1) + 1
            label_dict[label] = next_label_code

        cat = label_dict[label]  # 获取该标签对应的编码
        points = shape["points"]  # 获取形状的点
        # 将点转换为字符串格式,并按照图像的宽度和高度进行归一化
        points_str = ' '.join([f"{round(point[0]/json_info['imageWidth'], 6)} {round(point[1]/json_info['imageHeight'], 6)}" for point in points])
        annotation_str = f"{cat} {points_str}\n"  # 将编码和点字符串合并为一行注释
        annotations.append(annotation_str)  # 将该注释添加到列表中

    return annotations, label_dict  # 返回注释列表和更新后的标签字典

def process_directory_for_instance_segmentation(json_dir, txt_save_dir, label_dict=None):
    if not os.path.exists(txt_save_dir):
        os.makedirs(txt_save_dir)

    # 如果没有提供初始的label_dict,则从空字典开始
    final_label_dict = label_dict.copy() if label_dict is not None else {}

    json_files = [f for f in os.listdir(json_dir) if f.endswith(".json")]
    for json_file in json_files:
        json_path = os.path.join(json_dir, json_file)
        txt_path = os.path.join(txt_save_dir, json_file.replace(".json", ".txt"))

        # 每次解析时传递当前的final_label_dict
        annotations, updated_label_dict = parse_json_for_instance_segmentation(json_path, final_label_dict)

        # 更新final_label_dict,确保包括所有新标签及其编码
        final_label_dict.update(updated_label_dict)

        # 检查annotations是否为空,如果为空则跳过写入操作
        if annotations:
            with open(txt_path, "w") as file:
                file.writelines(annotations)

    # 保存最终的标签字典
    pd.DataFrame(list(final_label_dict.items()), columns=['原始Label', '编码后Label']).to_excel('label_codes.xlsx', index=False)
    return final_label_dict


# 传入参数为 实例分割 和 目标检测
query_type = "目标检测"
# 传入原始标签数据,比如xml、json格式文件所在的目录下
label_directory = './query_data/xy/Annotations/'
# 填写想转换的txt输出到哪里
output_directory = './query_data/xy/txt/'

anno_files = glob.glob(label_directory + "*")
file_type = anno_files[0].split(".")[-1]

label_dict = {
    # 如果想预设标签就在这里填入对应的键值对
}

if query_type == "实例分割":
    process_directory_for_instance_segmentation(label_directory, output_directory, label_dict)


elif query_type == "目标检测":
    if file_type == "json":
        query_json_to_txt(label_directory, output_directory)

    ## 处理xml格式文件
    elif file_type == "xml" or file_type == "XML":
        postfix = 'jpg'
        imgpath = 'query_data/xy/images'
        xmlpath = 'query_data/xy/Annotations'
        txtpath = 'query_data/xy/txt'

        if not os.path.exists(txtpath):
            os.makedirs(txtpath, exist_ok=True)

        file_list = glob.glob(xmlpath + "/*")
        error_file_list = []
        for i in range(0, len(file_list)):
            try:
                path = file_list[i]
                if ('.xml' in path) or ('.XML' in path):
                    convert_annotation(path, path.split("\\")[-1])
                    print(f'file {list[i]} convert success.')
                else:
                    print(f'file {list[i]} is not xml format.')
            except Exception as e:
                print(f'file {list[i]} convert error.')
                print(f'error message:\n{e}')
                error_file_list.append(list[i])
        print(f'this file convert failure\n{error_file_list}')
        print(f'Dataset Classes:{classes}')

脚本功能解析

处理XML文件

????????XML格式广泛用于多种图像标注工具中。我们的脚本首先解析XML文件,提取出其中的目标类别和边界框信息。然后,这些信息被转换为YOLO格式,并保存在相应的.txt文件中。

处理JSON文件

????????JSON格式也是一种常见的标注格式,尤其在实例分割任务中。脚本中包含的函数query_json_to_txt能够处理这些JSON文件,将它们转换为YOLO格式,并生成一个包含类别名称和编号的映射表。

实例分割数据处理

????????对于实例分割任务,我们的脚本还提供了解析JSON文件的功能,提取出多边形的坐标信息,并转换为YOLO格式。

使用指南

要使用此脚本,您只需按照如下步骤操作:

  1. 确定您的数据集类型(目标检测或实例分割)。
  2. 准备您的XML或JSON标注文件。
  3. 运行脚本进行转换。
  4. 使用生成的文本文件训练您的YOLOv8模型。

结语

????????数据预处理是机器学习项目中至关重要的一步。通过本文介绍的脚本,您可以轻松地将XML和JSON格式的标注文件转换为适合YOLOv8模型的格式。这不仅节省了大量时间,也为实现高效的目标检测任务奠定了基础。如果有哪里写的不够清晰,小伙伴本可以给评论或者留言,我这边会尽快的优化博文内容,另外如有需要,我这边可支持技术答疑与支持。

文章来源:https://blog.csdn.net/qq_42452134/article/details/135244845
本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。