2024-05-10 19:41:13 +02:00
from sklearn import tree as skltree
import pandas , os
import matplotlib . pyplot as plt
2024-05-11 14:18:40 +02:00
atributes = [ ' plant_water_level ' , ' growth ' , ' disease ' , ' fertility ' , ' tractor_water_level ' , ' temperature ' , ' rain ' , ' season ' , ' current_time ' ] #Columns in CSV file has to be in the same order
2024-05-10 19:41:13 +02:00
class Drzewo :
def __init__ ( self ) :
self . tree = self . treeLearn ( )
def treeLearn ( self ) :
csvdata = pandas . read_csv ( ' Data/dataTree.csv ' )
x = csvdata [ atributes ]
decision = csvdata [ ' action ' ]
self . tree = skltree . DecisionTreeClassifier ( )
2024-05-11 14:18:40 +02:00
self . tree = self . tree . fit ( x . values , decision )
2024-05-10 19:41:13 +02:00
def plotTree ( self ) :
plt . figure ( )
skltree . plot_tree ( self . tree , filled = True , feature_names = atributes )
plt . title ( " Drzewo decyzyjne wytrenowane na przygotowanych danych " )
2024-05-11 14:18:40 +02:00
plt . savefig ( ' tree.png ' )
2024-05-10 19:41:13 +02:00
plt . show ( )
def makeDecision ( self , values ) :
2024-05-11 14:18:40 +02:00
action = self . tree . predict ( [ values ] ) #0- nie podlewac, 1-podlewac
2024-05-11 14:30:29 +02:00
if ( action == [ 0 ] ) :
return " Nie podlewac "
if ( action == [ 1 ] ) :
return " Podlewac "