【转换格式脚本】COCO格式转YOLO

用法：可用于JSON格式的数据转换成YOLO训练格式（e.g. YOLOv5等）。只需要输入json文件路径和需要保存txt的路径即可。

import json
import argparse


# json中bbox的格式是top_x, top_y, w, h 
# 转成Yolo的格式是cen_x, cen_y, w, h的相对位置
def convert(img_size, box):
    dw = 1. / (img_size[0])
    dh = 1. / (img_size[1])
    x = box[0] + box[2] / 2.0 
    y = box[1] + box[3] / 2.0 
    w = box[2]
    h = box[3]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)

def main(args):
    # 读取 json 文件数据
    with open(args.json_dir, 'r') as load_f:
        content = json.load(load_f)

    img_hw = dict()
    for k in content['images']:
        file_name = k['id']
        img_hw[file_name] = [k["width"], k["height"]]
    # 循环处理
    for t in content['annotations']:
        tmp = t['image_id']
        filename = args.out_dir + tmp + '.txt'

        width = img_hw[tmp][0]
        height = img_hw[tmp][1]
        # import pdb;pdb.set_trace()
        
        with open(filename, mode='w') as fp:
            # 计算 yolo 数据格式所需要的中心点的 相对 x, y 坐标, w,h 的值
            bbox = convert((width, height), t['bbox'])
            fp = open(filename, mode="r+", encoding="utf-8")
            file_str = str(t['category_id']) + ' ' + ' '.join([str(a) for a in bbox])
            line_data = fp.readlines()

            if len(line_data) != 0:
                fp.write('\n' + file_str)
            else:
                fp.write(file_str)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Start convert.')
    parser.add_argument('--json_file', type=str, help='json file path')# json文件路径
    parser.add_argument('--output', type=str, help='output path', default='')# 输出的 txt 文件路径
    args = parser.parse_args()
    main(args)

更新于2022.3.17，上面的代码可能会出现生成的txt文件打不开的情况。更新的代码修改了保存逻辑，先转成list，再转成tuple，最后再转成txt文件保存需要的string。

import json
import argparse
import numpy as np


# json中bbox的格式是top_x, top_y, w, h 
# 转成Yolo的格式是cen_x, cen_y, w, h的相对位置
def convert(img_size, box):
    dw = 1. / (img_size[0])
    dh = 1. / (img_size[1])
    x = box[0] + box[2] / 2.0 
    y = box[1] + box[3] / 2.0 
    w = box[2]
    h = box[3]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)

def main(args):
    # 读取 json 文件数据
    with open(args.json_file, 'r') as load_f:
        content = json.load(load_f)

    img_hw = dict()
    for k in content['images']:
        file_name = k['id']
        # file_name = k['file_name'].rsplit('.',1)[0]
        img_hw[file_name] = [k["width"], k["height"], k["file_name"].rsplit('.', 1)[0]]
    # 循环处理
    for t in content['annotations']:
        tmp = t['image_id']
        filename = args.output + img_hw[tmp][2] + '.txt'

        width = img_hw[tmp][0]
        height = img_hw[tmp][1]

        # 计算 yolo 数据格式所需要的中心点的 相对 x, y 坐标, w,h 的值
        bbox = convert((width, height), t['bbox'])
        box = np.array(bbox, dtype=np.float64)

        categroy_id = t['category_id']

        if box[2] > 0 and box[3] > 0:
            line = tuple([categroy_id] + box.tolist())
        
            with open(filename, mode='a') as fp:
                fp.write(('%g ' * len(line)).rstrip() % line + '\n')

        else:
            print('bbox error')
            
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Start convert.')
    parser.add_argument('--json_file', type=str, help='json file path')# json文件路径
    parser.add_argument('--output', type=str, help='output path', default='')# 输出的 txt 文件路径
    args = parser.parse_args()
    main(args)