-
Notifications
You must be signed in to change notification settings - Fork 94
/
Copy pathconvert_yolo_prediction_to_caltech.py
88 lines (58 loc) · 2.81 KB
/
convert_yolo_prediction_to_caltech.py
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
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
from os import listdir
from os.path import isfile, join
import argparse
import cv2
import numpy as np
import sys
import os
import shutil
def save_file_predictions(root,predictions,stream):
for prediction in predictions:
(confidence,xmin,ymin,xmax,ymax) = map(float,prediction.split(' ')[1:]) #leaving out [0] because it corresponds to ID
w = xmax - xmin
h = ymax - ymin
stream.write('%f %f %f %f %f \n'%(xmin,ymin,w,h,confidence))
def open_file(root,id):
set_folder_name,video_folder_name,image_file_name = id.split('\\')[-3:]
#print 'image_file_name = %s'%(image_file_name)
#print 'video_folder_name = %s'%(video_folder_name)
#print 'set_folder_name = %s'%(set_folder_name)
video_dir = join(root,set_folder_name,video_folder_name)
if not os.path.isdir(video_dir):
os.makedirs(video_dir)
f = open(join(video_dir,image_file_name+'.txt'),'w')
print join(video_dir,image_file_name+'.txt')
return f
def main(argv):
parser = argparse.ArgumentParser()
parser.add_argument('-predictions_file', default = 'C:\\darknet_fire_detection\\build\\darknet\\x64\\results\\comp-caltech-voc2-8K.txt', # This dir name I cannot change cuz it will cause confusion
help='path to yolo predictions\n', )
parser.add_argument('-caltech_pedestrian_root', default = 'F:\\dataset\\CaltechPedestrians\\code\\data-USA\\res\\comp-caltech-voc2-8K',
help='where to save converted predictions')
args = parser.parse_args()
print "predictions_file you provided {}".format(args.predictions_file )
print "caltech_pedestrian_root you provided is {}".format(args.caltech_pedestrian_root)
if not os.path.exists(args.caltech_pedestrian_root):
os.mkdir(args.caltech_pedestrian_root)
f = open(args.predictions_file)
lines = [line.rstrip('\n') for line in f.readlines()]
subdirs = os.listdir(args.caltech_pedestrian_root)
if len(subdirs) >5:
print '%s is non empty please specify empty dir'%(args.caltech_pedestrian_root)
print subdirs
return
else:
print '%s is empty ;). Continuing'%(args.caltech_pedestrian_root)
single_file_predictions = []
id = lines[0].split()[0]
f = open_file(args.caltech_pedestrian_root, id)
for line in lines:
if line.split()[0] !=id:
save_file_predictions(args.caltech_pedestrian_root,single_file_predictions,f)
#start new file
id = line.split(' ')[0]
f = open_file(args.caltech_pedestrian_root, id) # if already created, then do nothing
single_file_predictions = []
single_file_predictions.append(line)
if __name__=="__main__":
main(sys.argv)