VOC (xml格式)转COCO(json格式)代码

VOC (xml格式)转COCO(json格式)代码如下：

import json,cv2,copy,shutil,sys,os,tqdm
import random

import numpy as np

import xml.etree.ElementTree as ET
from xml.etree.ElementTree import Element

# def show(img,bbox):
#     img_show=img.copy()
#     img_show = cv2.rectangle(img_show, (int(bbox[0]),int(bbox[1])), (int(bbox[0])+int(bbox[2]),int(bbox[1])+int(bbox[3])),color=(0,0,255), thickness=10)
#     cv2.imshow('picture',img_show)
#     cv2.waitKey(200)

if __name__=="__main__":
    # cv2.namedWindow('picture', 0)
    f_json = open("/YOLOX_model/annotations/instances_train2017.json",'w')
    sample = json.load(open(r'/YOLOX_model/instances_val2017.json','r'))
    my_dict ={}
    my_dict['info']=sample['info']
    my_dict['licenses'] = sample['licenses']
    my_dict['images']=[]
    my_dict['annotations']=[]

    my_dict['categories'] =[
        {'supercategory': 'smoke', 'id': 1, 'name': 'smoke'},
        {'supercategory':'fire','id':2,'name':'fire'},
    ]
    img_id = 0
    ann_id = 0
    root = "/YOLOX_model/JPEGImages/"     #图片的路径 
    #xml  = "/YOLOX_model/Annotations/"  xml的路径  
    img_list = os.listdir(root)
    img_list =[root +v for v in img_list if (v.endswith('.jpg') or v.endswith('.png') or v.endswith('.jpeg') )]

    for img_path in tqdm.tqdm(img_list):
        label_path = img_path.replace('/JPEGImages/','/Annotations/').replace('.jpg','.xml').replace('.png','.xml').replace('.jpeg','.xml')
        if os.path.exists(label_path) is False:
            print('%s does not exists'%label_path)
            continue
        ##image
        img = cv2.imread(img_path)
        if img is None:
            print('%s  read img failed. None '%img_path)
            continue
        image_dict = {'license': 4, 'file_name': img_path.replace(root,''), 'coco_url': '', 'height': img.shape[0], 'width': img.shape[1],
             'date_captured': '', 'flickr_url': '', 'id': img_id}
        is_related = True
        ##annotation
        tree = ET.parse(label_path)
        xml_root = tree.getroot()
        for obj in xml_root.findall('object'):
            name = obj.find('name').text
            if name not in ['smoke','fire']:
                print('%s--%s'%(label_path,name))
                continue
            if name=='smoke':
                cat_id = 1
            elif name=='fire':
                cat_id=2
            bnd_box = obj.find('bndbox')
            x_min= float(bnd_box.find('xmin').text)
            y_min= float(bnd_box.find('ymin').text)
            w=  float(bnd_box.find('xmax').text)-float(bnd_box.find('xmin').text)
            h=  float(bnd_box.find('ymax').text)-float(bnd_box.find('ymin').text)

            label_dict = {'segmentation': [], 'area':w*h, 'iscrowd': 0, 'image_id': img_id,
                          'bbox': [x_min, y_min, w, h], 'category_id': cat_id, 'id': ann_id}
            my_dict['annotations'].append(copy.deepcopy(label_dict))
            is_related=True
            ann_id+=1
        if is_related:  ##
            my_dict['images'].append(copy.deepcopy(image_dict))
        img_id +=1

    f_json.write(json.dumps(my_dict))
    f_json.flush()

#样本案例需要一个coco格式文件：
链接：https://download.youkuaiyun.com/download/weixin_45734021/90042374

sample = json.load(open(r'/YOLOX_model/instances_val2017.json','r'))

这个代码需要改动几个地方：保存的json文件地址，案例的json文件地址，类别。
xml保存到Annotations文件中，图片保存到JPEGImages文件中