Create a const for encoding and fix an issue in yolo format according…

… to pr#387

README.rst

@@ -12,7 +12,7 @@ LabelImg is a graphical image annotation tool.
 It is written in Python and uses Qt for its graphical interface.
 
 Annotations are saved as XML files in PASCAL VOC format, the format used
-by `ImageNet <http://www.image-net.org/>`__.
+by `ImageNet <http://www.image-net.org/>`__.  Besdies, it also supports YOLO format
 
 .. image:: https://raw.githubusercontent.com/tzutalin/labelImg/master/demo/demo3.jpg
      :alt: Demo Image

libs/constants.py

@@ -15,3 +15,4 @@
 FORMAT_PASCALVOC='PascalVOC'
 FORMAT_YOLO='YOLO'
 SETTING_DRAW_SQUARE = 'draw/square'
+DEFAULT_ENCODING = 'utf-8'

libs/pascal_voc_io.py

@@ -5,9 +5,10 @@
 from xml.etree.ElementTree import Element, SubElement
 from lxml import etree
 import codecs
+from libs.constants import DEFAULT_ENCODING
 
 XML_EXT = '.xml'
-ENCODE_METHOD = 'utf-8'
+ENCODE_METHOD = DEFAULT_ENCODING
 
 class PascalVocWriter:
 
@@ -84,11 +85,8 @@ def appendObjects(self, top):
         for each_object in self.boxlist:
             object_item = SubElement(top, 'object')
             name = SubElement(object_item, 'name')
-            try:
-                name.text = unicode(each_object['name'])
-            except NameError:
-                # Py3: NameError: name 'unicode' is not defined
-                name.text = each_object['name']
+            print (each_object['name'])
+            name.text = each_object['name']
             pose = SubElement(object_item, 'pose')
             pose.text = "Unspecified"
             truncated = SubElement(object_item, 'truncated')

libs/ustr.py

@@ -1,14 +1,15 @@
 import sys
+from libs.constants import DEFAULT_ENCODING
 
 def ustr(x):
     '''py2/py3 unicode helper'''
 
     if sys.version_info < (3, 0, 0):
         from PyQt4.QtCore import QString
         if type(x) == str:
-            return x.decode('utf-8')
+            return x.decode(DEFAULT_ENCODING)
         if type(x) == QString:
-            return unicode(x)
+            return unicode(x, DEFAULT_ENCODING)
         return x
     else:
-        return x  # py3
+        return x

libs/yolo_io.py

@@ -6,9 +6,10 @@
 from xml.etree.ElementTree import Element, SubElement
 from lxml import etree
 import codecs
+from libs.constants import DEFAULT_ENCODING
 
 TXT_EXT = '.txt'
-ENCODE_METHOD = 'utf-8'
+ENCODE_METHOD = DEFAULT_ENCODING
 
 class YOLOWriter:
 
@@ -39,7 +40,12 @@ def BndBox2YoloLine(self, box, classList=[]):
         w = float((xmax - xmin)) / self.imgSize[1]
         h = float((ymax - ymin)) / self.imgSize[0]
 
-        classIndex = classList.index(box['name'])
+        # PR387
+        boxName = box['name']
+        if boxName not in classList:
+            classList.append(boxName)
+
+        classIndex = classList.index(boxName)
 
         return classIndex, xcen, ycen, w, h
 
@@ -62,11 +68,11 @@ def save(self, classList=[], targetFile=None):
 
         for box in self.boxlist:
             classIndex, xcen, ycen, w, h = self.BndBox2YoloLine(box, classList)
-            print (classIndex, xcen, ycen, w, h)
+            # 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)
+        # print (classList)
+        # print (out_class_file)
         for c in classList:
             out_class_file.write(c+'\n')
 
@@ -89,12 +95,12 @@ def __init__(self, filepath, image, classListPath=None):
         else:
             self.classListPath = classListPath
 
-        print (filepath, self.classListPath)
+        # print (filepath, self.classListPath)
 
         classesFile = open(self.classListPath, 'r')
         self.classes = classesFile.read().strip('\n').split('\n')
 
-        print (self.classes)
+        # print (self.classes)
 
         imgSize = [image.height(), image.width(),
                       1 if image.isGrayscale() else 3]