16

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

 3 years ago
source link: https://enoch2090.me/2021/01/06/tech/%E6%8A%8AYOLO%E6%A0%87%E6%B3%A8%E6%96%87%E4%BB%B6%E8%BD%AC%E6%8D%A2%E4%B8%BACOCO%E6%A0%BC%E5%BC%8F/
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.
neoserver,ios ssh client

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

Enoch2090

2021-01-06

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

# 原先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)

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


report

Recommend

About Joyk


Aggregate valuable and interesting links.
Joyk means Joy of geeK