将COCO格式的json文件转化为DOTA格式标签的代码实现

用oriented bounding box网络训练自己的数据集时，发现基本是针对DOTA数据集格式的，自己的数据集是labelme多边形标注，再转化成COCO格式的，写了一个小脚本把它再转化成DOTA格式。 

DOTA数据集标签：

保存在具有相同文件名的文本文件中（image_name.txt）

标签格式为：

 

 

x1 y1 x2 y2 x3 y3 x4 y4 category difficult

转化思路：

1、获取COCO.json文件中的“annotation”信息，保存category_id、image_id、annotation_id；

2、通过image_id从json中的“images”里获取到文件名；通过category_id和自己数据集里的类别矩阵our_category获取该实例对应的类别；

3、对多边形轮廓点求最小外接矩形；

4、输出矩形四个坐标点的取整值、类别名称，difficult统一设置为0.

代码实现：

import json
import numpy as np
import cv2

our_category = ['BG', 'single', 'dashed', 'double', 'arrow', 'stop', 'crosswalk']

with open('./data/annotations/json/train_data.json', 'r', encoding="utf-8") as file:
    json_data = json.load(file)

for ann in json_data["annotations"]:
    category_id = ann["category_id"]
    category = our_category[category_id]
    image_id = ann["image_id"]
    ann_id = ann["id"]

    for img_file in json_data["images"]:
        if img_file["id"] == image_id:
            img_file_name = img_file["file_name"]

    seg_info = ann["segmentation"][0]
    length = len(seg_info)
    contours = []
    for i in range(length):
        if i % 2 == 0:
            contours.append([seg_info[i], seg_info[i+1]])

    contours = np.array(contours, dtype=np.float32)
    # 对每个轮廓点求最小外接矩形
    rect = cv2.minAreaRect(contours)
    # cv2.boxPoints可以将轮廓点转换为四个角点坐标
    box = cv2.boxPoints(rect)
    # 令四个角点坐标为顺时针
    startidx = box.sum(axis=1).argmin()
    box = np.roll(box, 4 - startidx, 0)
    # 在原图上画出预测的外接矩形
    # img_dir_path = 'data/train/'
    # image = cv2.imread(os.path.join(img_dir_path, img_file_name))
    # box = box.reshape((-1, 1, 2)).astype(np.int32)
    # img = cv2.polylines(image, [box], True, (0, 255, 0), 10)
    # cv2.imwrite('data/visual/' + str(image_id) + '_' + str(ann_id) + '.png', img)

    file = open('data/annotations/DOTA_txt/' + img_file_name.split(".")[0] + '.txt', 'a')
    for dot in box:
        d0 = int(dot[0])
        d1 = int(dot[1])
        file.write(str(d0))
        file.write(' ')
        file.write(str(d1))
        file.write(' ')

    file.write(category)
    file.write(' ')
    file.write('0\n')