Merge pull request #247 from isVoid/yolosupport

Yolo support

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
 ~~~~~~~~~~~~~~~~~~~~~~~~~~
 

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()])

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
 

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)

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)

resources.qrc

@@ -18,6 +18,8 @@
 <file alias="edit">icons/edit.png</file>
 <file alias="open">icons/open.png</file>
 <file alias="save">icons/save.png</file>
+<file alias="format_voc">icons/format_voc.png</file>
+<file alias="format_yolo">icons/format_yolo.png</file>
 <file alias="save-as">icons/save-as.png</file>
 <file alias="color">icons/color.png</file>
 <file alias="color_line">icons/color_line.png</file>