diff --git a/.gitignore b/.gitignore index c5c32eb8..d0618fbe 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ /Images !/Images/001 /Labels -!Labels/001 \ No newline at end of file +!Labels/001 +.DS_Store diff --git a/class.txt b/class.txt index a931f210..bdc8403c 100644 --- a/class.txt +++ b/class.txt @@ -1,3 +1,2 @@ -dog -cat -baby \ No newline at end of file +person +helmet \ No newline at end of file diff --git a/convert.py b/convert.py new file mode 100644 index 00000000..0ccd2fe8 --- /dev/null +++ b/convert.py @@ -0,0 +1,110 @@ +# -*- coding: utf-8 -*- +""" +Created on Wed Dec 9 14:55:43 2015 + +This script is to convert the txt annotation files to appropriate format needed by YOLO + +@author: Guanghan Ning +Email: gnxr9@mail.missouri.edu +""" + +import os +from os import walk, getcwd +from PIL import Image + +classes = ["helmet"] + +def convert(size, box): + dw = 1./size[0] + dh = 1./size[1] + x = (box[0] + box[1])/2.0 + y = (box[2] + box[3])/2.0 + w = box[1] - box[0] + h = box[3] - box[2] + x = x*dw + w = w*dw + y = y*dh + h = h*dh + return (x,y,w,h) + + +"""-------------------------------------------------------------------""" + +""" Configure Paths""" +cls = "helmet" + +#Raj: Changes in code to run it directory specific! working directory needs "images" folder where all images are there! +workingDirectory = "multiple-result" +mypath = "/Users/i305138/SAP/SKOM/BBox-Label-Tool/" + workingDirectory + "/labels/" + cls + "/" +outpath = "/Users/i305138/SAP/SKOM/BBox-Label-Tool/" + workingDirectory + "/labels-yolo/" + cls + "/" + +if cls not in classes: + exit(0) +cls_id = classes.index(cls) + +print ("Found class, class ID: " + str(cls_id)) + +wd = getcwd() +list_file = open('%s/%s/%s_list.txt'%(wd, workingDirectory, cls), 'w') + +""" Get input text file list """ +print ("My Path: " + mypath) +txt_name_list = [] +for (dirpath, dirnames, filenames) in walk(mypath): + txt_name_list.extend(filenames) + break +print(txt_name_list) +if '.DS_Store' in txt_name_list: txt_name_list.remove('.DS_Store') +print(txt_name_list) +""" Process """ +for txt_name in txt_name_list: + # txt_file = open("Labels/stop_sign/001.txt", "r") + + """ Open input text files """ + txt_path = mypath + txt_name + print("Input:" + txt_path) + txt_file = open(txt_path, "r") + lines = txt_file.read().split('\n') #for ubuntu, use "\r\n" instead of "\n" + print ("Source Lines: "+ str (lines)) + """ Open output text files """ + txt_outpath = outpath + txt_name + print("Output:" + txt_outpath) + txt_outfile = open(txt_outpath, "w") + + + """ Convert the data to YOLO format """ + ct = 0 + for line in lines: + #print('lenth of line is: ') + #print(len(line)) + #print('\n') + if(len(line) >= 2): + ct = ct + 1 + print(line + "\n") + elems = line.split(' ') + print(elems) + xmin = elems[0] + xmax = elems[2] + ymin = elems[1] + ymax = elems[3] + # + img_path = str('%s/%s/images/%s/%s.JPEG'%(wd, workingDirectory, cls, os.path.splitext(txt_name)[0])) + #t = magic.from_file(img_path) + #wh= re.search('(\d+) x (\d+)', t).groups() + im=Image.open(img_path) + w= int(im.size[0]) + h= int(im.size[1]) + #w = int(xmax) - int(xmin) + #h = int(ymax) - int(ymin) + # print(xmin) + print(w, h) + b = (float(xmin), float(xmax), float(ymin), float(ymax)) + bb = convert((w,h), b) + print(bb) + txt_outfile.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n') + + """ Save those images with bb into list""" + if(ct != 0): + list_file.write('%s/%s/images/%s/%s.JPEG\n'%(wd, workingDirectory, cls, os.path.splitext(txt_name)[0])) + +list_file.close() diff --git a/main.py b/main.py index df680920..00dbcb63 100644 --- a/main.py +++ b/main.py @@ -16,7 +16,7 @@ import random # colors for the bboxes -COLORS = ['red', 'blue', 'olive', 'teal', 'cyan', 'green', 'black'] +COLORS = ['red', 'blue', 'green', 'black'] # image sizes for the examples SIZE = 256, 256 @@ -40,6 +40,7 @@ def __init__(self, master): self.category = 0 self.imagename = '' self.labelfilename = '' + self.yoloLabelFileName = '' self.tkimg = None self.currentLabelclass = '' self.cla_can_temp = [] @@ -195,6 +196,7 @@ def loadDir(self, dbg = False): def loadImage(self): # load image imagepath = self.imageList[self.cur - 1] + print ('Image loaded: ' + imagepath) self.img = Image.open(imagepath) self.tkimg = ImageTk.PhotoImage(self.img) self.mainPanel.config(width = max(self.tkimg.width(), 400), height = max(self.tkimg.height(), 400)) @@ -206,6 +208,8 @@ def loadImage(self): self.imagename = os.path.split(imagepath)[-1].split('.')[0] labelname = self.imagename + '.txt' self.labelfilename = os.path.join(self.outDir, labelname) + yoloLabelName = 'yolo/' + self.imagename + '.txt' + self.yoloLabelFileName = os.path.join(self.outDir, yoloLabelName) bbox_cnt = 0 if os.path.exists(self.labelfilename): with open(self.labelfilename) as f: @@ -232,9 +236,29 @@ def saveImage(self): f.write('%d\n' %len(self.bboxList)) for bbox in self.bboxList: f.write(' '.join(map(str, bbox)) + '\n') + with open(self.yoloLabelFileName, 'w') as fy: + #Raj: Code to write the data in Yolo format: + w= int(self.img.size[0]) + h= int(self.img.size[1]) + sizeVar = (w,h) + for bbox in self.bboxList: + x1 = bbox[0] + x2 = bbox[2] + y1 = bbox[1] + y2 = bbox[3] + className = bbox[4] + b = (float(x1), float(x2), float(y1), float(y2)) + bb = self.convert(sizeVar, b) + #Raj: We will add the index of the current class since Yolo training will get this index and match it up in the *.names file for the proper class name + currentClassIndex = self.classcandidate['values'].index(className) + yoloList = [] + yoloList.append(str(currentClassIndex) + " " + " ".join([str(a) for a in bb])) + + for yoloLine in yoloList: + print ('Yolo line: ' + str(yoloLine)) + fy.write(str(yoloLine)+ '\n') print 'Image No. %d saved' %(self.cur) - def mouseClick(self, event): if self.STATE['click'] == 0: self.STATE['x'], self.STATE['y'] = event.x, event.y @@ -246,8 +270,22 @@ def mouseClick(self, event): self.bboxId = None self.listbox.insert(END, '%s : (%d, %d) -> (%d, %d)' %(self.currentLabelclass,x1, y1, x2, y2)) self.listbox.itemconfig(len(self.bboxIdList) - 1, fg = COLORS[(len(self.bboxIdList) - 1) % len(COLORS)]) + self.STATE['click'] = 1 - self.STATE['click'] + def convert(self, sizetuple, box): + dw = 1./sizetuple[0] + dh = 1./sizetuple[1] + x = (box[0] + box[1])/2.0 + y = (box[2] + box[3])/2.0 + w = box[1] - box[0] + h = box[3] - box[2] + x = x*dw + w = w*dw + y = y*dh + h = h*dh + return (x,y,w,h) + def mouseMove(self, event): self.disp.config(text = 'x: %d, y: %d' %(event.x, event.y)) if self.tkimg: