File size: 2,247 Bytes
9c29da9 |
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 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
# -*- coding: utf-8 -*-
from xml.dom import minidom
import os
import glob
lut={}
lut["pikachu"]=0
lut["charmander"]=1
lut["bulbasaur"]=2
lut["squirtle"]=3
lut["eevee"]=4
lut["other"]=5
lut["jigglypuff"]=6
def convert_coordinates(size, box):
if(size[0]==0 or size[1]==0):
return (0,0,0,0)
dw = 1.0/size[0]
dh = 1.0/size[1]
x = (box[0]+box[1])/2.0
y = (box[2]+box[3])/2.0
w = box[1]-box[0]
h = box[3]-box[2]
x = x*dw
w = w*dw
y = y*dh
h = h*dh
return (x,y,w,h)
def convert_xml2yolo(lut):
for fname in glob.glob("*.xml"):
xmldoc = minidom.parse(fname)
fname_out = (fname[:-4]+'.txt')
with open(fname_out, "w") as f:
itemlist = xmldoc.getElementsByTagName('object')
size = xmldoc.getElementsByTagName('size')[0]
width = int((size.getElementsByTagName('width')[0]).firstChild.data)
height = int((size.getElementsByTagName('height')[0]).firstChild.data)
for item in itemlist:
# get class label
classid = (item.getElementsByTagName('name')[0]).firstChild.data
if classid in lut:
label_str = str(lut[classid])
else:
label_str = "-1"
print ("warning: label '%s' not in look-up table" % classid)
# get bbox coordinates
xmin = ((item.getElementsByTagName('bndbox')[0]).getElementsByTagName('xmin')[0]).firstChild.data
ymin = ((item.getElementsByTagName('bndbox')[0]).getElementsByTagName('ymin')[0]).firstChild.data
xmax = ((item.getElementsByTagName('bndbox')[0]).getElementsByTagName('xmax')[0]).firstChild.data
ymax = ((item.getElementsByTagName('bndbox')[0]).getElementsByTagName('ymax')[0]).firstChild.data
b = (float(xmin), float(xmax), float(ymin), float(ymax))
bb = convert_coordinates((width,height), b)
#print(bb)
f.write(label_str + " " + " ".join([("%.6f" % a) for a in bb]) + '\n')
print ("wrote %s" % fname_out)
def main():
convert_xml2yolo(lut)
if __name__ == '__main__':
main() |