deep_lipid/treeviz.py at master · sheneman/deep_lipid · GitHub

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
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
###############################################################################
# treeviz.py
#
# sheneman@uidaho.edu
#
# A silly tool for taking a single trained decision tree from a random forest
# model and visualizing it using the GraphViz library.
#
# Usage:
# python treeviz.py [ --help | --verbose | --config=<YAML config file> ]
#
################################################################################

from PIL import Image, ImageFilter
import sys
import getopt
import yaml
import pickle
from os.path import isfile, join
import numpy
import pprint
from skimage import feature
from sklearn.ensemble import RandomForestClassifier
from sklearn.tree import export_graphviz
from subprocess import call

import preprocess


numpy.set_printoptions(threshold=sys.maxsize)


#################################################################################
#
# HANDLE Command line arguments
#
#
def usage():
	print("python treeviz.py [ --help | --verbose | --config=<YAML config filename> ] ")

try:
	opts, args = getopt.getopt(sys.argv[1:], "ho:v", ["help", "config="])
except getopt.GetoptError as err:
	print(err)  # will print something like "option -a not recognized"
	usage()
	sys.exit(2)

configfile = None
verbose = False

for o, a in opts:
	if o == "-v":
		verbose = True
	elif o in ("-h", "--help"):
		usage()
		sys.exit()
	elif o in ("-c", "--config"):
		configfile = a
	else:
		assert False, "unhandled option"

if(configfile == None):
	configfile="treeviz.yaml"

#################################################################################

#################################################################################
#
# Format and Example config YAML file:
#
# FORMAT:
# -------
#   model:	<path to the model to use for classification>
#   output:	<path to output image>
#
# EXAMPLE:
# --------
#   model:	"./models/foo.model"
#   output:	"./tree.png"
#
#################################################################################

cf = open(configfile, "r")
config = yaml.load(cf, Loader=yaml.FullLoader)
print("YAML CONFIG:")
for c in config:
	print("    [%s]:\"%s\"" %(c,config[c]))
print("\n")
cf.close()

# Load the classifier model
print("Loading classifier model = %s" %config["model"])
classifier = pickle.load(open(config["model"],'rb'))
print("Model %s Loaded!" %config["model"])

# Get the feature labels
feature_labels = preprocess.feature_labels()
print(feature_labels)

tree = classifier.estimators_[0]
export_graphviz(tree, out_file='./tree.dot',
                feature_names = feature_labels,
                class_names = ("Non-Lipid", "Lipid"),
                rounded = True, proportion = False,
                precision = 2, filled = True)


#call(['dot', '-Tpng', '/tmp/tree.dot', '-o', config["output"], '-Gdpi=600'])
print("Done!")

exit(0)