inteligentny-traktor/decisionTree/treemaker.py

77 lines
1.9 KiB
Python

#from sklearn.datasets import load_iris
from sklearn.tree import export_text
from sklearn.tree import DecisionTreeClassifier
import joblib
X1 = []
view = []
with open("decisionTree/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("decisionTree/learning_set.txt", "w") #zapisuje atrybuty s³ownie
for i in view:
f.write(str(i)+"\n")
f.close()
Y1 = []
with open("decisionTree/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/decisionTree.sav'
print("Model trained")
print("Decision tree:")
print(export_text(model, feature_names=labels))
joblib.dump(model, filename)