Merge pull request HumanSignal#80 from ChrisDal/master

Add Difficult checkbox

labelImg.py

@@ -139,10 +139,19 @@ def __init__(self, defaultFilename=None):
         listLayout.addWidget(self.labelList)
         self.editButton = QToolButton()
         self.editButton.setToolButtonStyle(Qt.ToolButtonTextBesideIcon)
+
+        # Add chris 
+        self.diffcButton = QCheckBox("Difficult")
+        self.diffcButton.setChecked(False) 
+        self.diffcButton.stateChanged.connect(self.btnstate)
+
         self.labelListContainer = QWidget()
         self.labelListContainer.setLayout(listLayout)
         listLayout.addWidget(self.editButton)  # , 0, Qt.AlignCenter)
+        # Add chris 
+        listLayout.addWidget(self.diffcButton) 
         listLayout.addWidget(self.labelList)
+
 
         self.dock = QDockWidget(u'Box Labels', self)
         self.dock.setObjectName(u'Labels')
@@ -376,6 +385,8 @@ def __init__(self, defaultFilename=None):
         self.fillColor = None
         self.zoom_level = 100
         self.fit_window = False
+        # Add Chris
+        self.difficult = False 
 
         # XXX: Could be completely declarative.
         # Restore application settings.
@@ -431,6 +442,8 @@ def __init__(self, defaultFilename=None):
         self.fillColor = QColor(settings.get('fill/color', Shape.fill_color))
         Shape.line_color = self.lineColor
         Shape.fill_color = self.fillColor
+        # Add chris 
+        Shape.Difficult = self.difficult
 
         def xbool(x):
             if isinstance(x, QVariant):
@@ -608,6 +621,33 @@ def fileitemDoubleClicked(self, item=None):
             filename = self.mImgList[currIndex]
             if filename:
                 self.loadFile(filename)
+
+    # Add chris       
+    def btnstate(self, item= None):
+        """ Function to handle difficult examples 
+        Update on each object """
+        if not self.canvas.editing():
+            return
+
+        item = self.currentItem()
+        if not item: # If not selected Item, take the first one 
+            item = self.labelList.item(self.labelList.count()-1)
+
+        difficult = self.diffcButton.isChecked()
+
+        try:
+            shape = self.itemsToShapes[item]
+        except:
+            pass
+        # Checked and Update 
+        try:
+            if difficult != shape.difficult:
+                shape.difficult = difficult
+                self.setDirty()
+            else:  # User probably changed item visibility
+                self.canvas.setShapeVisible(shape, item.checkState() == Qt.Checked)
+        except:
+            pass
 
     # React to canvas signals.
     def shapeSelectionChanged(self, selected=False):
@@ -646,17 +686,20 @@ def remLabel(self, shape):
 
     def loadLabels(self, shapes):
         s = []
-        for label, points, line_color, fill_color in shapes:
+        for label, points, line_color, fill_color, difficult in shapes:
             shape = Shape(label=label)
             for x, y in points:
                 shape.addPoint(QPointF(x, y))
+            shape.difficult = difficult
             shape.close()
             s.append(shape)
             self.addLabel(shape)
+
             if line_color:
                 shape.line_color = QColor(*line_color)
             if fill_color:
                 shape.fill_color = QColor(*fill_color)
+
         self.canvas.loadShapes(s)
 
     def saveLabels(self, annotationFilePath):
@@ -671,7 +714,9 @@ def format_shape(s):
                         if s.line_color != self.lineColor else None,
                         fill_color=s.fill_color.getRgb()
                         if s.fill_color != self.fillColor else None,
-                        points=[(p.x(), p.y()) for p in s.points])
+                        points=[(p.x(), p.y()) for p in s.points],
+                       # add chris
+                        difficult = s.difficult)
 
         shapes = [format_shape(shape) for shape in self.canvas.shapes]
         # Can add differrent annotation formats here
@@ -700,6 +745,9 @@ def labelSelectionChanged(self):
         if item and self.canvas.editing():
             self._noSelectionSlot = True
             self.canvas.selectShape(self.itemsToShapes[item])
+            shape = self.itemsToShapes[item]
+            # Add Chris
+            self.diffcButton.setChecked(shape.difficult)
 
     def labelItemChanged(self, item):
         shape = self.itemsToShapes[item]
@@ -721,6 +769,8 @@ def newShape(self):
                 parent=self, listItem=self.labelHist)
 
         text = self.labelDialog.popUp(text=self.prevLabelText)
