From 2783071c85e6bac2b8a696fa1037f4229fb8b60f Mon Sep 17 00:00:00 2001 From: Wang Yinghao Date: Thu, 1 Mar 2018 21:38:22 -0600 Subject: [PATCH] YOLO Write --- labelImg.py | 32 ++++++++--- libs/labelFile.py | 28 ++++++++++ libs/lib.py | 2 +- libs/yolo_io.py | 133 ++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 188 insertions(+), 7 deletions(-) create mode 100644 libs/yolo_io.py diff --git a/labelImg.py b/labelImg.py index a817a653..973c7686 100755 --- a/labelImg.py +++ b/labelImg.py @@ -38,6 +38,7 @@ from libs.labelFile import LabelFile, LabelFileError from libs.toolBar import ToolBar from libs.pascal_voc_io import PascalVocReader from libs.pascal_voc_io import XML_EXT +from libs.yolo_io import TXT_EXT from libs.ustr import ustr from libs.version import __version__ @@ -97,6 +98,8 @@ class MainWindow(QMainWindow, WindowMixin): # Save as Pascal voc xml self.defaultSaveDir = None self.usingPascalVocFormat = True + self.usingYoloFormat = False + # For loading all image under a directory self.mImgList = [] self.dirname = None @@ -232,6 +235,9 @@ class MainWindow(QMainWindow, WindowMixin): save = action('&Save', self.saveFile, 'Ctrl+S', 'save', u'Save labels to file', enabled=False) + save_format = action('&Format', self.change_format, + 'Ctrl+', 'format', u'Change save format', enabled=True) + saveAs = action('&Save As', self.saveFileAs, 'Ctrl+Shift+S', 'save-as', u'Save labels to a different file', enabled=False) @@ -324,7 +330,7 @@ class MainWindow(QMainWindow, WindowMixin): self.popLabelListMenu) # Store actions for further handling. - self.actions = struct(save=save, saveAs=saveAs, open=open, close=close, resetAll = resetAll, + self.actions = struct(save=save, save_format=save_format, saveAs=saveAs, open=open, close=close, resetAll = resetAll, lineColor=color1, create=create, delete=delete, edit=edit, copy=copy, createMode=createMode, editMode=editMode, advancedMode=advancedMode, shapeLineColor=shapeLineColor, shapeFillColor=shapeFillColor, @@ -363,7 +369,7 @@ class MainWindow(QMainWindow, WindowMixin): self.lastLabel = None addActions(self.menus.file, - (open, opendir, changeSavedir, openAnnotation, self.menus.recentFiles, save, saveAs, close, resetAll, quit)) + (open, opendir, changeSavedir, openAnnotation, self.menus.recentFiles, save, save_format, saveAs, close, resetAll, quit)) addActions(self.menus.help, (help, showInfo)) addActions(self.menus.view, ( self.autoSaving, @@ -383,11 +389,11 @@ class MainWindow(QMainWindow, WindowMixin): self.tools = self.toolbar('Tools') self.actions.beginner = ( - open, opendir, changeSavedir, openNextImg, openPrevImg, verify, save, None, create, copy, delete, None, + open, opendir, changeSavedir, openNextImg, openPrevImg, verify, save, save_format, None, create, copy, delete, None, zoomIn, zoom, zoomOut, fitWindow, fitWidth) self.actions.advanced = ( - open, opendir, changeSavedir, openNextImg, openPrevImg, save, None, + open, opendir, changeSavedir, openNextImg, openPrevImg, save, save_format, None, createMode, editMode, None, hideAll, showAll) @@ -466,6 +472,14 @@ class MainWindow(QMainWindow, WindowMixin): ## Support Functions ## + def change_format(self): + self.usingPascalVocFormat = not self.usingPascalVocFormat + self.usingYoloFormat = not self.usingYoloFormat + print ("changing_format") + print(self.actions.save_format) + if self.usingPascalVocFormat: self.actions.save_format.setText("PascalVOC") + if self.usingYoloFormat: self.actions.save_format.setText("YOLO") + def noShapes(self): return not self.itemsToShapes @@ -733,9 +747,15 @@ class MainWindow(QMainWindow, WindowMixin): # Can add differrent annotation formats here try: if self.usingPascalVocFormat is True: + annotationFilePath += XML_EXT print ('Img: ' + self.filePath + ' -> Its xml: ' + annotationFilePath) self.labelFile.savePascalVocFormat(annotationFilePath, shapes, self.filePath, self.imageData, self.lineColor.getRgb(), self.fillColor.getRgb()) + elif self.usingYoloFormat is True: + annotationFilePath += TXT_EXT + print ('Img: ' + self.filePath + ' -> Its txt: ' + annotationFilePath) + self.labelFile.saveYoloFormat(annotationFilePath, shapes, self.filePath, self.imageData, self.labelHist, + self.lineColor.getRgb(), self.fillColor.getRgb()) else: self.labelFile.save(annotationFilePath, shapes, self.filePath, self.imageData, self.lineColor.getRgb(), self.fillColor.getRgb()) @@ -1197,13 +1217,13 @@ class MainWindow(QMainWindow, WindowMixin): if self.defaultSaveDir is not None and len(ustr(self.defaultSaveDir)): if self.filePath: imgFileName = os.path.basename(self.filePath) - savedFileName = os.path.splitext(imgFileName)[0] + XML_EXT + savedFileName = os.path.splitext(imgFileName)[0] savedPath = os.path.join(ustr(self.defaultSaveDir), savedFileName) self._saveFile(savedPath) else: imgFileDir = os.path.dirname(self.filePath) imgFileName = os.path.basename(self.filePath) - savedFileName = os.path.splitext(imgFileName)[0] + XML_EXT + savedFileName = os.path.splitext(imgFileName)[0] savedPath = os.path.join(imgFileDir, savedFileName) self._saveFile(savedPath if self.labelFile else self.saveFileDialog()) diff --git a/libs/labelFile.py b/libs/labelFile.py index 7918aba2..9a3c54ee 100644 --- a/libs/labelFile.py +++ b/libs/labelFile.py @@ -8,6 +8,7 @@ except ImportError: from base64 import b64encode, b64decode from libs.pascal_voc_io import PascalVocWriter +from libs.yolo_io import YOLOWriter from libs.pascal_voc_io import XML_EXT import os.path import sys @@ -55,6 +56,33 @@ class LabelFile(object): writer.save(targetFile=filename) return + def saveYoloFormat(self, filename, shapes, imagePath, imageData, classList, + lineColor=None, fillColor=None, databaseSrc=None): + imgFolderPath = os.path.dirname(imagePath) + imgFolderName = os.path.split(imgFolderPath)[-1] + imgFileName = os.path.basename(imagePath) + #imgFileNameWithoutExt = os.path.splitext(imgFileName)[0] + # Read from file path because self.imageData might be empty if saving to + # Pascal format + image = QImage() + image.load(imagePath) + imageShape = [image.height(), image.width(), + 1 if image.isGrayscale() else 3] + writer = YOLOWriter(imgFolderName, imgFileName, + imageShape, localImgPath=imagePath) + writer.verified = self.verified + + for shape in shapes: + points = shape['points'] + label = shape['label'] + # Add Chris + difficult = int(shape['difficult']) + bndbox = LabelFile.convertPoints2BndBox(points) + writer.addBndBox(bndbox[0], bndbox[1], bndbox[2], bndbox[3], label, difficult) + + writer.save(targetFile=filename, classList=classList) + return + def toggleVerify(self): self.verified = not self.verified diff --git a/libs/lib.py b/libs/lib.py index da172dc2..13673cb8 100644 --- a/libs/lib.py +++ b/libs/lib.py @@ -75,7 +75,7 @@ def fmtShortcut(text): def generateColorByText(text): - s = str(ustr(text)) + s = str(ustr(text)) hashCode = int(hashlib.sha256(s.encode('utf-8')).hexdigest(), 16) r = int((hashCode / 255) % 255) g = int((hashCode / 65025) % 255) diff --git a/libs/yolo_io.py b/libs/yolo_io.py new file mode 100644 index 00000000..36e75927 --- /dev/null +++ b/libs/yolo_io.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python +# -*- coding: utf8 -*- +import sys +import os +from xml.etree import ElementTree +from xml.etree.ElementTree import Element, SubElement +from lxml import etree +import codecs + +TXT_EXT = '.txt' +ENCODE_METHOD = 'utf-8' + +class YOLOWriter: + + def __init__(self, foldername, filename, imgSize, databaseSrc='Unknown', localImgPath=None): + self.foldername = foldername + self.filename = filename + self.databaseSrc = databaseSrc + self.imgSize = imgSize + self.boxlist = [] + self.localImgPath = localImgPath + self.verified = False + + def addBndBox(self, xmin, ymin, xmax, ymax, name, difficult): + bndbox = {'xmin': xmin, 'ymin': ymin, 'xmax': xmax, 'ymax': ymax} + bndbox['name'] = name + bndbox['difficult'] = difficult + self.boxlist.append(bndbox) + + def BndBox2YoloLine(self, box, classList=[]): + xmin = box['xmin'] + xmax = box['xmax'] + ymin = box['ymin'] + ymax = box['ymax'] + + xcen = (xmin + xmax) / 2 / self.imgSize[1] + ycen = (ymin + ymax) / 2 / self.imgSize[0] + + w = (xmax - xmin) / self.imgSize[1] + h = (ymax - ymin) / self.imgSize[0] + + classIndex = classList.index(box['name']) + + return classIndex, xcen, ycen, w, h + + def save(self, classList=[], targetFile=None): + + out_file = None #Update yolo .txt + out_class_file = None #Update class list .txt + + if targetFile is None: + out_file = open( + self.filename + TXT_EXT, 'w', encoding=ENCODE_METHOD) + classesFile = os.path.join(os.path.dirname(os.path.abspath(self.filename)), "classes.txt") + out_class_file = open(classesFile, 'w') + + else: + out_file = codecs.open(targetFile, 'w', encoding=ENCODE_METHOD) + classesFile = os.path.join(os.path.dirname(os.path.abspath(targetFile)), "classes.txt") + out_class_file = open(classesFile, 'w') + + + for box in self.boxlist: + classIndex, xcen, ycen, w, h = self.BndBox2YoloLine(box, classList) + print (classIndex, xcen, ycen, w, h) + out_file.write("%d %.6f %.6f %.6f %.6f\n" % (classIndex, xcen, ycen, w, h)) + + print (classList) + print (out_class_file) + for c in classList: + out_class_file.write(c+'\n') + + out_class_file.close() + out_file.close() + + + +class YoloReader: + + def __init__(self, filepath, imgSize, classListPath=None): + # shapes type: + # [labbel, [(x1,y1), (x2,y2), (x3,y3), (x4,y4)], color, color, difficult] + self.shapes = [] + self.filepath = filepath + + if classListPath == None: + dir_path = os.path.dirname(os.path.realpath(self.filepath)) + self.classListPath = os.path.join(dir_path, "classes.txt") + else: + self.classListPath = classListPath + + classesFile = open(self.classListPath, 'r') + self.classes = classesFile.read().split('\n') + + self.imgSize = imgSize + + self.verified = False + try: + self.parseYoloFormat() + except: + pass + + def getShapes(self): + return self.shapes + + def addShape(self, label, xmin, ymin, xmax, ymax, difficult): + + points = [(xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax)] + self.shapes.append((label, points, None, None, difficult)) + + def yoloLine2Shape(self, classIndex, xcen, ycen, w, h): + label = self.classes[classIndex] + + xmin = min(float(xcen) - float(w) / 2, 0) + xmax = max(float(xcen) + float(w) / 2, 1) + ymin = min(float(ycen) - float(h) / 2, 0) + ymax = max(float(ycen) + float(h) / 2, 1) + + xmin = int(imgSize[1] * xmin) + xmax = int(imgSize[1] * xmax) + ymin = int(imgSize[0] * ymin) + ymax = int(imgSize[0] * ymax) + + return label, xmin, ymin, xmax, ymax + + def parseYoloFormat(self): + bndBoxFile = open(self.filepath, 'r') + for bndBox in bndBoxFile: + classIndex, xcen, ycen, w, h = bndBox.split(' ') + label, xmin, ymin, xmax, ymax = yoloLine2Shape(classIndex, xcen, ycen, w, h) + + # Caveat: difficult flag is discarded when saved as yolo format. + self.addShape(label, xmin, ymin, xmax, ymax, False)