将yolov5检测结果保存成json文件
将yolov5检测结果保存成json文件,每帧图片对应一个json文件,json文件中包含图片名称、检测到的bbox信息(分类结果、中心坐标和置信度)。函数json_add以图片名、该图像内的所有bbox信息,以及bbox数量为输入。在yolov5的detect.py里获取每帧中的bbox信息和数量后,直接调用函数json_add即可。
# -*- coding: UTF-8 -*-
import os
import json
import random
import time
from PIL import Image
class json_result:
def __init__(self):
self.W=1920
self.H=1080
self.save_json_path=./task11_result/师兄不上自己上/
self.class_name=["pedestrain","non-motor-vehicle","motor-vehicle"]
def get_time_stamp(self):
ct = time.time()
local_time = time.localtime(ct)
data_head = time.strftime("%Y-%m-%d %H:%M:%S", local_time)
data_secs = (ct - int(ct)) * 1000
time_stamp = "%s.%03d" % (data_head, data_secs)
#print(time_stamp)
stamp = ("".join(time_stamp.split()[0].split("-"))+"".join(time_stamp.split()[1].split(":"))).replace(., )
#print(stamp)
return stamp
def json_add(self,name,data,count):
write_json_context=dict()
write_json_context[info]=[]
write_json_context[annotations]=[]
#image name
img_context={}
img_context[image_name]=name
write_json_context[info].append(img_context)
for i in range(count):
bbox_dict = {}
cls = data[i][0]
x = data[i][1]
y = data[i][2]
w = data[i][3]
h = data[i][4]
score = data[i][5]
cls,x,y,w,h,score=int(cls),float(x),float(y),float(w),float(h),float(score)
x=x*self.W
y=y*self.H
w=w*self.W
h=h*self.H
bbox_dict[category_name]=self.class_name[cls]
bbox_dict[score]=score
bbox_dict[bbox]=[x,y,w,h]
write_json_context[annotations].append(bbox_dict)
json_name = str(os.path.splitext(name)[0])
json_name=json_name +.json
json_name = os.path.join(self.save_json_path,json_name)
with open(json_name,w) as fw:
json.dump(write_json_context,fw)
if __name__==__main__:
j=json_result()
t1=float(j.get_time_stamp())
print(t1)
name=1998.png
count=2
data=[[0,0.940625,0.208333,0.11875,0.238889,0.5],[0,0.933333,0.193056,0.13125,0.225,0.7]]
j.json_add(name,data,count)
t2=float(j.get_time_stamp())
print(t2)
t=float(t2-t1)
print(time:,t)