+        # Add Chris
+        self.diffcButton.setChecked(False)
         if text is not None:
             self.prevLabelText = text
             self.addLabel(self.canvas.setLastLabel(text))
@@ -838,7 +888,12 @@ def loadFile(self, filePath=None):
                         self.loadPascalXMLByFilename(xmlPath)
 
             self.setWindowTitle(__appname__ + ' ' + filePath)
-
+
+            # Default : select last item if there is at least one item 
+            if self.labelList.count():
+                self.labelList.setCurrentItem(self.labelList.item(self.labelList.count()-1))
+                self.labelList.setItemSelected(self.labelList.item(self.labelList.count()-1), True)
+
             self.canvas.setFocus(True)
             return True
         return False

libs/labelFile.py

@@ -45,8 +45,10 @@ def savePascalVocFormat(self, filename, shapes, imagePath, imageData,
         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)
+            writer.addBndBox(bndbox[0], bndbox[1], bndbox[2], bndbox[3], label, difficult)
 
         writer.save(targetFile=filename)
         return

libs/pascal_voc_io.py

@@ -2,6 +2,8 @@
 # -*- coding: utf8 -*-
 import _init_path
 import sys
+from xml.etree import ElementTree
+from xml.etree.ElementTree import Element, SubElement
 from lxml import etree
 import codecs
 
@@ -10,7 +12,7 @@
 
 class PascalVocWriter:
 
-    def __init__(self, foldername, filename, imgSize, databaseSrc='Unknown', localImgPath=None):
+    def __init__(self, foldername, filename, imgSize,databaseSrc='Unknown', localImgPath=None):
         self.foldername = foldername
         self.filename = filename
         self.databaseSrc = databaseSrc
