有时候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)
当然这是将旋转框数据集转换为正框数据集,会出现很多误差,但是比如yolov5-obb一些算法计算iou时都是采用的诸如此类的方法。