欢迎访问本站,希望站内信息可以帮助到您!感兴趣的,可以加入我的知识星球,一起学习,一起进步!

COCO 格式转 YOLO

YOLO 迷途小书童 1评论

环境

前言

前文 MS COCO数据集 介绍过了 COCO 数据集,COCO 是将所有图片对应的标注信息写在了一个 json 文件里,如下

coco to yolo

因此要将 coco 格式的数据集转换成 yolo 格式,本质上就是去解析这个 json 文件,将对应图片的标注信息提取出来,写入 txt 文件中

实操

这里还是使用我们熟悉的 roboflow 平台上的口罩数据集,下载地址是 https://public.roboflow.com/object-detection/mask-wearing/4

下载的格式选择 COCO

coco to yolo

解压后,进入 train 文件夹,可以看到 json 文件和图片文件

coco to yolo

在同级目录创建 python 脚本 coco2yolo.py,输入如下代码,少量注释嵌入到代码中了

import os
import json
from tqdm import tqdm
import argparse

def convert(size, box):
    '''
    size: 图片的宽和高(w,h)
    box格式: x,y,w,h
    返回值:x_center/image_width y_center/image_height width/image_width height/image_height
    '''

    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = box[0] + box[2] / 2.0
    y = box[1] + box[3] / 2.0
    w = box[2]
    h = box[3]

    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('--json_file', default='test.json',
                        type=str, help="coco file path")
    parser.add_argument('--save_dir', default='labels', type=str,
                        help="where to save .txt labels")
    arg = parser.parse_args()

    data = json.load(open(arg.json_file, 'r'))

    # 如果存放txt文件夹不存在,则创建
    if not os.path.exists(arg.save_dir):
        os.makedirs(arg.save_dir)

    id_map = {}

    # 解析目标类别,也就是 categories 字段,并将类别写入文件 classes.txt 中
    with open(os.path.join(arg.save_dir, 'classes.txt'), 'w') as f:
        for i, category in enumerate(data['categories']):
            f.write(f"{category['name']}\n")
            id_map[category['id']] = i

    for img in tqdm(data['images']):

        # 解析 images 字段,分别取出图片文件名、图片的宽和高、图片id
        filename = img["file_name"]
        img_width = img["width"]
        img_height = img["height"]
        img_id = img["id"]
        head, tail = os.path.splitext(filename)

        # txt文件名,与对应图片名只有后缀名不一样
        txt_name = head + ".txt"
        f_txt = open(os.path.join(arg.save_dir, txt_name), 'w')

        for ann in data['annotations']:
            if ann['image_id'] == img_id:
                box = convert((img_width, img_height), ann["bbox"])

                # 写入txt,共5个字段
                f_txt.write("%s %s %s %s %s\n" % (
                    id_map[ann["category_id"]], box[0], box[1], box[2], box[3]))

        f_txt.close()

执行上述代码

python coco2yolo.py --json_file _annotations.coco.json --save_dir labels

完成后,还是在同级目录,就能看到 labels 文件夹及里面的 txt 标注文件了

coco to yolo

关联阅读

喜欢 (1)

您必须 登录 才能发表评论!

(1)个小伙伴在吐槽
  1. 报错了个KeyError:'categories'
    匿名2022-04-14 21:19