diff --git a/README.rst b/README.rst index 8d36e139..ef0fc69c 100644 --- a/README.rst +++ b/README.rst @@ -135,7 +135,7 @@ You can pull the image which has all of the installed and required dependencies. Usage ----- -Steps +Steps (PascalVOC) ~~~~~ 1. Build and launch using the instructions above. @@ -150,6 +150,27 @@ The annotation will be saved to the folder you specify. You can refer to the below hotkeys to speed up your workflow. +Steps (YOLO) +~~~~~ + +1. In ``data/predefined_classes.txt`` define the list of classes that will be used for your training. + +2. Build and launch using the instructions above. + +3. Right below "Save" button in toolbar, click "PascalVOC" button to switch to YOLO format. + +4. You may use Open/OpenDIR to process single or multiple images. When finished with single image, click save. + +A txt file of yolo format will be saved in the same folder as your image with same name. A file named "classes.txt" is saved to that folder too. "classes.txt" defines the list of class names that your yolo label refers to. + +Note: + +- Your label list shall not change in the middle of processing a list of images. When you save a image, classes.txt will also get updated, while previous annotations will not be updated. + +- You shouldn't use "default class" function when saving to YOLO format, it will not be referred. + +- When saving as YOLO format, "difficult" flag is discarded. + Create pre-defined classes ~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/icons/format_voc.png b/icons/format_voc.png new file mode 100644 index 00000000..cb15e439 Binary files /dev/null and b/icons/format_voc.png differ diff --git a/icons/format_yolo.png b/icons/format_yolo.png new file mode 100644 index 00000000..ca9acc71 Binary files /dev/null and b/icons/format_yolo.png differ diff --git a/labelImg.py b/labelImg.py index a817a653..02b5df64 100755 --- a/labelImg.py +++ b/labelImg.py @@ -38,6 +38,8 @@ 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 YoloReader +from libs.yolo_io import TXT_EXT from libs.ustr import ustr from libs.version import __version__ @@ -97,6 +99,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 +236,9 @@ class MainWindow(QMainWindow, WindowMixin): save = action('&Save', self.saveFile, 'Ctrl+S', 'save', u'Save labels to file', enabled=False) + save_format = action('&PascalVOC', self.change_format, + 'Ctrl+', 'format_voc', 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 +331,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 +370,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 +390,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) @@ -465,6 +472,22 @@ class MainWindow(QMainWindow, WindowMixin): self.openDirDialog(dirpath=self.filePath) ## Support Functions ## + def set_format(self, save_format): + if save_format == 'PascalVOC': + self.actions.save_format.setText("PascalVOC") + self.actions.save_format.setIcon(newIcon("format_voc")) + self.usingPascalVocFormat = True + self.usingYoloFormat = False + + elif save_format == 'YOLO': + self.actions.save_format.setText("YOLO") + self.actions.save_format.setIcon(newIcon("format_yolo")) + self.usingPascalVocFormat = False + self.usingYoloFormat = True + + def change_format(self): + if self.usingPascalVocFormat: self.set_format("YOLO") + elif self.usingYoloFormat: self.set_format("PascalVOC") def noShapes(self): return not self.itemsToShapes @@ -733,9 +756,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()) @@ -948,16 +977,27 @@ class MainWindow(QMainWindow, WindowMixin): self.toggleActions(True) # Label xml file and show bound box according to its filename - if self.usingPascalVocFormat is True: - if self.defaultSaveDir is not None: - basename = os.path.basename( - os.path.splitext(self.filePath)[0]) + XML_EXT - xmlPath = os.path.join(self.defaultSaveDir, basename) + # if self.usingPascalVocFormat is True: + if self.defaultSaveDir is not None: + basename = os.path.basename( + os.path.splitext(self.filePath)[0]) + xmlPath = os.path.join(self.defaultSaveDir, basename + XML_EXT) + txtPath = os.path.join(self.defaultSaveDir, basename + TXT_EXT) + + """Annotation file priority: + PascalXML > YOLO + """ + if os.path.isfile(xmlPath): self.loadPascalXMLByFilename(xmlPath) - else: - xmlPath = os.path.splitext(filePath)[0] + XML_EXT - if os.path.isfile(xmlPath): - self.loadPascalXMLByFilename(xmlPath) + elif os.path.isfile(txtPath): + self.loadYOLOTXTByFilename(txtPath) + else: + xmlPath = os.path.splitext(filePath)[0] + XML_EXT + txtPath = os.path.splitext(filePath)[0] + TXT_EXT + if os.path.isfile(xmlPath): + self.loadPascalXMLByFilename(xmlPath) + elif os.path.isfile(txtPath): + self.loadYOLOTXTByFilename(txtPath) self.setWindowTitle(__appname__ + ' ' + filePath) @@ -1197,13 +1237,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()) @@ -1320,11 +1360,27 @@ class MainWindow(QMainWindow, WindowMixin): if os.path.isfile(xmlPath) is False: return + self.set_format("PascalVOC") + tVocParseReader = PascalVocReader(xmlPath) shapes = tVocParseReader.getShapes() self.loadLabels(shapes) self.canvas.verified = tVocParseReader.verified + def loadYOLOTXTByFilename(self, txtPath): + if self.filePath is None: + return + if os.path.isfile(txtPath) is False: + return + + self.set_format("YOLO") + tYoloParseReader = YoloReader(txtPath, self.image) + shapes = tYoloParseReader.getShapes() + print (shapes) + self.loadLabels(shapes) + self.canvas.verified = tYoloParseReader.verified + + def inverted(color): return QColor(*[255 - v for v in color.getRgb()]) 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..f585af08 --- /dev/null +++ b/libs/yolo_io.py @@ -0,0 +1,140 @@ +#!/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, image, classListPath=None): + # shapes type: + # [labbel, [(x1,y1), (x2,y2), (x3,y3), (x4,y4)], color, color, difficult] + self.shapes = [] + self.filepath = filepath + + if classListPath is None: + dir_path = os.path.dirname(os.path.realpath(self.filepath)) + self.classListPath = os.path.join(dir_path, "classes.txt") + else: + self.classListPath = classListPath + + print (filepath, self.classListPath) + + classesFile = open(self.classListPath, 'r') + self.classes = classesFile.read().strip('\n').split('\n') + + print (self.classes) + + imgSize = [image.height(), image.width(), + 1 if image.isGrayscale() else 3] + + 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[int(classIndex)] + + xmin = max(float(xcen) - float(w) / 2, 0) + xmax = min(float(xcen) + float(w) / 2, 1) + ymin = max(float(ycen) - float(h) / 2, 0) + ymax = min(float(ycen) + float(h) / 2, 1) + + xmin = int(self.imgSize[1] * xmin) + xmax = int(self.imgSize[1] * xmax) + ymin = int(self.imgSize[0] * ymin) + ymax = int(self.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 = self.yoloLine2Shape(classIndex, xcen, ycen, w, h) + + # Caveat: difficult flag is discarded when saved as yolo format. + self.addShape(label, xmin, ymin, xmax, ymax, False) diff --git a/resources.qrc b/resources.qrc index f6b30625..805fe421 100644 --- a/resources.qrc +++ b/resources.qrc @@ -18,6 +18,8 @@ icons/edit.png icons/open.png icons/save.png +icons/format_voc.png +icons/format_yolo.png icons/save-as.png icons/color.png icons/color_line.png