File size: 1,526 Bytes
7206ed3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
import json
import numpy as np
import pdb
dict_pred = {0: 'animal', 1: 'person', 2: 'vehicle'}
def save_results(md_results, dlc_outputs,map_label_id_to_str,thr,output_file = 'dowload_predictions.json'):
"""
write json
"""
info = {}
## info megaDetector
info['file']= md_results.files[0]
number_bb = len(md_results.xyxy[0].tolist())
info['number_of_bb'] = number_bb
number_bb_thr = len(dlc_outputs)
labels = [n for n in map_label_id_to_str.values()]
#pdb.set_trace()
new_index = []
for i in range(number_bb):
corner_x1,corner_y1,corner_x2,corner_y2,confidence, _ = md_results.xyxy[0].tolist()[i]
if confidence > thr:
new_index.append(i)
for i in range(number_bb_thr):
aux={}
corner_x1,corner_y1,corner_x2,corner_y2,confidence, _ = md_results.xyxy[0].tolist()[new_index[i]]
aux['corner_1'] = (corner_x1,corner_y1)
aux['corner_2'] = (corner_x2,corner_y2)
aux['predict MD'] = md_results.names[0]
aux['confidence MD'] = confidence
## info dlc
kypts = []
for s in dlc_outputs[i]:
aux1 = []
for j in s:
aux1.append(float(j))
kypts.append(aux1)
aux['dlc_pred'] = dict(zip(labels,kypts))
info['bb_' + str(new_index[i]) ]=aux
with open(output_file, 'w') as f:
json.dump(info, f, indent=1)
print('Output file saved at {}'.format(output_file))
return output_file
|