Advertisement
Guest User

Untitled

a guest
Jun 19th, 2019
75
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 11.79 KB | None | 0 0
  1. #-------------------------------------------------------------------------------
  2. # Name: Object bounding box label tool
  3. # Purpose: Label object bboxes for ImageNet Detection data
  4.  
  5. from __future__ import division
  6. from Tkinter import *
  7. import tkMessageBox
  8. from PIL import Image, ImageTk
  9. import os
  10. import glob
  11. import random
  12.  
  13. # colors for the bboxes
  14. COLORS = ['red', 'blue', 'yellow', 'pink', 'cyan', 'green', 'black', 'purple', 'orange', 'aqua']
  15. # image sizes for the examples
  16. SIZE = 256, 256
  17.  
  18. class LabelTool():
  19. def __init__(self, master):
  20. # set up the main frame
  21. self.parent = master
  22. self.parent.title("LabelTool")
  23. self.frame = Frame(self.parent)
  24. self.frame.pack(fill=BOTH, expand=1)
  25. self.parent.resizable(width = FALSE, height = FALSE)
  26.  
  27. # initialize global state
  28. self.imageDir = ''
  29. self.imageList= []
  30. self.egDir = ''
  31. self.egList = []
  32. self.outDir = ''
  33. self.cur = 0
  34. self.total = 0
  35. self.category = 0
  36. self.imagename = ''
  37. self.labelfilename = ''
  38. self.tkimg = None
  39.  
  40. # initialize mouse state
  41. self.STATE = {}
  42. self.STATE['click'] = 0
  43. self.STATE['x'], self.STATE['y'] = 0, 0
  44.  
  45. # reference to bbox
  46. self.bboxIdList = []
  47. self.bboxId = None
  48. self.bboxList = []
  49. self.hl = None
  50. self.vl = None
  51.  
  52. # ----------------- GUI stuff ---------------------
  53. # dir entry & load
  54. self.label = Label(self.frame, text = "Image Dir:")
  55. self.label.grid(row = 0, column = 0, sticky = E)
  56. self.entry = Entry(self.frame)
  57. self.entry.grid(row = 0, column = 1, sticky = W+E)
  58. self.ldBtn = Button(self.frame, text = "Load", command = self.loadDir)
  59. self.ldBtn.grid(row = 0, column = 2, sticky = W+E)
  60.  
  61. # main panel for labeling
  62. self.mainPanel = Canvas(self.frame, cursor='tcross')
  63. self.mainPanel.bind("<Button-1>", self.mouseClick)
  64. self.mainPanel.bind("<Motion>", self.mouseMove)
  65. self.parent.bind("<Escape>", self.cancelBBox) # press <Espace> to cancel current bbox
  66. self.parent.bind("s", self.cancelBBox)
  67. self.parent.bind("a", self.prevImage) # press 'a' to go backforward
  68. self.parent.bind("d", self.nextImage) # press 'd' to go forward
  69. self.mainPanel.grid(row = 1, column = 1, rowspan = 4, sticky = W+N)
  70.  
  71. # showing bbox info & delete bbox
  72. self.lb1 = Label(self.frame, text = 'Bounding boxes:')
  73. self.lb1.grid(row = 1, column = 2, sticky = W+N)
  74. self.listbox = Listbox(self.frame, width = 22, height = 12)
  75. self.listbox.grid(row = 2, column = 2, sticky = N)
  76. self.btnDel = Button(self.frame, text = 'Delete', command = self.delBBox)
  77. self.btnDel.grid(row = 3, column = 2, sticky = W+E+N)
  78. self.btnClear = Button(self.frame, text = 'ClearAll', command = self.clearBBox)
  79. self.btnClear.grid(row = 4, column = 2, sticky = W+E+N)
  80.  
  81. # control panel for image navigation
  82. self.ctrPanel = Frame(self.frame)
  83. self.ctrPanel.grid(row = 5, column = 1, columnspan = 2, sticky = W+E)
  84. self.prevBtn = Button(self.ctrPanel, text='<< Prev', width = 10, command = self.prevImage)
  85. self.prevBtn.pack(side = LEFT, padx = 5, pady = 3)
  86. self.nextBtn = Button(self.ctrPanel, text='Next >>', width = 10, command = self.nextImage)
  87. self.nextBtn.pack(side = LEFT, padx = 5, pady = 3)
  88. self.progLabel = Label(self.ctrPanel, text = "Progress: / ")
  89. self.progLabel.pack(side = LEFT, padx = 5)
  90. self.tmpLabel = Label(self.ctrPanel, text = "Go to Image No.")
  91. self.tmpLabel.pack(side = LEFT, padx = 5)
  92. self.idxEntry = Entry(self.ctrPanel, width = 5)
  93. self.idxEntry.pack(side = LEFT)
  94. self.goBtn = Button(self.ctrPanel, text = 'Go', command = self.gotoImage)
  95. self.goBtn.pack(side = LEFT)
  96.  
  97. # example pannel for illustration
  98. self.egPanel = Frame(self.frame, border = 10)
  99. self.egPanel.grid(row = 1, column = 0, rowspan = 5, sticky = N)
  100. self.tmpLabel2 = Label(self.egPanel, text = "Examples:")
  101. self.tmpLabel2.pack(side = TOP, pady = 5)
  102. self.egLabels = []
  103. for i in range(2):
  104. self.egLabels.append(Label(self.egPanel))
  105. self.egLabels[-1].pack(side = TOP)
  106.  
  107. # display mouse position
  108. self.disp = Label(self.ctrPanel, text='')
  109. self.disp.pack(side = RIGHT)
  110.  
  111. self.frame.columnconfigure(1, weight = 1)
  112. self.frame.rowconfigure(4, weight = 1)
  113.  
  114.  
  115. def loadDir(self, dbg = False):
  116. if not dbg:
  117. s = self.entry.get()
  118. self.parent.focus()
  119. self.category = int(s)
  120.  
  121. else:
  122. s = r'C:/Users/Werda Farooq/Desktop/box-label-tool-master'
  123.  
  124. print ('%d' %(int(s)))
  125. self.imageDir = os.path.join(r'./Images', '%03d' %(self.category))
  126. self.imageList = glob.glob(os.path.join(self.imageDir, '*.jpeg'))
  127.  
  128. print '%s t1'%(self.imageDir)
  129. if len(self.imageList) == 0:
  130. print 'No .jpg images found in the specified dir!'
  131. return
  132.  
  133. # default to the 1st image in the collection
  134. self.cur = 1
  135. self.total = len(self.imageList)
  136. print '%s t2'%(self.imageList)
  137. # set up output dir
  138. self.outDir = os.path.join(r'./Labels', '%03d' %(self.category))
  139. if not os.path.exists(self.outDir):
  140. os.mkdir(self.outDir)
  141. print'%s t3' %(self.outDir)
  142. # load example bboxes
  143. self.egDir = os.path.join(r'./Examples', '%03d' %(self.category))
  144. if not os.path.exists(self.egDir):
  145. return
  146. filelist = glob.glob(os.path.join(self.egDir, '*.png'))
  147. self.tmp = []
  148. self.egList = []
  149. random.shuffle(filelist)
  150. for (i, f) in enumerate(filelist):
  151. if i == 3:
  152. break
  153. im = Image.open(f)
  154. r = min(SIZE[0] / im.size[0], SIZE[1] / im.size[1])
  155. new_size = int(r * im.size[0]), int(r * im.size[1])
  156. self.tmp.append(im.resize(new_size, Image.ANTIALIAS))
  157. self.egList.append(ImageTk.PhotoImage(self.tmp[-1]))
  158. self.egLabels[i].config(image = self.egList[-1], width = SIZE[0], height = SIZE[1])
  159.  
  160. self.loadImage()
  161. print '%d images loaded from %s' %(self.total, s)
  162.  
  163. def loadImage(self):
  164. # load image
  165. imagepath = self.imageList[self.cur - 1]
  166. self.img = Image.open(imagepath)
  167. self.tkimg = ImageTk.PhotoImage(self.img)
  168. self.mainPanel.config(width = max(self.tkimg.width(), 400), height = max(self.tkimg.height(), 400))
  169. self.mainPanel.create_image(0, 0, image = self.tkimg, anchor=NW)
  170. self.progLabel.config(text = "%04d/%04d" %(self.cur, self.total))
  171.  
  172. # load labels
  173. self.clearBBox()
  174. self.imagename = os.path.split(imagepath)[-1].split('.')[0]
  175. labelname = self.imagename + '.txt'
  176. self.labelfilename = os.path.join(self.outDir, labelname)
  177. bbox_cnt = 0
  178. if os.path.exists(self.labelfilename):
  179. with open(self.labelfilename) as f:
  180. for (i, line) in enumerate(f):
  181. if i == 0:
  182. bbox_cnt = int(line.strip())
  183. continue
  184. tmp = [int(t.strip()) for t in line.split()]
  185.  
  186. self.bboxList.append(tuple(tmp))
  187. tmpId = self.mainPanel.create_rectangle(tmp[0], tmp[1],
  188. tmp[2], tmp[3],
  189. width = 2,
  190. outline = COLORS[(len(self.bboxList)-1) % len(COLORS)])
  191. self.bboxIdList.append(tmpId)
  192. self.listbox.insert(END, '(%d, %d) -> (%d, %d)' %(tmp[0], tmp[1], tmp[2], tmp[3]))
  193. self.listbox.itemconfig(len(self.bboxIdList) - 1, fg = COLORS[(len(self.bboxIdList) - 1) % len(COLORS)])
  194.  
  195. def saveImage(self):
  196. with open(self.labelfilename, 'w') as f:
  197. f.write('%dn' %len(self.bboxList))
  198. for bbox in self.bboxList:
  199. f.write(' '.join(map(str, bbox)) + 'n')
  200. print 'Image No. %d saved' %(self.cur)
  201.  
  202.  
  203. def mouseClick(self, event):
  204. if self.STATE['click'] == 0:
  205. self.STATE['x'], self.STATE['y'] = event.x, event.y
  206. else:
  207. x1, x2 = min(self.STATE['x'], event.x), max(self.STATE['x'], event.x)
  208. y1, y2 = min(self.STATE['y'], event.y), max(self.STATE['y'], event.y)
  209. self.bboxList.append((x1, y1, x2, y2))
  210. self.bboxIdList.append(self.bboxId)
  211. self.bboxId = None
  212. self.listbox.insert(END, '(%d, %d) -> (%d, %d)' %(x1, y1, x2, y2))
  213. self.listbox.itemconfig(len(self.bboxIdList) - 1, fg = COLORS[(len(self.bboxIdList) - 1) % len(COLORS)])
  214. self.STATE['click'] = 1 - self.STATE['click']
  215.  
  216. def mouseMove(self, event):
  217. self.disp.config(text = 'x: %d, y: %d' %(event.x, event.y))
  218. if self.tkimg:
  219. if self.hl:
  220. self.mainPanel.delete(self.hl)
  221. self.hl = self.mainPanel.create_line(0, event.y, self.tkimg.width(), event.y, width = 2)
  222. if self.vl:
  223. self.mainPanel.delete(self.vl)
  224. self.vl = self.mainPanel.create_line(event.x, 0, event.x, self.tkimg.height(), width = 2)
  225. if 1 == self.STATE['click']:
  226. if self.bboxId:
  227. self.mainPanel.delete(self.bboxId)
  228. self.bboxId = self.mainPanel.create_rectangle(self.STATE['x'], self.STATE['y'],
  229. event.x, event.y,
  230. width = 2,
  231. outline = COLORS[len(self.bboxList) % len(COLORS)])
  232.  
  233. def cancelBBox(self, event):
  234. if 1 == self.STATE['click']:
  235. if self.bboxId:
  236. self.mainPanel.delete(self.bboxId)
  237. self.bboxId = None
  238. self.STATE['click'] = 0
  239.  
  240. def delBBox(self):
  241. sel = self.listbox.curselection()
  242. if len(sel) != 1 :
  243. return
  244. idx = int(sel[0])
  245. self.mainPanel.delete(self.bboxIdList[idx])
  246. self.bboxIdList.pop(idx)
  247. self.bboxList.pop(idx)
  248. self.listbox.delete(idx)
  249.  
  250. def clearBBox(self):
  251. for idx in range(len(self.bboxIdList)):
  252. self.mainPanel.delete(self.bboxIdList[idx])
  253. self.listbox.delete(0, len(self.bboxList))
  254. self.bboxIdList = []
  255. self.bboxList = []
  256.  
  257. def prevImage(self, event = None):
  258. self.saveImage()
  259. if self.cur > 1:
  260. self.cur -= 1
  261. self.loadImage()
  262.  
  263. def nextImage(self, event = None):
  264. self.saveImage()
  265. if self.cur < self.total:
  266. self.cur += 1
  267. self.loadImage()
  268.  
  269. def gotoImage(self):
  270. idx = int(self.idxEntry.get())
  271. if 1 <= idx and idx <= self.total:
  272. self.saveImage()
  273. self.cur = idx
  274. self.loadImage()
  275.  
  276.  
  277. if __name__ == '__main__':
  278. root = Tk()
  279. tool = LabelTool(root)
  280. root.resizable(width = True, height = True)
  281. root.mainloop()
  282.  
  283.  
  284. i expect the output as:
  285.  
  286. rawdata/25.bmp 12 710 363 760 413 741 393 789 441 543 398 617 472 1083 529 1191 637 658 443 729 514 10 606 104 700 754 502 906 654 266 278 308 320 467 318 510 361 693 337 743 387 1035 333 1076 374 26 570 74 618
  287.  
  288. rawdata/50.bmp 12 710 363 760 413 740 393 788 441 543 400 618 475 1083 529 1191 637 10 606 104 700 754 502 906 654 661 447 730 516 267 273 310 316 466 319 509 362 693 337 743 387 1035 333 1076 374 27 570 75 618
  289.  
  290. rawdata/75.bmp 12 710 363 760 413 741 393 789 441 543 401 620 478 1083 529 1191 637 10 606 104 700 751 500 905 654 661 450 731 520 506 275 546 315 465 320 509 364 693 337 743 387 1035 332 1077 374 27 569 75 617
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement