import pandas as pd import json from sklearn import tree import pydotplus from sklearn.tree import DecisionTreeClassifier import matplotlib.pyplot as plt import matplotlib.image as pltimg import pickle class DecisionTree: def __init__(self, doCreation): self.data = pd.read_csv("C:\\Users\\kratu\\PycharmProjects\\Projekt_AI-Automatyczny_saper\\out.csv") if doCreation: self.mapData() features = ['bomb_type', 'detonation_duration', 'size', 'detonation_area', 'defusable'] X = self.data[features] y = self.data['action'] dtree = DecisionTreeClassifier() dtree = dtree.fit(X, y) decision_tree_model_pkl = open('tree.pkl', 'wb') pickle.dump(dtree, decision_tree_model_pkl) decision_tree_model_pkl.close() decision_tree_model_pkl = open('Engine/tree.pkl', 'rb') self.dtree = pickle.load(decision_tree_model_pkl) def getTree(self): return self.dtree def mapData(self): d = {'Atomic Bomb': 0, 'Claymore': 1, 'Land Mine': 2, 'Chemical Bomb': 3, 'Decoy': 4} self.data['bomb_type'] = self.data['bomb_type'].map(d) d = {'immediate': 0, 'short': 1, 'long': 2, 'none': 3} self.data['detonation_duration'] = self.data['detonation_duration'].map(d) d = {'small': 0, 'medium': 1, 'large': 2} self.data['size'] = self.data['size'].map(d) d = {'small': 0, 'large': 1} self.data['detonation_area'] = self.data['detonation_area'].map(d) d = {'no': 0, 'yes': 1} self.data['defusable'] = self.data['defusable'].map(d) d = {'detonate': 0, 'poligon': 1, 'defuse': 2} self.data['action'] = self.data['action'].map(d) def mapAction(self, action): d = {0 : 'detonate', 1 : 'poligon', 2 : 'defuse'} return d.get(action) if __name__ == "__main__": # data = pd.read_csv("C:\\Users\\kratu\\PycharmProjects\\Projekt_AI-Automatyczny_saper\\out.csv") DecisionTree(True) # with open('C:\\Users\\kratu\\PycharmProjects\\Projekt_AI-Automatyczny_saper\\DecisionTree.json', 'w') as fp: # json.dump(tree, fp)