生成VOC2007标注格式

在目标检测、文本检测人脸检测的项目时候，会有一些标注给的是txt文件，而最经济实用的方法是，将txt标注文件转成VOC格式（XML），否则自己实现一个数据输入层。下面贴出转化的代码：
  1 import cv2
  2 from xml.dom.minidom import Document
  3 from collections import defaultdict
  4 
  5 src_img_dir = 'VOCdevkit/VOC2012/JPEGImages/'
  6 src_txt_dir = 'train.txt'
  7 src_xml_dir = 'xml/'
  8 
  9 img_labels = defaultdict(list)
 10 f = open(src_txt_dir, 'r')
 11 for i in open(src_txt_dir):
 12     line = f.readline().strip('\n')
 13     x = line.split('.jpg ')
 14     key = x[0]
 15     value = x[1]
 16     img_labels[key].append(value)
 17 
 18 for key in img_labels:
 19     if(len(img_labels[key])>1):
 20         #print (key)
 21         #print (img_labels[key])
 22         pass
 23 
 24     doc = Document()
 25     annotation = doc.createElement('annotation')
 26     doc.appendChild(annotation)
 27 
 28     folder = doc.createElement('folder')
 29     folder.appendChild(doc.createTextNode('VOC2012'))
 30     annotation.appendChild(folder)
 31 
 32     filename = doc.createElement('filename')
 33     filename.appendChild(doc.createTextNode(key+'.jpg'))
 34     annotation.appendChild(filename)
 35 
 36     source = doc.createElement('source')
 37     database = doc.createElement('database')
 38     database.appendChild(doc.createTextNode('Baidu'))
 39     source.appendChild(database)
 40     annotation.appendChild(source)
 41 
 42     # obtain the size of the image
 43     imagefile = src_img_dir + key + '.jpg'
 44     img = cv2.imread(imagefile)
 45     imgSize = img.shape
 46 
 47     size = doc.createElement('size')
 48     width = doc.createElement('width')
 49     width.appendChild(doc.createTextNode(str(imgSize[1])))
 50     size.appendChild(width)
 51     height = doc.createElement('height')
 52     height.appendChild(doc.createTextNode(str(imgSize[0])))
 53     size.appendChild(height)
 54     depth = doc.createElement('depth')
 55     depth.appendChild(doc.createTextNode(str(imgSize[2])))
 56     size.appendChild(depth)
 57     annotation.appendChild(size)
 58 
 59     segmented = doc.createElement('segmented')
 60     segmented.appendChild(doc.createTextNode(str(0)))
 61     annotation.appendChild(segmented)
 62 
 63     # write the coordinates of the b-box
 64     for b_box in list(img_labels[key]):
 65         #print b_box
 66         x = b_box.split(' ')
 67         if(int(x[1])<0):
 68             x[1] = '0'
 69         if (int(x[2]) < 0):
 70             x[2] = '0'
 71 
 72         object = doc.createElement('object')
 73         name = doc.createElement('name')
 74         name.appendChild(doc.createTextNode('n' + x[0]))
 75         #name.appendChild(doc.createTextNode(x[0]))
 76         object.appendChild(name)
 77 
 78         difficult = doc.createElement('difficult')
 79         difficult.appendChild(doc.createTextNode('0'))
 80         object.appendChild(difficult)
 81 
 82         truncated = doc.createElement('truncated')
 83         truncated.appendChild(doc.createTextNode('0'))
 84         object.appendChild(truncated)
 85 
 86         pose = doc.createElement('pose')
 87         pose.appendChild(doc.createTextNode('undefined'))
 88         object.appendChild(pose)
 89 
 90         bndbox = doc.createElement('bndbox')
 91         xmin = doc.createElement('xmin')
 92         xmin.appendChild(doc.createTextNode(x[1]))
 93         bndbox.appendChild(xmin)
 94         object.appendChild(bndbox)
 95         ymin = doc.createElement('ymin')
 96         ymin.appendChild(doc.createTextNode(x[2]))
 97         bndbox.appendChild(ymin)
 98         xmax = doc.createElement('xmax')
 99         xmax.appendChild(doc.createTextNode(x[3]))
100         bndbox.appendChild(xmax)
101         ymax = doc.createElement('ymax')
102         ymax.appendChild(doc.createTextNode(x[4]))
103         bndbox.appendChild(ymax)
104         annotation.appendChild(object)
105 
106     with open(src_xml_dir + key + '.xml', 'wb') as f:
107         f.write(doc.toprettyxml(indent='\t', encoding='utf-8'))
生成VOC2007标注格式

猜你喜欢
