YOLO Write

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)