Projekt_AI-Automatyczny_saper/Engine/DecisionTree.py

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\\Gabriel\\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('C:\\Users\\Gabriel\\Projekt_AI-Automatyczny_saper\\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)