YOLO格式标注数据转COCO标注数据

这里仅仅考虑person类别，如果考虑其他类别，则需要增加类别信息，稍作调整即可。

import json
import os
import imagesize
import copy


def txt_to_json(img_dir,annotation_dir,json_path,img_format='.jpg',annotation_format='.txt'):
    # json 文件主要两项内容
    json_dict = dict()
    annotations = list()
    images = list()
    categories = list()
    # 一个标签和一张图
    one_annotation = dict()
    one_image = dict()
    annotation_bbox_id = 0
    for file in os.listdir(annotation_dir):
        if file.endswith(annotation_format):
            # 读取图片信息：长宽，整合到one_image当中
            one_image['file_name'] = file.split('.')[0]+img_format
            one_image['id'] = file.split('.')[0]
            one_image['width'],one_image['height'] = imagesize.get(os.path.join(img_dir,one_image['file_name']))
            # 读取txt文件的内容，for循环整合到one_annotation当中
            with open(os.path.join(annotation_dir,file), 'r') as f:
                for line in f.readlines():
                    line = line.strip('\n')  # 去掉每一行中的换行符
                    [categories_id,x,y,w,h] = line.split(' ')
                    w = one_image['width']*float(w)
                    h = one_image['height']*float(h)
                    x_min = one_image['width']*float(x) - w/2.0
                    y_min = one_image['height']*float(y) - h/2.0
                    one_annotation['segmentation'] = []
                    one_annotation['area'] = w*h
                    one_annotation['iscrowd'] = 0
                    one_annotation['image_id'] = one_image['id']
                    one_annotation["bbox"] = [x_min,y_min,w,h]
                    one_annotation["category_id"] = 1
                    one_annotation["id"] = annotation_bbox_id                   # 这里的id就是一个编号，每一个人的编号都不相同
                    annotation_bbox_id = annotation_bbox_id + 1
                    annotations.append(copy.deepcopy(one_annotation))
            images.append(copy.deepcopy(one_image))
    category = dict()
    category['supercategory'] = 'person'
    category['person'] = 'person'
    category['id'] = 1
    categories.append(category)
    json_dict['annotations'] = annotations
    json_dict['images'] = images
    json_dict['categories'] = categories
    # 将获取的xml内容写入到json文件中
    with open(json_path, 'w') as f:
        f.write(json.dumps(json_dict, indent=1, separators=(',', ':')))


if __name__ == '__main__':
    img_dir = './train/'
    annotation_dir = './train/'
    json_path = './annotations/train.json'
    txt_to_json(img_dir,annotation_dir,json_path)