-
Notifications
You must be signed in to change notification settings - Fork 0
/
convert_yolodarknetdata.py
74 lines (60 loc) · 1.82 KB
/
convert_yolodarknetdata.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
#data looks like:
#00110592,car,465,111,507,135
#00110592,non-motorized_vehicle,197,187,318,269
#Now we need to generate the label files that Darknet uses. Darknet wants a .txt file for each image with a line for each ground truth object in the image that looks like:
#make file per picture (0011...txt) with:
#<object-class> <x> <y> <width> <height>
import csv
#DIT HEEFT LENNART AL GEMAAKT, IK MOET IETS ANDERS DOEN, VRAGEN
import pickle
def classnumber(veh_class):
#print(veh_class)
dictionary = {
"articulated_truck": 1,
"bicycle": 2,
"bus": 3,
"car": 4,
"motorcycle": 5,
"motorized_vehicle": 6,
"non-motorized_vehicle": 7,
"pedestrian": 8,
"pickup_truck": 9,
"single_unit_truck": 10,
"work_van": 11
}
return str(dictionary[veh_class])
def makedict():
file = open("data/MIO-TCD-Localization/gt_train.txt","r",encoding="UTF-16")
filedictionary = {}
for line in file.readlines():
params = line.rstrip("\n").split(",")
picture = params[0]
type = params[1]
x1 = eval(params[2])
y1 = eval(params[3])
x2 = eval(params[4])
y2 = eval(params[5])
width = x2-x1
height = y2-y1
xmid = round((x1 + x2)/2 + 1)
ymid = round((y1 + y2)/2 + 1)
newline = classnumber(type) + " " + str(xmid) + " " + str(ymid) + " " + str(width) + " " + str(height) + "\n"
if picture not in filedictionary:
filedictionary[picture] = [newline]
else:
filedictionary[picture].append(newline)
pickle.dump(filedictionary, open("filedictionary.pickle",'wb'))
print("Dict made")
def main():
#makedict()
filedictionary = pickle.load(open("filedictionary.pickle","rb"))
for picture in filedictionary:
filename = picture+".txt"
file = open("yolodarknetdata/"+filename, "w")
for item in filedictionary[picture]:
if item != filedictionary[picture][-1]:
file.write(item+"\n")
else:
file.write(item)
file.close()
main()