YOLO Write

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

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,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)