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

,

实验室之前用的模型是基于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)

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