Cleanup the cascade generator

* Add command line args
* Add cascade name arg
* Add cascade info arg
This commit is contained in:
iabdalkader 2015-07-09 00:21:15 +02:00
parent ff611f3706
commit 50e95cf3d5

View File

@ -1,96 +1,160 @@
#!/usr/bin/env python #!/usr/bin/env python2.7
import sys,os import sys,os
import struct import struct
import argparse
from xml.dom import minidom from xml.dom import minidom
def print_help_exit():
print "usage: cascade.py <cv_haar_cascade.xml> <num stages>"
sys.exit(1)
if len(sys.argv)!= 3: def cascade_info(path):
print_help_exit() #parse xml file
xmldoc = minidom.parse(path)
n_stages = int(sys.argv[2]) trees = xmldoc.getElementsByTagName('trees')
#parse xml file n_stages = len(trees)
xmldoc = minidom.parse(sys.argv[1])
trees = xmldoc.getElementsByTagName('trees') # read stages
max_stages = len(trees) stages = [len(t.childNodes)/2 for t in trees][0:n_stages]
if n_stages > max_stages: stage_threshold = xmldoc.getElementsByTagName('stage_threshold')[0:n_stages]
raise Exception("The max number of stages is: %d"%(max_stages))
# read stages # total number of features
stages = [len(t.childNodes)/2 for t in trees][0:n_stages] n_features = sum(stages)
stage_threshold = xmldoc.getElementsByTagName('stage_threshold')[0:n_stages] # read features threshold
threshold = xmldoc.getElementsByTagName('threshold')[0:n_features]
#theres one of each per feature
alpha1 = xmldoc.getElementsByTagName('left_val')[0:n_features]
alpha2 = xmldoc.getElementsByTagName('right_val')[0:n_features]
# total number of features #read rectangles
n_features = sum(stages) feature = xmldoc.getElementsByTagName('rects')[0:n_features]
# read features threshold
threshold = xmldoc.getElementsByTagName('threshold')[0:n_features]
#theres one of each per feature
alpha1 = xmldoc.getElementsByTagName('left_val')[0:n_features]
alpha2 = xmldoc.getElementsByTagName('right_val')[0:n_features]
#read rectangles #read cascade size
feature = xmldoc.getElementsByTagName('rects')[0:n_features] size = (map(int, xmldoc.getElementsByTagName('size')[0].childNodes[0].nodeValue.split()))
fout = open(os.path.basename(path).split('.')[0]+".cascade", "w")
#read cascade size n_rectangles = 0
size = (map(int, xmldoc.getElementsByTagName('size')[0].childNodes[0].nodeValue.split())) for f in feature:
fout = open(os.path.basename(sys.argv[1]).split('.')[0]+".cascade", "w") rects = f.getElementsByTagName('_')
n_rectangles = n_rectangles + len(rects)
n_rectangles = 0 #print some cascade info
for f in feature: print("size:%dx%d"%(size[0], size[1]))
rects = f.getElementsByTagName('_') print("stages:%d"%len(stages))
n_rectangles = n_rectangles + len(rects) print("features:%d"%n_features)
print("rectangles:%d"%n_rectangles)
#print some cascade info def cascade_binary(path, n_stages, name):
print "size:%dx%d"%(size[0], size[1]) #parse xml file
print "stages:%d"%len(stages) xmldoc = minidom.parse(path)
print "features:%d"%n_features
print "rectangles:%d"%n_rectangles
# write detection window size trees = xmldoc.getElementsByTagName('trees')
fout.write(struct.pack('i', size[0])) max_stages = len(trees)
fout.write(struct.pack('i', size[1])) if n_stages > max_stages:
raise Exception("The max number of stages is: %d"%(max_stages))
# write num stages if n_stages == 0:
fout.write(struct.pack('i', len(stages))) n_stages = max_stages
# write num feat in stages # read stages
for s in stages: stages = [len(t.childNodes)/2 for t in trees][0:n_stages]
fout.write(struct.pack('B', s)) # uint8_t stage_threshold = xmldoc.getElementsByTagName('stage_threshold')[0:n_stages]
# write stages thresholds # total number of features
for t in stage_threshold: n_features = sum(stages)
fout.write(struct.pack('h', int(float(t.childNodes[0].nodeValue)*256))) #int16_t
# write features threshold 1 per feature # read features threshold
for t in threshold: threshold = xmldoc.getElementsByTagName('threshold')[0:n_features]
fout.write(struct.pack('h', int(float(t.childNodes[0].nodeValue)*4096))) #int16_t
# write alpha1 1 per feature # theres one of each per feature
for a in alpha1: alpha1 = xmldoc.getElementsByTagName('left_val')[0:n_features]
fout.write(struct.pack('h', int(float(a.childNodes[0].nodeValue)*256))) #int16_t alpha2 = xmldoc.getElementsByTagName('right_val')[0:n_features]
# write alpha2 1 per feature # read rectangles
for a in alpha2: feature = xmldoc.getElementsByTagName('rects')[0:n_features]
fout.write(struct.pack('h', int(float(a.childNodes[0].nodeValue)*256))) #int16_t
# write num_rects per feature # read cascade size
for f in feature: size = (map(int, xmldoc.getElementsByTagName('size')[0].childNodes[0].nodeValue.split()))
rects = f.getElementsByTagName('_')
fout.write(struct.pack('B', len(rects))) # uint8_t
# write rects weights 1 per rectangle # open output file with the specified name or xml file name
for f in feature: if not name:
rects = f.getElementsByTagName('_') name = os.path.basename(path).split('.')[0]
for r in rects: fout = open(name+".cascade", "w")
l = map(int, r.childNodes[0].nodeValue[:-1].split())
fout.write(struct.pack('b', l[4])) #int8_t NOTE: multiply by 4096
# write rects n_rectangles = 0
for f in feature: for f in feature:
rects = f.getElementsByTagName('_') rects = f.getElementsByTagName('_')
for r in rects: n_rectangles = n_rectangles + len(rects)
l = map(int, r.childNodes[0].nodeValue[:-1].split())
fout.write(struct.pack('BBBB',l[0], l[1], l[2], l[3])) #uint8_t # write detection window size
fout.write(struct.pack('i', size[0]))
fout.write(struct.pack('i', size[1]))
# write num stages
fout.write(struct.pack('i', len(stages)))
# write num feat in stages
for s in stages:
fout.write(struct.pack('B', s)) # uint8_t
# write stages thresholds
for t in stage_threshold:
fout.write(struct.pack('h', int(float(t.childNodes[0].nodeValue)*256))) #int16_t
# write features threshold 1 per feature
for t in threshold:
fout.write(struct.pack('h', int(float(t.childNodes[0].nodeValue)*4096))) #int16_t
# write alpha1 1 per feature
for a in alpha1:
fout.write(struct.pack('h', int(float(a.childNodes[0].nodeValue)*256))) #int16_t
# write alpha2 1 per feature
for a in alpha2:
fout.write(struct.pack('h', int(float(a.childNodes[0].nodeValue)*256))) #int16_t
# write num_rects per feature
for f in feature:
rects = f.getElementsByTagName('_')
fout.write(struct.pack('B', len(rects))) # uint8_t
# write rects weights 1 per rectangle
for f in feature:
rects = f.getElementsByTagName('_')
for r in rects:
l = map(int, r.childNodes[0].nodeValue[:-1].split())
fout.write(struct.pack('b', l[4])) #int8_t NOTE: multiply by 4096
# write rects
for f in feature:
rects = f.getElementsByTagName('_')
for r in rects:
l = map(int, r.childNodes[0].nodeValue[:-1].split())
fout.write(struct.pack('BBBB',l[0], l[1], l[2], l[3])) #uint8_t
# print cascade info
print("size:%dx%d"%(size[0], size[1]))
print("stages:%d"%len(stages))
print("features:%d"%n_features)
print("rectangles:%d"%n_rectangles)
print("binary cascade generated")
def main():
# CMD args parser
parser = argparse.ArgumentParser(description='haar cascade generator')
parser.add_argument("-i", "--info", action = "store_true", help = "print cascade info and exit")
parser.add_argument("-n", "--name", action = "store", help = "set cascade name", default = "")
parser.add_argument("-s", "--stages", action = "store", help = "set the maximum number of stages", type = int, default=0)
parser.add_argument("file", action = "store", help = "OpenCV xml cascade file path")
# Parse CMD args
args = parser.parse_args()
if args.info:
# print cascade info and exit
cascade_info(args.file)
return
# generate a binary cascade from the xml cascade
cascade_binary(args.file, args.stages, args.name)
if __name__ == '__main__':
main()