diff --git a/main.py b/main.py index ef33b9f..7eb4d69 100644 --- a/main.py +++ b/main.py @@ -725,6 +725,41 @@ while True: actTake = True else: actTake = False + if event.key == pygame.K_3: + project = 0 + tickTime = 0 + import pandas as pd + from sklearn.tree import DecisionTreeClassifier + from sklearn.model_selection import train_test_split + from sklearn import metrics + + col_names = ['idle','first','good'] + data = pd.read_csv("results.csv", header=None, names=col_names) + data = data.iloc[1:] + data.head() + + feature_cols = ['idle','first'] + X = data[feature_cols] + y = data.good + + X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=1) + + clf = DecisionTreeClassifier(criterion="gini", max_depth=4) + + clf = clf.fit(X_train,y_train) + + from sklearn.externals.six import StringIO + from IPython.display import Image + from sklearn.tree import export_graphviz + import pydotplus + dot_data = StringIO() + export_graphviz(clf, out_file=dot_data, + filled=True, rounded=True, + special_characters=True, feature_names = feature_cols,class_names=['0','1']) + graph = pydotplus.graph_from_dot_data(dot_data.getvalue()) + graph.write_png('results.png') + Image(graph.create_png()) + print("graph done") if event.key == pygame.K_F4: pygame.quit() if event.key == pygame.K_F5: