From ed787d4f3a56b80f61ee40ac7b93b56d3a8b7512 Mon Sep 17 00:00:00 2001
From: tzutalin <tzu.ta.lin@gmail.com>
Date: Wed, 9 Dec 2015 21:29:26 +0800
Subject: [PATCH] Add a function which can load xml file

---
 labelImg.py                                   | 25 ++++++++++--
 libs/labelFile.py                             |  2 +-
 ...{pascal_voc_writer.py => pascal_voc_io.py} | 39 +++++++++++++++++++
 3 files changed, 61 insertions(+), 5 deletions(-)
 rename libs/{pascal_voc_writer.py => pascal_voc_io.py} (75%)

diff --git a/labelImg.py b/labelImg.py
index d202ba49..4fdd3d97 100755
--- a/labelImg.py
+++ b/labelImg.py
@@ -22,6 +22,7 @@ from labelDialog import LabelDialog
 from colorDialog import ColorDialog
 from labelFile import LabelFile, LabelFileError
 from toolBar import ToolBar
+from pascal_voc_io import PascalVocReader
 
 __appname__ = 'labelImg'
 
@@ -139,6 +140,9 @@ class MainWindow(QMainWindow, WindowMixin):
         changeSavedir = action('&Change default saved Annotation dir', self.changeSavedir,
                 'Ctrl+r', 'open', u'Change default saved Annotation dir')
 
+        openAnnotation = action('&Open Annotation', self.openAnnotation,
+                'Ctrl+q', 'openAnnotation', u'Open Annotation')
+
         openNextImg = action('&Next Image', self.openNextImg,
                 'n', 'next', u'Open Next')
 
@@ -263,7 +267,7 @@ class MainWindow(QMainWindow, WindowMixin):
                 labelList=labelMenu)
 
         addActions(self.menus.file,
-                (open, opendir,changeSavedir, self.menus.recentFiles, save, saveAs, close, None, quit))
+                (open, opendir,changeSavedir, openAnnotation, self.menus.recentFiles, save, saveAs, close, None, quit))
         addActions(self.menus.help, (help,))
         addActions(self.menus.view, (
             labels, advancedMode, None,
@@ -797,6 +801,22 @@ class MainWindow(QMainWindow, WindowMixin):
         self.statusBar().showMessage('%s . Annotation will be saved to %s' %('Change saved folder', self.defaultSaveDir))
         self.statusBar().show()
 
+    def openAnnotation(self, _value=False):
+        path = os.path.dirname(unicode(self.filename))\
+                if self.filename else '.'
+
+        formats = ['*.%s' % unicode(fmt).lower()\
+                for fmt in QImageReader.supportedImageFormats()]
+        filters = "Open Annotation XML file (%s)" % \
+                ' '.join(formats + ['*.xml'])
+        filename = unicode(QFileDialog.getOpenFileName(self,
+            '%s - Choose a xml file' % __appname__, path, filters))
+
+        tVocParseReader = PascalVocReader(filename)
+        shapes = tVocParseReader.getShapes()
+        self.loadLabels(shapes)
+        return
+
     def openDir(self, _value=False):
         if not self.mayContinue():
             return
@@ -834,8 +854,6 @@ class MainWindow(QMainWindow, WindowMixin):
             if filename:
                 self.loadFile(filename)
 
-
-
     def openNextImg(self, _value=False):
         # Proceding next image without dialog if having any label
         if self.autoSaving is True and self.defaultSaveDir is not None:
@@ -1039,7 +1057,6 @@ class Settings(object):
 def inverted(color):
     return QColor(*[255 - v for v in color.getRgb()])
 
-
 def read(filename, default=None):
     try:
         with open(filename, 'rb') as f:
diff --git a/libs/labelFile.py b/libs/labelFile.py
index 41fc7ce4..bf984352 100644
--- a/libs/labelFile.py
+++ b/libs/labelFile.py
@@ -3,7 +3,7 @@ import os.path
 import numpy
 import cv2
 import sys
-from pascal_voc_writer import PascalVocWriter
+from pascal_voc_io import PascalVocWriter
 from base64 import b64encode, b64decode
 
 class LabelFileError(Exception):
diff --git a/libs/pascal_voc_writer.py b/libs/pascal_voc_io.py
similarity index 75%
rename from libs/pascal_voc_writer.py
rename to libs/pascal_voc_io.py
index 241c7f22..05ce8749 100644
--- a/libs/pascal_voc_writer.py
+++ b/libs/pascal_voc_io.py
@@ -100,6 +100,45 @@ class PascalVocWriter:
         out_file.write(self.prettify(root))
         out_file.close()
 
+
+class PascalVocReader:
+
+    def __init__(self, filepath):
+        ## shapes type:
+        ## [labbel, [(x1,y1), (x2,y2), (x3,y3), (x4,y4)], color, color]
+        self.shapes=[]
+        self.filepath = filepath
+        self.parseXML()
+
+    def getShapes(self):
+        return self.shapes
+
+    def addShape(self, label, rect):
+        xmin = rect[0]
+        ymin = rect[1]
+        xmax = rect[2]
+        ymax = rect[3]
+        points = [(xmin,ymin), (xmin,ymax), (xmax, ymax), (xmax, ymin)]
+        self.shapes.append((label, points, None, None))
+
+    def parseXML(self):
+        assert self.filepath.endswith('.xml'), "Unsupport file format"
+        xmltree = ElementTree.parse(self.filepath).getroot()
+        filename = xmltree.find('filename').text
+
+        for object_iter in xmltree.findall('object'):
+           rects = []
+           bndbox = object_iter.find("bndbox")
+           rects.append([int(it.text) for it in bndbox])
+           label = object_iter.find('name').text
+
+           for rect in rects:
+               self.addShape(label, rect)
+        return True
+
+
+# tempParseReader = PascalVocReader('test.xml')
+# print tempParseReader.getShapes()
 """
 # Test
 tmp = PascalVocWriter('temp','test', (10,20,3))