You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
99 lines
2.5 KiB
Python
99 lines
2.5 KiB
Python
indent_prefix = 2 * ' '
|
|
|
|
|
|
def indent_line(line):
|
|
return indent_prefix + line
|
|
|
|
|
|
def encode_branch(branch):
|
|
if isinstance(branch, dict):
|
|
return encode_tree(branch)
|
|
elif isinstance(branch, (int, float)):
|
|
return [ indent_line('return {};'.format(branch)) ]
|
|
else:
|
|
return [ indent_line('return "{}";'.format(branch)) ]
|
|
|
|
# if value at index is smaller than threshold
|
|
# enter left, else right
|
|
|
|
|
|
def encode_tree(tree):
|
|
lines = [
|
|
"print(\"Considering feature {index}, is it smaller than {threshold}?\");".format(index=tree['index'], threshold=tree['value']),
|
|
"// timeout code",
|
|
"if (r[{index}] < {threshold}) {{".format(index=tree['index'], threshold=tree['value']),
|
|
indent_line("print(\"Yes\");"),
|
|
*encode_branch(tree['left']),
|
|
"}",
|
|
"else {",
|
|
indent_line("print(\"No\");"),
|
|
*encode_branch(tree['right']),
|
|
"}"
|
|
]
|
|
|
|
return map(indent_line, lines)
|
|
|
|
|
|
def make_classifier (tree):
|
|
lines = [
|
|
"#pragma once",
|
|
"#include <cstdarg>",
|
|
"namespace PublishingHouse",
|
|
"{",
|
|
*map(indent_line, [
|
|
"namespace RandomForest",
|
|
"{",
|
|
*map(indent_line, [
|
|
"class DecisionTree",
|
|
"{",
|
|
"public:",
|
|
*map(indent_line, [
|
|
"char* predict(float *r)",
|
|
"{",
|
|
*encode_tree(tree),
|
|
"}",
|
|
]),
|
|
"private:",
|
|
"};"
|
|
]),
|
|
"}"
|
|
]),
|
|
"}",
|
|
]
|
|
|
|
return('\n'.join(lines))
|
|
|
|
if __name__ == '__main__':
|
|
import json
|
|
import os.path
|
|
import glob
|
|
basepath = os.path.dirname(os.path.realpath(__file__))
|
|
globpath = os.path.realpath(os.path.join(basepath, 'random_forest_model_*-trees.json'))
|
|
|
|
models = glob.glob(globpath)
|
|
|
|
# Search for exported models
|
|
print("Found:")
|
|
for k, modelpath in enumerate(models):
|
|
print("[{}] {}".format(k, modelpath))
|
|
|
|
model_key = int(input("Which model?\n"))
|
|
|
|
modelpath = models[model_key]
|
|
|
|
# Open model
|
|
with open(modelpath, 'r') as file_in:
|
|
# Parse the forest
|
|
forest = json.load(file_in)
|
|
modelname, _ = os.path.splitext(os.path.basename(modelpath))
|
|
classifiernamepattern = 'classifier_{}_tree_{{}}.h'.format(modelname)
|
|
|
|
# Walk through the forest and visualize the trees
|
|
for idx, tree in enumerate(forest):
|
|
print('Transforming tree {} of {}'.format(idx, len(forest)))
|
|
code = make_classifier(tree)
|
|
|
|
with open(os.path.join(basepath, classifiernamepattern.format(idx)), 'w') as h:
|
|
h.write(code)
|
|
|
|
print('Classifiers placed in: {}'.format(basepath)) |