【目标检测数据集处理】DOTA数据集转换coco格式

有时候dota数据集格式无法使用，需要转换为coco数据集。代码如下

import os
import cv2
import json
from tqdm import tqdm
import numpy as np


def txt2xywh(txt_file_path, image_width, image_height):
    with open(txt_file_path, 'r') as f:
        lines = f.readlines()
    boxes = []
    for line in lines:
        parts = line.strip().split(',')
        category = parts[0].split()
        x1, y1, x2, y2, x3, y3, x4, y4, confidence = category[1:]
        x1 = float(x1)
        x2 = float(x2)
        x3 = float(x3)
        x4 = float(x4)
        y1 = float(y1)
        y2 = float(y2)
        y3 = float(y3)
        y4 = float(y4)
        xmin = float(min(x1, x2, x3, x4))
        xmax = float(max(x1, x2, x3, x4))
        ymin = float(min(y1, y2, y3, y4))
        ymax = float(max(y1, y2, y3, y4))
        width = xmax - xmin
        height = ymax - ymin
        x_min = xmin / image_width
        y_min = ymin / image_height
        width = width / image_width
        height = height / image_height
        boxes.append([x_min, y_min, width, height])
    return int(category[0]), float(confidence), boxes


def yolo2txt(originLabelsDir, originImagesDir, saveTempTxt, img_suffix='png'):
    """
    将yolo的标签全部存入一个txt文件中
    将yolo格式的标签：classId, xCenter, yCenter, w, h转换为
    coco格式：classId, xMin, yMim, xMax, yMax格式
    coco的id编号从1开始计算，所以这里classId应该从1开始计算
    最终annos.txt中每行为imageName, classId, xMin, yMim, xMax, yMax, 一个bbox对应一行
    originLabelsDir: 原始yolo标签路径
    originImagesDir: 原始图像路径
    ssaveTempTxt: 保存txt路径
    )
    """
    txtFileList = os.listdir(originLabelsDir)
    with open("data.json", "w") as f:
        # data = json.load(f)
        for txtFile in tqdm(txtFileList, desc="generating COCO format"):
            classes, score, bbox = txt2xywh(originLabelsDir+'/'+txtFile, 1024, 1024)
            data={
                "image_id": txtFile,
                "bbox":bbox[0],
                "score":score,
                "category_id":classes
            }
            json.dump(data,f)
            break
        


if __name__ == "__main__":
    # 原始标签路径
    originLabelsDir = '/Users/aoxin/CODE/python/dota2coco/labels'
    # 转换后的文件保存路径
    saveTempTxt = '/Users/aoxin/CODE/python/dota2coco/annos_temp.txt'
    # 原始标签对应的图片路径
    originImagesDir = '/Users/aoxin/CODE/python/dota2coco/labels' # meaningless
    yolo2txt(originLabelsDir, originImagesDir, saveTempTxt)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74

当然这是将旋转框数据集转换为正框数据集，会出现很多误差，但是比如yolov5-obb一些算法计算iou时都是采用的诸如此类的方法。