forked from HouchangX-AI/Helmet-Detection
-
Notifications
You must be signed in to change notification settings - Fork 2
/
voc_annotation.py
34 lines (26 loc) · 1.19 KB
/
voc_annotation.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
import xml.etree.ElementTree as ET
from os import getcwd
sets=['train', 'val', 'test']
classes = ["person", "hat"]
def convert_annotation( image_id, list_file):
in_file = open('/Users/huike/Downloads/Faster-R-CNN-Keras-master/data/Annotations/%s.xml'%(image_id))
tree=ET.parse(in_file)
root = tree.getroot()
for obj in root.iter('object'):
difficult = obj.find('difficult').text
cls = obj.find('name').text
if cls not in classes: # or int(difficult)==1:
continue
cls_id = classes.index(cls)
xmlbox = obj.find('bndbox')
b = (int(xmlbox.find('xmin').text), int(xmlbox.find('ymin').text), int(xmlbox.find('xmax').text), int(xmlbox.find('ymax').text))
list_file.write(" " + ",".join([str(a) for a in b]) + ',' + str(cls_id))
wd = getcwd()
for image_set in sets:
image_ids = open('/Users/huike/Downloads/Faster-R-CNN-Keras-master/data/ImageSets/Main/%s.txt'%(image_set)).read().strip().split()
list_file = open('%s.txt'%( image_set), 'w')
for image_id in image_ids:
list_file.write('JPEGImages/%s.jpg'%( image_id))
convert_annotation(image_id, list_file)
list_file.write('\n')
list_file.close()