廣東工業智造創新大賽數據格式轉VOC

#!/usr/bin/env python
# coding: utf-8
# @author: lwh
# @file: mouse_find.py
# @time: 2021/1/06 15:54
from os import getcwd
from xml.etree import ElementTree as ET

import os
import json
import math

# 創建xml文件的函數
def create_tree(image_name, h, w):
    global annotation
    # 創建樹根annotation
    annotation = ET.Element('annotation')
    # 創建一級分支folder
    folder = ET.SubElement(annotation, 'folder')
    # 添加folder標簽內容
    folder.text = (img_dir)

    # 創建一級分支filename
    filename = ET.SubElement(annotation, 'filename')
    filename.text = image_name

    # 創建一級分支path
    path = ET.SubElement(annotation, 'path')

    path.text = getcwd() + '\{}\{}'.format(img_dir, image_name)  # 用於返回當前工作目錄

    # 創建一級分支source
    source = ET.SubElement(annotation, 'source')
    # 創建source下的二級分支database
    database = ET.SubElement(source, 'database')
    database.text = 'Unknown'

    # 創建一級分支size
    size = ET.SubElement(annotation, 'size')
    # 創建size下的二級分支圖像的寬、高及depth
    width = ET.SubElement(size, 'width')
    width.text = str(w)
    height = ET.SubElement(size, 'height')
    height.text = str(h)
    depth = ET.SubElement(size, 'depth')
    depth.text = '3'

    # 創建一級分支segmented
    segmented = ET.SubElement(annotation, 'segmented')
    segmented.text = '0'


# 定義一個創建一級分支object的函數
def create_object(root, xi, yi, xa, ya, obj_name):  # 參數依次，樹根，xmin，ymin，xmax，ymax
    # 創建一級分支object
    _object = ET.SubElement(root, 'object')
    # 創建二級分支
    name = ET.SubElement(_object, 'name')
    # print(obj_name)
    name.text = str(obj_name)
    pose = ET.SubElement(_object, 'pose')
    pose.text = 'Unspecified'
    truncated = ET.SubElement(_object, 'truncated')
    truncated.text = '0'
    difficult = ET.SubElement(_object, 'difficult')
    difficult.text = '0'
#     # 創建bndbox
    bndbox = ET.SubElement(_object, 'bndbox')
    xmin = ET.SubElement(bndbox, 'xmin')
    xmin.text = '%s' % xi
    ymin = ET.SubElement(bndbox, 'ymin')
    ymin.text = '%s' % yi
    xmax = ET.SubElement(bndbox, 'xmax')
    xmax.text = '%s' % xa
    ymax = ET.SubElement(bndbox, 'ymax')
    ymax.text = '%s' % ya


def json_to_xml(json_path, xml_dir):
    with open(json_path, 'r') as load_f:
        load_list = json.load(load_f)
        xml_list = []
        for load_dict in load_list:
            img_name = load_dict["name"]
            h = load_dict["image_height"]
            w = load_dict["image_width"]
            category = load_dict["category"]
            xmin = str(math.floor(load_dict["bbox"][0]))  #向下取整，保證目標被完全框住
            ymin = str(math.floor(load_dict["bbox"][1]))  #向下取整，保證目標被完全框住
            xmax = str(math.ceil(load_dict["bbox"][2]))  #向上取整，保證目標被完全框住=
            ymax = str(math.ceil(load_dict["bbox"][3]))  #向上取整，保證目標被完全框住
            #判斷是否出現該圖片得xml文件
            xml_name = img_name.split(".")[0] + ".xml"

            # print(xml_name)
            # xml_list = os.listdir(xml_dir)

            if xml_name in xml_list:
                print(xml_name)
                xml_path = os.path.join(xml_dir, xml_name)
                doc = ET.parse(xml_path)
                root = doc.getroot()
                create_object(root, xmin, ymin, xmax, ymax, category)
                doc.write(xml_path, encoding="utf-8", xml_declaration=True)
            else:
                xml_list.append(xml_name)
                create_tree(img_name, h, w)
                create_object(annotation, xmin, ymin, xmax, ymax,category)
                tree = ET.ElementTree(annotation)
                tree.write('{}\{}.xml'.format(xml_dir, img_name.strip('.jpg')))




if __name__ == '__main__':
    img_dir = r"F:\Intelligent_manufacturing\tile_round1_train_20201231\train_imgs"
    json_path = r"F:\Intelligent_manufacturing\tile_round1_train_20201231\train_annos.json"  # 該目錄為存放json文件的路徑
    xml_dir = r"F:\Intelligent_manufacturing\tile_round1_train_20201231\train_xmls"  # 該目錄為放xml文件的路徑
    json_to_xml(json_path, xml_dir)