From 2708b4bbb89ab13412260f27ecd7b26bdf525d5b Mon Sep 17 00:00:00 2001 From: ChrisDal Date: Tue, 25 Apr 2017 09:39:15 +0200 Subject: [PATCH] Add Truncated - difficult --- libs/pascal_voc_io.py | 100 ++++++++++++++++++++---------------------- 1 file changed, 48 insertions(+), 52 deletions(-) diff --git a/libs/pascal_voc_io.py b/libs/pascal_voc_io.py index 043f7b20..4f0e35bc 100644 --- a/libs/pascal_voc_io.py +++ b/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 @@ XML_EXT = '.xml' 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 @@ class PascalVocWriter: 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,26 +40,26 @@ class PascalVocWriter: 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: @@ -65,38 +67,46 @@ class PascalVocWriter: 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 @@ class PascalVocReader: 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 @@ class PascalVocReader: 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() -"""