[人工智能]将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)