将coco格式数据集转换为yolo格式数据集
将coco格式的数据集转换成yolo格式
1. YOLO格式数据集和COCO格式的区别
yolo数据集的标注文件是.txt文件,在label文件夹中每一个.txt文件对应数据集中的一张图片
其中每个.txt文件中的每一行代表图片中的一个目标。
而coco数据集的标注文件是.json文件,全部的数据标注文件由三个.json文件组成:train.json val.json test.json,其中每个.json文件中包含全部的数据集图片中的所有目标(注意是所有目标不是数据集中的所有张图片)
YOLO格式数据集和COCO数据集的文件夹组成
COCO数据集:
|-dataset-cooc
|-|-annotations
|-|-|-train.json
|-|-|-val.json
|-|-|-test.json
|-images
|-|-001.jpg
|-|-002.jpg
…
YOLO数据集:(每个.txt文件和.jpg文件是一一对应的)
|-dataset-yolo
|-|-labels
|-|-|-train
|-|-|-|-001.txt
|-|-|-val
|-|-|-|-006.txt
|-|-|-test
|-|-|-|-009.txt
|-|-images
|-|-|-train
|-|-|-|-001.jpg
|-|-|-val
|-|-|-|-006.jpg
|-|-|-test
|-|-|-|-009.jpg
…
将YOLO格式 的数据集转换为COCO格式
注意:
COCO2017官方的数据集中categories id 是不连续的,这在yolo读取的时候会出问题,所以需要重新映射一下,这个代码会按id从小到大映射到0~79之间。(如果是自己的数据集,也会重新映射)
代码:
"""
author: Wu
https://github.com/Weifeng-Chen/DL_tools/issues/3
2021/1/24
COCO 格式的数据集转化为 YOLO 格式的数据集,源代码采取遍历方式,太慢,
这里改进了一下时间复杂度,从O(nm)改为O(n+m),但是牺牲了一些内存占用
--json_path 输入的json文件路径
--save_path 保存的文件夹名字,默认为当前目录下的labels。
"""
import os
import json
from tqdm import tqdm
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--json_path', default='./instances_val2017.json',type=str, help="input: coco format(json)")
parser.add_argument('--save_path', default='./labels', type=str, help="specify where to save the output dir of labels")
arg = parser.parse_args()
def convert(size, box):
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__':
json_file = arg.json_path # COCO Object Instance 类型的标注
ana_txt_save_path = arg.save_path # 保存的路径
data = json.load(open(json_file, 'r'))
if not os.path.exists(ana_txt_save_path):
os.makedirs(ana_txt_save_path)
id_map = {} # coco数据集的id不连续!重新映射一下再输出!
for i, category in enumerate(data['categories']):
id_map[category['id']] = i
# 通过事先建表来降低时间复杂度
max_id = 0
for img in data['images']:
max_id = max(max_id, img['id'])
# 注意这里不能写作 [[]]*(max_id+1),否则列表内的空列表共享地址
img_ann_dict = [[] for i in range(max_id+1)]
for i, ann in enumerate(data['annotations']):
img_ann_dict[ann['image_id']].append(i)
for img in tqdm(data['images']):
filename = img["file_name"]
img_width = img["width"]
img_height = img["height"]
img_id = img["id"]
head, tail = os.path.splitext(filename)
ana_txt_name = head + ".txt" # 对应的txt名字,与jpg一致
f_txt = open(os.path.join(ana_txt_save_path, ana_txt_name), 'w')
'''for ann in data['annotations']:
if ann['image_id'] == img_id:
box = convert((img_width, img_height), ann["bbox"])
f_txt.write("%s %s %s %s %s\n" % (id_map[ann["category_id"]], box[0], box[1], box[2], box[3]))'''
# 这里可以直接查表而无需重复遍历
for ann_id in img_ann_dict[img_id]:
ann = data['annotations'][ann_id]
box = convert((img_width, img_height), ann["bbox"])
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()
# 旧版,很慢hhh
# """
# COCO 格式的数据集转化为 YOLO 格式的数据集
# --json_path 输入的json文件路径
# --save_path 保存的文件夹名字,默认为当前目录下的labels。
# """
# import os
# import json
# from tqdm import tqdm
# import argparse
# parser = argparse.ArgumentParser()
# parser.add_argument('--json_path', default='./instances_val2017.json',type=str, help="input: coco format(json)")
# parser.add_argument('--save_path', default='./labels', type=str, help="specify where to save the output dir of labels")
# arg = parser.parse_args()
# def convert(size, box):
# 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__':
# json_file = arg.json_path # COCO Object Instance 类型的标注
# ana_txt_save_path = arg.save_path # 保存的路径
# data = json.load(open(json_file, 'r'))
# if not os.path.exists(ana_txt_save_path):
# os.makedirs(ana_txt_save_path)
# id_map = {} # coco数据集的id不连续!重新映射一下再输出!
# with open(os.path.join(ana_txt_save_path, 'classes.txt'), 'w') as f:
# # 写入classes.txt
# for i, category in enumerate(data['categories']):
# f.write(f"{category['name']}\n")
# id_map[category['id']] = i
# # print(id_map)
# for img in tqdm(data['images']):
# filename = img["file_name"]
# img_width = img["width"]
# img_height = img["height"]
# img_id = img["id"]
# head, tail = os.path.splitext(filename)
# ana_txt_name = head + ".txt" # 对应的txt名字,与jpg一致
# f_txt = open(os.path.join(ana_txt_save_path, ana_txt_name), 'w')
# for ann in data['annotations']:
# if ann['image_id'] == img_id:
# box = convert((img_width, img_height), ann["bbox"])
# 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_path $JSON_FILE_PATH --save_path $LABEL_SAVE_PATH
$JSON_FILE_PATH是json文件的地址。
$JSON_FILE_PATH是输出目录(默认为工作目录下的./labels目录。
参考链接:
更多推荐
所有评论(0)