From 50e95cf3d5e8fc70fa7ec30072be6f5f1f93df57 Mon Sep 17 00:00:00 2001 From: iabdalkader Date: Thu, 9 Jul 2015 00:21:15 +0200 Subject: [PATCH] Cleanup the cascade generator * Add command line args * Add cascade name arg * Add cascade info arg --- usr/openmv-cascade.py | 210 +++++++++++++++++++++++++++--------------- 1 file changed, 137 insertions(+), 73 deletions(-) diff --git a/usr/openmv-cascade.py b/usr/openmv-cascade.py index 91573aed5..b11e21db1 100755 --- a/usr/openmv-cascade.py +++ b/usr/openmv-cascade.py @@ -1,96 +1,160 @@ -#!/usr/bin/env python +#!/usr/bin/env python2.7 import sys,os import struct +import argparse from xml.dom import minidom -def print_help_exit(): - print "usage: cascade.py " - sys.exit(1) -if len(sys.argv)!= 3: - print_help_exit() +def cascade_info(path): + #parse xml file + xmldoc = minidom.parse(path) -n_stages = int(sys.argv[2]) -#parse xml file -xmldoc = minidom.parse(sys.argv[1]) + trees = xmldoc.getElementsByTagName('trees') + n_stages = len(trees) -trees = xmldoc.getElementsByTagName('trees') -max_stages = len(trees) -if n_stages > max_stages: - raise Exception("The max number of stages is: %d"%(max_stages)) + # read stages + stages = [len(t.childNodes)/2 for t in trees][0:n_stages] + stage_threshold = xmldoc.getElementsByTagName('stage_threshold')[0:n_stages] -# read stages -stages = [len(t.childNodes)/2 for t in trees][0:n_stages] -stage_threshold = xmldoc.getElementsByTagName('stage_threshold')[0:n_stages] + # total number of features + n_features = sum(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 -n_features = sum(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] + #read rectangles + feature = xmldoc.getElementsByTagName('rects')[0:n_features] -#read rectangles -feature = xmldoc.getElementsByTagName('rects')[0:n_features] + #read cascade size + size = (map(int, xmldoc.getElementsByTagName('size')[0].childNodes[0].nodeValue.split())) + fout = open(os.path.basename(path).split('.')[0]+".cascade", "w") -#read cascade size -size = (map(int, xmldoc.getElementsByTagName('size')[0].childNodes[0].nodeValue.split())) -fout = open(os.path.basename(sys.argv[1]).split('.')[0]+".cascade", "w") + n_rectangles = 0 + for f in feature: + rects = f.getElementsByTagName('_') + n_rectangles = n_rectangles + len(rects) -n_rectangles = 0 -for f in feature: - rects = f.getElementsByTagName('_') - n_rectangles = n_rectangles + len(rects) + #print some 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 some 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 +def cascade_binary(path, n_stages, name): + #parse xml file + xmldoc = minidom.parse(path) -# write detection window size -fout.write(struct.pack('i', size[0])) -fout.write(struct.pack('i', size[1])) + trees = xmldoc.getElementsByTagName('trees') + max_stages = len(trees) + if n_stages > max_stages: + raise Exception("The max number of stages is: %d"%(max_stages)) -# write num stages -fout.write(struct.pack('i', len(stages))) + if n_stages == 0: + n_stages = max_stages -# write num feat in stages -for s in stages: - fout.write(struct.pack('B', s)) # uint8_t + # read stages + stages = [len(t.childNodes)/2 for t in trees][0:n_stages] + stage_threshold = xmldoc.getElementsByTagName('stage_threshold')[0:n_stages] -# write stages thresholds -for t in stage_threshold: - fout.write(struct.pack('h', int(float(t.childNodes[0].nodeValue)*256))) #int16_t + # total number of features + n_features = sum(stages) -# write features threshold 1 per feature -for t in threshold: - fout.write(struct.pack('h', int(float(t.childNodes[0].nodeValue)*4096))) #int16_t + # read features threshold + threshold = xmldoc.getElementsByTagName('threshold')[0:n_features] -# write alpha1 1 per feature -for a in alpha1: - fout.write(struct.pack('h', int(float(a.childNodes[0].nodeValue)*256))) #int16_t + # theres one of each per feature + alpha1 = xmldoc.getElementsByTagName('left_val')[0:n_features] + alpha2 = xmldoc.getElementsByTagName('right_val')[0:n_features] -# write alpha2 1 per feature -for a in alpha2: - fout.write(struct.pack('h', int(float(a.childNodes[0].nodeValue)*256))) #int16_t + # read rectangles + feature = xmldoc.getElementsByTagName('rects')[0:n_features] -# write num_rects per feature -for f in feature: - rects = f.getElementsByTagName('_') - fout.write(struct.pack('B', len(rects))) # uint8_t + # read cascade size + size = (map(int, xmldoc.getElementsByTagName('size')[0].childNodes[0].nodeValue.split())) -# 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 + # open output file with the specified name or xml file name + if not name: + name = os.path.basename(path).split('.')[0] + fout = open(name+".cascade", "w") -# 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 + n_rectangles = 0 + for f in feature: + rects = f.getElementsByTagName('_') + n_rectangles = n_rectangles + len(rects) + + # 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()