@@ -19,14 +21,14 @@ def __init__(self, foldername, filename, imgSize, databaseSrc='Unknown', localIm
         self.localImgPath = localImgPath
         self.verified = False
 
+
     def prettify(self, elem):
         """
             Return a pretty-printed XML string for the Element.
         """
-        rough_string = etree.tostring(elem, encoding='UTF-8')
-        rough_string = str(rough_string, encoding="UTF-8")
-        root = etree.XML(rough_string)
-        return etree.tostring(root, encoding='UTF-8', pretty_print=True)
+        rough_string = ElementTree.tostring(elem, 'utf8')
+        root = etree.fromstring(rough_string)
+        return etree.tostring(root, pretty_print=True)
 
     def genXML(self):
         """
@@ -38,65 +40,73 @@ def genXML(self):
                 self.imgSize is None:
             return None
 
-        top = etree.Element('annotation')
+        top = Element('annotation')
         top.set('verified', 'yes' if self.verified else 'no')
 
-        folder = etree.SubElement(top, 'folder')
+        folder = SubElement(top, 'folder')
         folder.text = self.foldername
 
-        filename = etree.SubElement(top, 'filename')
+        filename = SubElement(top, 'filename')
         filename.text = self.filename
 
-        localImgPath = etree.SubElement(top, 'path')
+        localImgPath = SubElement(top, 'path')
         localImgPath.text = self.localImgPath
 
-        source = etree.SubElement(top, 'source')
-        database = etree.SubElement(source, 'database')
+        source = SubElement(top, 'source')
+        database = SubElement(source, 'database')
         database.text = self.databaseSrc
 
-        size_part = etree.SubElement(top, 'size')
-        width = etree.SubElement(size_part, 'width')
-        height = etree.SubElement(size_part, 'height')
-        depth = etree.SubElement(size_part, 'depth')
+        size_part = SubElement(top, 'size')
+        width = SubElement(size_part, 'width')
+        height = SubElement(size_part, 'height')
+        depth = SubElement(size_part, 'depth')
         width.text = str(self.imgSize[1])
         height.text = str(self.imgSize[0])
         if len(self.imgSize) == 3:
             depth.text = str(self.imgSize[2])
         else:
             depth.text = '1'
 
-        segmented = etree.SubElement(top, 'segmented')
+        segmented = SubElement(top, 'segmented')
         segmented.text = '0'
         return top
 
-    def addBndBox(self, xmin, ymin, xmax, ymax, name):
+    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 appendObjects(self, top):
         for each_object in self.boxlist:
-            object_item = etree.SubElement(top, 'object')
-            name = etree.SubElement(object_item, 'name')
+            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']
-            pose = etree.SubElement(object_item, 'pose')
+            pose = SubElement(object_item, 'pose')
             pose.text = "Unspecified"
-            truncated = etree.SubElement(object_item, 'truncated')
-            truncated.text = "0"
-            difficult = etree.SubElement(object_item, 'difficult')
-            difficult.text = "0"
-            bndbox = etree.SubElement(object_item, 'bndbox')
-            xmin = etree.SubElement(bndbox, 'xmin')
+            truncated = SubElement(object_item, 'truncated')
+            # max == height or min 
+            if int(each_object['ymax']) == int(self.imgSize[0]) or (int(each_object['ymin'])== 1):
+		        truncated.text = "1"
+            # max == width or min 
+            elif (int(each_object['xmax'])==int(self.imgSize[1])) or (int(each_object['xmin'])== 1):
+                truncated.text = "1"
+            else:
+                truncated.text = "0"
+            difficult = SubElement(object_item, 'Difficult')
+            difficult.text = str( bool(each_object['difficult']) & 1 )
+            bndbox = SubElement(object_item, 'bndbox')
+            xmin = SubElement(bndbox, 'xmin')
             xmin.text = str(each_object['xmin'])
-            ymin = etree.SubElement(bndbox, 'ymin')
+            ymin = SubElement(bndbox, 'ymin')
             ymin.text = str(each_object['ymin'])
-            xmax = etree.SubElement(bndbox, 'xmax')
+            xmax = SubElement(bndbox, 'xmax')
             xmax.text = str(each_object['xmax'])
-            ymax = etree.SubElement(bndbox, 'ymax')
+            ymax = SubElement(bndbox, 'ymax')
             ymax.text = str(each_object['ymax'])
 
     def save(self, targetFile=None):
@@ -118,7 +128,7 @@ class PascalVocReader:
 
     def __init__(self, filepath):
         # shapes type:
-        # [labbel, [(x1,y1), (x2,y2), (x3,y3), (x4,y4)], color, color]
+        # [labbel, [(x1,y1), (x2,y2), (x3,y3), (x4,y4)], color, color, difficult]
         self.shapes = []
         self.filepath = filepath
         self.verified = False
@@ -127,24 +137,19 @@ def __init__(self, filepath):
     def getShapes(self):
         return self.shapes
 
-    def addShape(self, label, bndbox):
+    def addShape(self, label, bndbox, difficult):
         xmin = int(bndbox.find('xmin').text)
         ymin = int(bndbox.find('ymin').text)
         xmax = int(bndbox.find('xmax').text)
         ymax = int(bndbox.find('ymax').text)
         points = [(xmin, ymin), (xmax, ymin), (xmax, ymax), (xmin, ymax)]
-        self.shapes.append((label, points, None, None))
+
+        self.shapes.append((label, points, None, None, difficult))
 
     def parseXML(self):
         assert self.filepath.endswith('.xml'), "Unsupport file format"
-        content = None
-        with open(self.filepath, 'r') as xmlFile:
-            content = xmlFile.read()
-
-        if content is None:
-            return False
-
-        xmltree = etree.XML(content)
+        parser = etree.XMLParser(encoding='utf-8')
+        xmltree = ElementTree.parse(self.filepath, parser=parser).getroot()
         filename = xmltree.find('filename').text
         try:
             verified = xmltree.attrib['verified']
@@ -156,16 +161,7 @@ def parseXML(self):
         for object_iter in xmltree.findall('object'):
             bndbox = object_iter.find("bndbox")
             label = object_iter.find('name').text
-            self.addShape(label, bndbox)
+            # Add chris
+            difficult = bool(int(object_iter.find('Difficult').text))
+            self.addShape(label, bndbox, difficult)
         return True
-
-
-# tempParseReader = PascalVocReader('test.xml')
-# print tempParseReader.getShapes()
-"""
-# Test
-tmp = PascalVocWriter('temp','test', (10,20,3))
-tmp.addBndBox(10,10,20,30,'chair')
-tmp.addBndBox(1,1,600,600,'car')
-tmp.save()
-"""

libs/shape.py

@@ -36,11 +36,12 @@ class Shape(object):
     point_size = 8
     scale = 1.0
 
-    def __init__(self, label=None, line_color=None):
+    def __init__(self, label=None, line_color=None,difficult = False):
         self.label = label
         self.points = []
         self.fill = False
         self.selected = False
+        self.difficult = difficult
 
         self._highlightIndex = None
         self._highlightMode = self.NEAR_VERTEX
@@ -171,6 +172,7 @@ def copy(self):
             shape.line_color = self.line_color
         if self.fill_color != Shape.fill_color:
             shape.fill_color = self.fill_color
+        shape.difficult = self.difficult 
         return shape
 
     def __len__(self):