把YOLO标注文件转换为COCO格式.md

实验室之前用的模型是基于Darknet的YOLO。这两天需要试一些其它模型,所以我准备给标注文件做个格式转换。网上找了半天都找不到个给人用的YOLO转COCO脚本,遂自己写了一个。分享一下,大家可能用得上。

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
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
# 原先YOLO标注存放结构:
# - ...
# |
# +- FOLDER # 存放所有图片和对应的标注
# | |
# | + a.jpg # 图片
# | |
# | + a.txt # 图片对应的标注文件
# |
# + train.txt # 记录训练/测试/验证的分割文件,里面包含FOLDER下的图片文件名
# -----------------------------------------------------------------
# 转化后在 FOLDER 同级下添加train.json,存放所有 train.txt 包含的文件名的图片对应的标注。其它不作改动,方便兼容。

import cv2
import json
import os.path as osp

# ----使用前在此设置----
origDir = "../dataset/" # 即上面的 FOLDER
origList = "../dataset/train_split.txt" # 即上面的 train.txt
dumpDir = "../dataset/train.json" # 即上面的 train.json
CLASSES = [] # 在这里添加所有要用的类

# ----工具函数----
global id
id = 0
global anno_id
annoId = 0
def incrId():
global id
id += 1
return id
def incrAnnoId():
global annoId
annoId += 1
return annoId

# ----生成COCO下标注文件的Dict----
COCO_OBJ = {}
COCO_OBJ["categories"] = [
{"id": incrId(), "name": x, "supercategory": ""} for x in CLASSES]
COCO_OBJ["images"] = []
COCO_OBJ["annotations"] = []
id = 0
annoId = 0
fileList = []
with open(origList, "r") as f:
fileList = [x.replace("\n", "") for x in f.readlines()] # 根据 train.txt 的格式标准,你可能需要改动这一行来适应。
for img in fileList:
if (".jpg" in img):
print(img)
anno = img.replace(".jpg", ".txt")
annoPath = osp.join(origDir, anno)
imgPath = osp.join(origDir, img)
imgObj = cv2.imread(imgPath)
width = imgObj.shape[1]
height = imgObj.shape[0]
imgStruct = { # 有需求自行在此添加
"id": incrId(),
"width": width,
"height": height,
"file_name": img,
"license": 0,
"flickr_url": "",
"coco_url": "",
"date_captured": 0
}
COCO_OBJ["images"].append(imgStruct)
annoLines = []
with open(annoPath, "r") as a:
annoLines = a.readlines()
for annoLine in annoLines:
annoLine = annoLine.replace("\n", "")
annoComponents = annoLine.split(" ")
if len(annoComponents) < 5:
continue
c_x = float(annoComponents[1])*width
c_y = float(annoComponents[2])*height
w = float(annoComponents[3])*width
h = float(annoComponents[4])*height
annoStruct = { # 有需求自行在此添加
"id": incrAnnoId(),
"image_id": id,
"category_id": int(annoComponents[0]),
"segmentation": [],
"area": float(int(w)*int(h)),
"bbox": [
float(int(c_x - w/2) + 1), # x
float(int(c_y - h/2) + 1), # y
float(int(w) - 1), # w
float(int(h) - 1) # h
]
}
COCO_OBJ["annotations"].append(annoStruct)

with open(dumpDir, "w") as f:
json.dump(COCO_OBJ, f)

这里我对坐标转换的处理取了个上下整,尽量保证不溢出图像边缘。