# -*- coding: utf-8 -*-
# !/usr/bin/env python

import os
import json
from tqdm import tqdm
from imutils import paths
import xml.etree.ElementTree as ET

def xml2coco(name_id, xml_dir, json_file):
    # COCO JSON 文件的初始化数据结构
    coco_dict = {
        "images": [],
        "annotations": [],
        "categories": []
    }
    # 定义类别名称和 ID 映射关系
    class_map = {}
    with open(name_id, 'r') as f:
        for line in f.readlines():
            name, id = line.strip().split(':')
            class_map[name] = id
    xml_list = list(paths.list_files(xml_dir, '.xml'))
    ret_coco_dict = getxml2coco(xml_list, class_map, coco_dict)
    # 保存 COCO JSON 文件
    with open(json_file, 'w') as f:
        json.dump(ret_coco_dict, f)

def getxml2coco(xml_list, class_map, coco_dict):
    # 遍历 XML 文件夹中的所有文件
    count = 0
    for xml_path in tqdm(xml_list):
        count += 1
        # print(xml_path)
        # 解析 XML 文件内容
        tree = ET.parse(xml_path)
        root = tree.getroot()

        # 提取图像信息
        image_id = root.find('filename').text.split('.')[0]
        width = int(root.find('size/width').text)
        height = int(root.find('size/height').text)

        # 添加图像信息到 COCO 数据结构
        image_info = {
            "id": count,
            "file_name": f"{image_id}.jpg",
            "height": height,
            "width": width
        }
        coco_dict["images"].append(image_info)

        # 遍历每个物体标注
        for obj in root.findall('object'):
            # 提取物体类别和边界框信息
            class_name = obj.find('name').text
            bbox = obj.find('bndbox')
            xmin, ymin, xmax, ymax = [int(bbox.find(t).text) for t in ['xmin', 'ymin', 'xmax', 'ymax']]
            if class_name not in class_map.keys():
                continue
            category_id = class_map[class_name]
            box_w, box_h = xmax - xmin, ymax - ymin
            # 添加物体标注到 COCO 数据结构
            annotation_info = {
                "id": len(coco_dict["annotations"]),
                "image_id": count,
                "category_id": category_id,
                "bbox": [xmin, ymin, box_w, box_h],
                "area": box_w * box_h,
                "iscrowd": 0
            }
            coco_dict["annotations"].append(annotation_info)

    # 添加分类信息到 COCO 数据结构
    for class_name, class_id in class_map.items():
        category_info = {
            "id": class_id,
            "name": class_name,
            "supercategory": ""
        }
        coco_dict["categories"].append(category_info)

    return coco_dict


if __name__ == '__main__':
    txt = r'/xxx/name_id.txt'
    xml_dir = 'xxxx/Custom_coco/train/images/'
    json_file = 'xxxx/train.json'
    xml2coco(txt, xml_dir, json_file)

其中：

txt中内容格式如下，即标注对应的数字：

'tv':0
'dog':1
'fish':2

xml_dir为voc格式的xml所在路径

json_file为保存json所在路径