diff --git a/usr/openmv-cascade.py b/usr/openmv-cascade.py index 1378f4a4b..66955944d 100755 --- a/usr/openmv-cascade.py +++ b/usr/openmv-cascade.py @@ -4,11 +4,54 @@ import struct import argparse from xml.dom import minidom +def cascade_info_universal(path): + xmldoc = minidom.parse(path) + old_format = xmldoc.getElementsByTagName('stageNum').length == 0 + if old_format: + print("Parsing old XML format..") + cascade_info_old(path) + else: + print("Parsing new XML format..") + cascade_info(path) def cascade_info(path): #parse xml file xmldoc = minidom.parse(path) + n_stages = int(xmldoc.getElementsByTagName('stageNum')[0].childNodes[0].nodeValue) + + # read stages + stages_elements = xmldoc.getElementsByTagName('stages') + stages = [] + for node in stages_elements[0].childNodes: + if node.nodeType is 1: + stages.append(int(node.getElementsByTagName('maxWeakCount')[0].childNodes[0].nodeValue)) + stage_threshold = xmldoc.getElementsByTagName('stageThreshold')[0:n_stages] + + # total number of features + n_features = sum(stages) + + #read rectangles + feature = xmldoc.getElementsByTagName('rects')[0:n_features] + + #read cascade size + size = [int(xmldoc.getElementsByTagName('width')[0].childNodes[0].nodeValue), int(xmldoc.getElementsByTagName('height')[0].childNodes[0].nodeValue)] + + 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) + +def cascade_info_old(path): + #parse xml file + xmldoc = minidom.parse(path) + trees = xmldoc.getElementsByTagName('trees') n_stages = len(trees) @@ -29,7 +72,6 @@ def cascade_info(path): #read cascade size size = (map(int, xmldoc.getElementsByTagName('size')[0].childNodes[0].nodeValue.split())) - fout = open(os.path.basename(path).split('.')[0]+".cascade", "w") n_rectangles = 0 for f in feature: @@ -42,10 +84,127 @@ def cascade_info(path): print("features:%d"%n_features) print("rectangles:%d"%n_rectangles) +def cascade_binary_universal(path, n_stages, name): + xmldoc = minidom.parse(path) + old_format = xmldoc.getElementsByTagName('stageNum').length == 0 + if old_format: + print("Converting old XML format..") + cascade_binary_old(path, n_stages, name) + else: + print("Converting new XML format..") + cascade_binary(path, n_stages, name) + def cascade_binary(path, n_stages, name): #parse xml file xmldoc = minidom.parse(path) + max_stages = int(xmldoc.getElementsByTagName('stageNum')[0].childNodes[0].nodeValue) + if n_stages > max_stages: + raise Exception("The max number of stages is: %d"%(max_stages)) + + if n_stages == 0: + n_stages = max_stages + + # read stages + stages_elements = xmldoc.getElementsByTagName('stages') + stages = [] + for node in stages_elements[0].childNodes: + if node.nodeType is 1: + stages.append(int(node.getElementsByTagName('maxWeakCount')[0].childNodes[0].nodeValue)) + stage_threshold = xmldoc.getElementsByTagName('stageThreshold')[0:n_stages] + + # total number of features + n_features = sum(stages) + + # read features threshold + internal_nodes = xmldoc.getElementsByTagName('internalNodes')[0:n_features] + + # theres one of each per feature + leaf_values = xmldoc.getElementsByTagName('leafValues')[0:n_features] + + alpha1 = [] + alpha2 = [] + for val in leaf_values: + alpha1.append(val.childNodes[0].nodeValue.split()[0]) + alpha2.append(val.childNodes[0].nodeValue.split()[1]) + + # read rectangles + feature = xmldoc.getElementsByTagName('rects')[0:n_features] + + # read cascade size + size = [int(xmldoc.getElementsByTagName('width')[0].childNodes[0].nodeValue), int(xmldoc.getElementsByTagName('height')[0].childNodes[0].nodeValue)] + + # 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") + + 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 internal_nodes: + fout.write(struct.pack('h', int(float(t.childNodes[0].nodeValue.split()[3])*4096))) #int16_t + + # write alpha1 1 per feature + for a in alpha1: + fout.write(struct.pack('h', int(float(a)*256))) #int16_t + + # write alpha2 1 per feature + for a in alpha2: + fout.write(struct.pack('h', int(float(a)*256))) #int16_t + + # write num_rects per feature + for f in internal_nodes: + idx = int(f.childNodes[0].nodeValue.split()[2]) + rects = feature[idx].getElementsByTagName('_') + fout.write(struct.pack('B', len(rects))) # uint8_t + + # write rects weights 1 per rectangle + for f in internal_nodes: + idx = int(f.childNodes[0].nodeValue.split()[2]) + rects = feature[idx].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 internal_nodes: + idx = int(f.childNodes[0].nodeValue.split()[2]) + rects = feature[idx].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 cascade_binary_old(path, n_stages, name): + #parse xml file + xmldoc = minidom.parse(path) + trees = xmldoc.getElementsByTagName('trees') max_stages = len(trees) if n_stages > max_stages: @@ -242,7 +401,7 @@ def main(): if args.info: # print cascade info and exit - cascade_info(args.file) + cascade_info_universal(args.file) return if args.header: @@ -251,7 +410,7 @@ def main(): return # generate a binary cascade from the xml cascade - cascade_binary(args.file, args.stages, args.name) + cascade_binary_universal(args.file, args.stages, args.name) if __name__ == '__main__': main()