-
Notifications
You must be signed in to change notification settings - Fork 20
/
euclid_yolo_kitti_converter.py
304 lines (248 loc) · 13.3 KB
/
euclid_yolo_kitti_converter.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
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
#-------------------------------------------------------------------------------
# Euclid - Labelling tool
# Create and label bounding boxes
# prabindh@yahoo.com, 2016
# Initial code taken from github.com/puzzledqs/BBox-Label-Tool
# Significantly modified to add more image types, image folders, labelling saves, and format, and format selection
# Currently supports 8 classes, and Kitti and YOLO(darknet) output formats
# Python 2.7
# pip install pillow
# pip install image
#
#-------------------------------------------------------------------------------
#-------------------------------------------------------------------------------
# The DetectNet/ Kitti Database format
# Taken from https://github.com/NVIDIA/DIGITS/blob/master/digits/extensions/data/objectDetection/README.md#label-format
# All values (numerical or strings) are separated via spaces,
# each row corresponds to one object. The 15 columns represent:
#
#Values Name Description
#----------------------------------------------------------------------------
# 1 type Describes the type of object: 'Car', 'Van', 'Truck',
# 'Pedestrian', 'Person_sitting', 'Cyclist', 'Tram',
# 'Misc' or 'DontCare'
# 1 truncated Float from 0 (non-truncated) to 1 (truncated), where
# truncated refers to the object leaving image boundaries
# 1 occluded Integer (0,1,2,3) indicating occlusion state:
# 0 = fully visible, 1 = partly occluded
# 2 = largely occluded, 3 = unknown
# 1 alpha Observation angle of object, ranging [-pi..pi]
# 4 bbox 2D bounding box of object in the image (0-based index):
# contains left, top, right, bottom pixel coordinates
# 3 dimensions 3D object dimensions: height, width, length (in meters)
# 3 location 3D object location x,y,z in camera coordinates (in meters)
# 1 rotation_y Rotation ry around Y-axis in camera coordinates [-pi..pi]
# 1 score Only for results: Float, indicating confidence in
# detection, needed for p/r curves, higher is better.
#-------------------------------------------------------------------------------
#-------------------------------------------------------------------------------
# The YOLO format
# All values (numerical or strings) are separated via spaces,
# each row corresponds to one object. The 5 columns represent:
#
#Values Name Description
#----------------------------------------------------------------------------
# 1 Class ID Describes the class number of object, as an integer number (0 based)
# 1 Center_X Float from 0 to 1, X coordinate of b-box center, normalised to image width
# 1 Center_Y Float from 0 to 1, Y coordinate of b-box center, normalised to image height
# 1 Bbox_Width Float from 0 to 1, Width of b-box, normalised to image width
# 1 Bbox_Height Float from 0 to 1, Height of b-box, normalised to image height
#-------------------------------------------------------------------------------
import sys
if sys.version_info[0] < 3:
from Tkinter import *
import tkMessageBox
import tkFileDialog
else:
from tkinter import *
import tkinter.messagebox as tkMessageBox
import tkinter.filedialog as tkFileDialog
from PIL import Image, ImageTk
import os
import glob
import random
# Usage
USAGE = " \
1. Select a Directory of labels, or Enter the path directly and click Load\n \
2. Click Convert. The labels will be converted and stored in current directory \n \
"
# Object Classes (No spaces in name)
CLASSES = ['Class0', 'Class1', 'Class2', 'Class3', 'Class4', 'Class5', 'Class6', 'Class7']
class EuclidConverter():
def askDirectory(self):
self.imageDir = tkFileDialog.askdirectory()
self.entry.insert(0, self.imageDir)
self.loadDir(self)
def loadDir(self, dbg = False):
self.imageDir = self.entry.get()
self.parent.focus()
if not os.path.isdir(self.imageDir):
tkMessageBox.showerror("Folder error", message = "The specified directory doesn't exist!")
return
#get label list
labelFileTypes = ('*.txt') # the tuple of file types
self.labelFileList = []
# load labels
for files in labelFileTypes:
self.labelFileList.extend(glob.glob(os.path.join(self.imageDir, files.lower())) )
if (False == self.is_windows):
self.labelFileList.extend(glob.glob(os.path.join(self.imageDir, files)) )
if len(self.labelFileList) == 0:
tkMessageBox.showerror("Label files not found", message = "No labels (.txt) found in folder!")
self.updateStatus( 'No label files found in the specified dir!')
return
# Change title
self.parent.title("Euclid Label Converter (" + self.imageDir + ") " + str(len(self.labelFileList)) + " label files")
def showHelp(self, event):
tkMessageBox.showinfo("Help", USAGE)
def __init__(self, master):
# set up the main frame
self.parent = master
self.parent.title("Euclid Labeller (Press F1 for Help)")
self.frame = Frame(self.parent)
self.frame.pack(fill=BOTH, expand=1)
self.parent.resizable(width = TRUE, height = TRUE)
self.is_windows = hasattr(sys, 'getwindowsversion')
# initialize global state
self.imageDir = ''
self.imageList= []
self.outDir = ''
self.cur = 0
self.total = 0
self.imagename = ''
self.labelfilename = ''
self.currLabelMode = 'YOLO' #'KITTI' #'YOLO' # Other modes TODO
self.imagefilename = ''
self.tkimg = None
# initialize mouse state
self.STATE = {}
self.STATE['click'] = 0
self.STATE['x'], self.STATE['y'] = 0, 0
self.currentMouseX = 0;
self.currentMouseY = 0;
#colors
self.redColor = self.blueColor = self.greenColor = 128
# reference to bbox
self.bboxIdList = []
self.bboxId = None
self.bboxList = []
self.currClassLabel = 0
self.classLabelList = []
self.hl = None
self.vl = None
self.parent.bind("<F1>", self.showHelp) # press <F1> to show help
# ----------------- GUI stuff ---------------------
# dir entry & load File control panel
self.FileControlPanelFrame = Frame(self.frame)
self.FileControlPanelFrame.grid(row = 0, column = 0, sticky = W)
self.FileControlPanelLabel = Label(self.FileControlPanelFrame, text = '1. Select a directory (or) Enter input path, and click Load')
self.FileControlPanelLabel.grid(row = 0, column = 0, sticky = W+N)
self.browserBtn = Button(self.FileControlPanelFrame, text = "Select Dir", command = self.askDirectory)
self.browserBtn.grid(row = 1, column = 0, sticky = N)
self.entry = Entry(self.FileControlPanelFrame)
self.entry.grid(row = 1, column = 1, sticky = N)
self.ldBtn = Button(self.FileControlPanelFrame, text = "Load", command = self.loadDir)
self.ldBtn.grid(row = 1, column = 2, sticky = N)
self.ConvertBtn = Button(self.FileControlPanelFrame, text = "Convert", command = self.ConvertLabels)
self.ConvertBtn.grid(row = 2, column = 0, sticky = N)
# control panel for image navigation
self.ctrPanel = Frame(self.frame)
self.ctrPanel.grid(row = 1, column = 0, columnspan = 2, sticky = W+N)
self.progLabel = Label(self.ctrPanel, text = "Progress: [ 0 / 0 ]")
self.progLabel.pack(side = LEFT, padx = 5)
# Status panel for image navigation
self.statusPanel = Frame(self.frame)
self.statusPanel.grid(row = 2, column = 0, columnspan = 3, sticky = W)
self.statusText = StringVar()
self.statusLabel = Label(self.statusPanel, textvariable = self.statusText)
self.statusLabel.grid(row = 0, column = 0, sticky = W+E+N)
self.updateStatus("Directory not selected.")
# display mouse position
self.disp = Label(self.ctrPanel, text='')
self.disp.pack(side = RIGHT)
self.frame.columnconfigure(1, weight = 1)
self.frame.rowconfigure(4, weight = 1)
def GetBoundariesFromYoloFile(self, centerX, centerY, width, height, imageWidth, imageHeight):
topLeftX = (int)(centerX*imageWidth - (width*imageWidth)/2)
topLeftY = (int)(centerY*imageHeight - (height*imageHeight)/2)
bottomRightX = (int)(centerX*imageWidth + (width*imageWidth)/2)
bottomRightY = (int)(centerY*imageHeight + (height*imageHeight)/2)
return topLeftX, topLeftY, bottomRightX, bottomRightY
def convert2Yolo(self, image, boxCoords):
invWidth = 1./image[0]
invHeight = 1./image[1]
x = invWidth * (boxCoords[0] + boxCoords[2])/2.0
y = invHeight * (boxCoords[1] + boxCoords[3])/2.0
boxWidth = invWidth * (boxCoords[2] - boxCoords[0])
boxHeight = invHeight * (boxCoords[3] - boxCoords[1])
return (x,y,boxWidth,boxHeight)
def updateStatus(self, newStatus):
self.statusText.set("Status: " + newStatus)
def YoloLabelWriteOut(self, filename, classname, topLeftX, topLeftY, bottomRightX, bottomRightY):
labelFile = open(filename, "a+")
#TODO - get image w/h
yoloOut = self.convert2Yolo([1000,1000],
[topLeftX, topLeftY, bottomRightX, bottomRightY] );
labelFile.write('%s' % CLASSES[classname])
labelFile.write(' %.7f %.7f %.7f %.7f\n' % (yoloOut[0], yoloOut[1], yoloOut[2], yoloOut[3]))
labelFile.close()
def KittiLabelWriteOut(self, filename, classname, topLeftX, topLeftY, bottomRightX, bottomRightY):
labelFile = open(filename+ "Kitti.txt", "a+")
#TODO - get image w/h
#print(classname,playerNameArray[0])
labelFile.write('%s' % CLASSES[classname])
labelFile.write(' 0.0 0 0.0 ')
labelFile.write('%.2f %.2f %.2f %.2f' % (topLeftX, topLeftY, bottomRightX, bottomRightY))
labelFile.write(' 0.0 0.0 0.0 0.0 0.0 0.0 0.0')
labelFile.write('\n')
labelFile.close()
# Convert label formats
def ConvertLabels(self):
#get image list
labelFileTypes = ('*.txt') # the tuple of file types
self.labelFileList = []
# load labels
for files in labelFileTypes:
self.labelFileList.extend(glob.glob(os.path.join(self.imageDir, files.lower())) )
if (False == self.is_windows):
self.labelFileList.extend(glob.glob(os.path.join(self.imageDir, files)) )
if len(self.labelFileList) == 0:
tkMessageBox.showerror("Label files not found", message = "No labels (.txt) found in folder!")
self.updateStatus( 'No label files found in the specified dir!')
return
# Change title
self.parent.title("Euclid Label Converter (" + self.imageDir + ") " + str(len(self.labelFileList)) + " label files")
# default to the 1st image in the collection
self.cur = 1
self.total = len(self.labelFileList)
# set up output dir
self.outDir = os.path.join(self.imageDir + '/ConvertedLabelData')
if not os.path.exists(self.outDir):
os.mkdir(self.outDir)
self.updateStatus( '%d label files located from %s' %(self.total, self.imageDir))
# Convert one file
#self.labelfilename = self.labelFileList[self.cur - 1]
for file in range(self.total-1):
self.labelfilename = self.labelFileList[file]
#print(self.labelfilename, file)
lastPartFileName, lastPartFileExtension = os.path.splitext(os.path.split(self.labelfilename)[-1])
newFileName = self.outDir + '/' + lastPartFileName + lastPartFileExtension
bbox_cnt = 0
if os.path.exists(self.labelfilename):
with open(self.labelfilename) as f:
for (i, line) in enumerate(f):
bbox_cnt = len(line)
tmp = [elements.strip() for elements in line.split()]
# for each line, convert and writeout
if(len(tmp) > 5):
self.currLabelMode='KITTI'
self.YoloLabelWriteOut(newFileName, int(tmp[0]), float(tmp[4]), float(tmp[5]), float(tmp[6]), float(tmp[7]) )
elif(len(tmp) == 5):
self.currLabelMode='YOLO'
bbTuple = self.GetBoundariesFromYoloFile(float(tmp[1]),float(tmp[2]), float(tmp[3]),float(tmp[4]),
self.tkimg.width(), self.tkimg.height() )
self.KittiLabelWriteOut(self.labelfilename, int(tmp[0]), bbTuple[0], bbTuple[1], bbTuple[2], bbTuple[3] )
if __name__ == '__main__':
root = Tk()
tool = EuclidConverter(root)
root.mainloop()