<s> import sys <EOL> from xml . etree import ElementTree <EOL> from xml . etree . ElementTree import Element , SubElement <EOL> from xml . dom import minidom <EOL> from lxml import etree <EOL> class PascalVocWriter : <EOL> def __init__ ( self , foldername , filename , imgSize , databaseSrc = '<STR_LIT>' , localImgPath = None ) : <EOL> self . foldername = foldername <EOL> self . filename = filename <EOL> self . databaseSrc = databaseSrc <EOL> self . imgSize = imgSize <EOL> self . boxlist = [ ] <EOL> self . localImgPath = localImgPath <EOL> def prettify ( self , elem ) : <EOL> """<STR_LIT>""" <EOL> rough_string = ElementTree . tostring ( elem , '<STR_LIT:utf8>' ) <EOL> reparsed = minidom . parseString ( rough_string ) <EOL> return reparsed . toprettyxml ( indent = "<STR_LIT:\t>" ) <EOL> def genXML ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . filename is None or self . foldername is None or self . imgSize is None or len ( self . boxlist ) <= <NUM_LIT:0> : <EOL> return None <EOL> top = Element ( '<STR_LIT>' ) <EOL> folder = SubElement ( top , '<STR_LIT>' ) <EOL> folder . text = self . foldername <EOL> filename = SubElement ( top , '<STR_LIT:filename>' ) <EOL> filename . text = self . filename <EOL> localImgPath = SubElement ( top , '<STR_LIT:path>' ) <EOL> localImgPath . text = self . localImgPath <EOL> source = SubElement ( top , '<STR_LIT:source>' ) <EOL> database = SubElement ( source , '<STR_LIT>' ) <EOL> database . text = self . databaseSrc <EOL> size_part = SubElement ( top , '<STR_LIT:size>' ) <EOL> width = SubElement ( size_part , '<STR_LIT:width>' ) <EOL> height = SubElement ( size_part , '<STR_LIT>' ) <EOL> depth = SubElement ( size_part , '<STR_LIT>' ) <EOL> width . text = str ( self . imgSize [ <NUM_LIT:1> ] ) <EOL> height . text = str ( self . imgSize [ <NUM_LIT:0> ] ) <EOL> if len ( self . imgSize ) == <NUM_LIT:3> : <EOL> depth . text = str ( self . imgSize [ <NUM_LIT:2> ] ) <EOL> else : <EOL> depth . text = '<STR_LIT:1>' <EOL> segmented = SubElement ( top , '<STR_LIT>' ) <EOL> segmented . text = '<STR_LIT:0>' <EOL> return top <EOL> def addBndBox ( self , xmin , ymin , xmax , ymax , name ) : <EOL> bndbox = { '<STR_LIT>' : xmin , '<STR_LIT>' : ymin , '<STR_LIT>' : xmax , '<STR_LIT>' : ymax } <EOL> bndbox [ '<STR_LIT:name>' ] = name <EOL> self . boxlist . append ( bndbox ) ; <EOL> def appendObjects ( self , top ) : <EOL> for each_object in self . boxlist : <EOL> object_item = SubElement ( top , '<STR_LIT:object>' ) <EOL> name = SubElement ( object_item , '<STR_LIT:name>' ) <EOL> name . text = str ( each_object [ '<STR_LIT:name>' ] ) <EOL> pose = SubElement ( object_item , '<STR_LIT>' ) <EOL> pose . text = "<STR_LIT>" <EOL> truncated = SubElement ( object_item , '<STR_LIT>' ) <EOL> truncated . text = "<STR_LIT:0>" <EOL> difficult = SubElement ( object_item , '<STR_LIT>' ) <EOL> difficult . text = "<STR_LIT:0>" <EOL> bndbox = SubElement ( object_item , '<STR_LIT>' ) <EOL> xmin = SubElement ( bndbox , '<STR_LIT>' ) <EOL> xmin . text = str ( each_object [ '<STR_LIT>' ] ) <EOL> ymin = SubElement ( bndbox , '<STR_LIT>' ) <EOL> ymin . text = str ( each_object [ '<STR_LIT>' ] ) <EOL> xmax = SubElement ( bndbox , '<STR_LIT>' ) <EOL> xmax . text = str ( each_object [ '<STR_LIT>' ] ) <EOL> ymax = SubElement ( bndbox , '<STR_LIT>' ) <EOL> ymax . text = str ( each_object [ '<STR_LIT>' ] ) <EOL> def save ( self , targetFile = None ) : <EOL> root = self . genXML ( ) <EOL> self . appendObjects ( root ) <EOL> out_file = None <EOL> if targetFile is None : <EOL> out_file = open ( self . filename + '<STR_LIT>' , '<STR_LIT:w>' ) <EOL> else : <EOL> out_file = open ( targetFile , '<STR_LIT:w>' ) <EOL> out_file . write ( self . prettify ( root ) ) <EOL> out_file . close ( ) <EOL> class PascalVocReader : <EOL> def __init__ ( self , filepath ) : <EOL> self . shapes = [ ] <EOL> self . filepath = filepath <EOL> self . parseXML ( ) <EOL> def getShapes ( self ) : <EOL> return self . shapes <EOL> def addShape ( self , label , rect ) : <EOL> xmin = rect [ <NUM_LIT:0> ] <EOL> ymin = rect [ <NUM_LIT:1> ] <EOL> xmax = rect [ <NUM_LIT:2> ] <EOL> ymax = rect [ <NUM_LIT:3> ] <EOL> points = [ ( xmin , ymin ) , ( xmin , ymax ) , ( xmax , ymax ) , ( xmax , ymin ) ] <EOL> self . shapes . append ( ( label , points , None , None ) ) <EOL> def parseXML ( self ) : <EOL> assert self . filepath . endswith ( '<STR_LIT>' ) , "<STR_LIT>" <EOL> xmltree = ElementTree . parse ( self . filepath ) . getroot ( ) <EOL> filename = xmltree . find ( '<STR_LIT:filename>' ) . text <EOL> for object_iter in xmltree . findall ( '<STR_LIT:object>' ) : <EOL> rects = [ ] <EOL> bndbox = object_iter . find ( "<STR_LIT>" ) <EOL> rects . append ( [ int ( it . text ) for it in bndbox ] ) <EOL> label = object_iter . find ( '<STR_LIT:name>' ) . text <EOL> for rect in rects : <EOL> self . addShape ( label , rect ) <EOL> return True <EOL> """<STR_LIT>""" </s>
