將檢測的物體轉換成xml文件

# -*- coding: utf-8 -*-
import re
import shutil
import string
import numpy as np

from skimage import io

headstr = """\
<annotation>
    <folder>VOC2007</folder>
    <filename>%06d.jpg</filename>
    <source>
        <database>My Database</database>
        <annotation>PASCAL VOC2007</annotation>
        <image>flickr</image>
        <flickrid>NULL</flickrid>
    </source>
    <owner>
        <flickrid>NULL</flickrid>
        <name>company</name>
    </owner>
    <size>
        <width>%d</width>
        <height>%d</height>
        <depth>%d</depth>
    </size>
    <segmented>0</segmented>
"""
objstr = """\
    <object>
        <name>%s</name>
        <pose>Unspecified</pose>
        <truncated>0</truncated>
        <difficult>0</difficult>
        <bndbox>
            <xmin>%d</xmin>
            <ymin>%d</ymin>
            <xmax>%d</xmax>
            <ymax>%d</ymax>
        </bndbox>
    </object>
"""
tailstr = '''
</annotation>
'''

txt_dir = '/new/test1.txt'

xml_dir = '/new/xml/image8/'

img_dir = '/new/piuture8/'

labelname = 'person'


def clear_dir():
    if shutil.os.path.exists(xml_dir + 'Annotations'):
        shutil.rmtree(xml_dir + 'Annotations')

    shutil.os.mkdir(xml_dir + 'Annotations')


regex = re.compile('\s+')

def excute_datasets(file_list):
    #f = open(voc_dir + 'ImageSets/Main/' + datatype + '.txt', 'a')
    for anno_info in file_list:
        anno_info_arr = regex.split(string.strip(anno_info))
        filename = anno_info_arr[0]
        im = io.imread(img_dir+filename)
        idx = int(filename[:-4])
        head = headstr % (idx, im.shape[1], im.shape[0], im.shape[2])

        boxes = np.array(anno_info_arr[2:]).reshape(-1,4)
        tmp = []
        for bbx_info in boxes:
            bbx = [string.atoi(bbx_info[i]) for i in range(len(bbx_info))]
            tmp.append(bbx)
        writexml(idx, head, tmp, tailstr)

        idx += 1




def writexml(idx, head, bbxes, trailstr):
    filename = xml_dir + 'Annotations/%06d.xml' % (idx)
    f = open(filename, 'w')
    f.write(head)
    for bbx in bbxes:
        f.write(objstr % (labelname, bbx[0], bbx[1], bbx[2], bbx[3]))
    f.write(trailstr)
    f.close()


if __name__ == '__main__':
    clear_dir()

    file = open(txt_dir)
    file_list = file.readlines()
    excute_datasets(file_list)