用法:可用于JSON格式的数据转换成YOLO训练格式(e.g. YOLOv5等)。只需要输入json文件路径和需要保存txt的路径即可。
import json
import argparse
# json中bbox的格式是top_x, top_y, w, h
# 转成Yolo的格式是cen_x, cen_y, w, h的相对位置
def convert(img_size, box):
dw = 1. / (img_size[0])
dh = 1. / (img_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)
def main(args):
# 读取 json 文件数据
with open(args.json_dir, 'r') as load_f:
content = json.load(load_f)
img_hw = dict()
for k in content['images']:
file_name = k['id']
img_hw[file_name] = [k["width"], k["height"]]
# 循环处理
for t in content['annotations']:
tmp = t['image_id']
filename = args.out_dir + tmp + '.txt'
width = img_hw[tmp][0]
height = img_hw[tmp][1]
# import pdb;pdb.set_trace()
with open(filename, mode='w') as fp:
# 计算 yolo 数据格式所需要的中心点的 相对 x, y 坐标, w,h 的值
bbox = convert((width, height), t['bbox'])
fp = open(filename, mode="r+", encoding="utf-8")
file_str = str(t['category_id']) + ' ' + ' '.join([str(a) for a in bbox])
line_data = fp.readlines()
if len(line_data) != 0:
fp.write('\n' + file_str)
else:
fp.write(file_str)
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Start convert.')
parser.add_argument('--json_file', type=str, help='json file path')# json文件路径
parser.add_argument('--output', type=str, help='output path', default='')# 输出的 txt 文件路径
args = parser.parse_args()
main(args)
更新于2022.3.17,上面的代码可能会出现生成的txt文件打不开的情况。更新的代码修改了保存逻辑,先转成list,再转成tuple,最后再转成txt文件保存需要的string。
import json
import argparse
import numpy as np
# json中bbox的格式是top_x, top_y, w, h
# 转成Yolo的格式是cen_x, cen_y, w, h的相对位置
def convert(img_size, box):
dw = 1. / (img_size[0])
dh = 1. / (img_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)
def main(args):
# 读取 json 文件数据
with open(args.json_file, 'r') as load_f:
content = json.load(load_f)
img_hw = dict()
for k in content['images']:
file_name = k['id']
# file_name = k['file_name'].rsplit('.',1)[0]
img_hw[file_name] = [k["width"], k["height"], k["file_name"].rsplit('.', 1)[0]]
# 循环处理
for t in content['annotations']:
tmp = t['image_id']
filename = args.output + img_hw[tmp][2] + '.txt'
width = img_hw[tmp][0]
height = img_hw[tmp][1]
# 计算 yolo 数据格式所需要的中心点的 相对 x, y 坐标, w,h 的值
bbox = convert((width, height), t['bbox'])
box = np.array(bbox, dtype=np.float64)
categroy_id = t['category_id']
if box[2] > 0 and box[3] > 0:
line = tuple([categroy_id] + box.tolist())
with open(filename, mode='a') as fp:
fp.write(('%g ' * len(line)).rstrip() % line + '\n')
else:
print('bbox error')
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Start convert.')
parser.add_argument('--json_file', type=str, help='json file path')# json文件路径
parser.add_argument('--output', type=str, help='output path', default='')# 输出的 txt 文件路径
args = parser.parse_args()
main(args)