inteligentny-traktor/decisionTree/treemaker.py

# -*- coding: utf-8 -*-


# from sklearn.datasets import load_iris
from sklearn.tree import export_text

from sklearn.tree import DecisionTreeClassifier
import joblib

X1 = []
view = []
with open("database.txt", 'r') as f:
    for line in f:
        line = line.strip()
        test_list = [int(i) for i in line]
        x = []
        if line[0] == "0":
            x.append("clear")
        else:
            x.append("rainy")
        if line[1] == "0":
            x.append("not planted")
        else:
            x.append("planted")
        if line[2] == "0":
            x.append("bad")
        else:
            x.append("good")
        if line[3] == "0":
            x.append("good")
        else:
            x.append("too strong")
        if line[4] == "0":
            x.append("no snow")
        else:
            x.append("snow")
        if line[5] == "0":
            x.append("dry")
        else:
            x.append("moist")
        if line[6] == "0":
            x.append("healthy")
        else:
            x.append("rot")
        if line[7] == "0":
            x.append("morning")
        elif line[7] == "1":
            x.append("noon")
        elif line[7] == "2":
            x.append("sunset")
        else:
            x.append("night")
        view.append(x)
        X1.append(test_list)

f = open("learning_set.txt", "w")              # zapisuje atrybuty s<>ownie
for i in view:
    f.write(str(i)+"\n")
f.close()


Y1 = []
with open("decissions.txt", 'r') as f:         # czyta decyzje
    for line in f:
        line = line.strip()
        test = int(line)
        Y1.append(test)

dataset = X1
decision = Y1
labels = ['Rain', 'Plant', 'Temperature', 'Sun', 'Snow', 'Moisture', 'Rotten', 'Time']
model = DecisionTreeClassifier(random_state=0, max_depth=20).fit(dataset, decision)
filename = 'decisionTree.sav'
print("Model trained")
print("Decision tree:")
print(export_text(model, feature_names=labels))
joblib.dump(model, filename)