Update main.py

main.py

@@ -5,13 +5,15 @@ from tractor import Tractor
 from kolejka import Stan, Odwiedzone
 from queue import Queue
 import pandas as pd
-
+from neuralnetwork import load_model
 
 data = pd.read_csv('dane.csv')
 
 
 from decisiontree import train_decision_tree
 
+model_path = 'model.pth'
+neuralnetwork_model = load_model(model_path)
 
 model, feature_columns = train_decision_tree(data)
 
@@ -60,7 +62,7 @@ def main():
     clock = pygame.time.Clock()
     board = Board()
     board.load_images()
-    tractor = Tractor(4, 4, model, feature_columns)
+    tractor = Tractor(4, 4, model, feature_columns, neuralnetwork_model)
 
     while run:
         clock.tick(fps)
@@ -75,7 +77,7 @@ def main():
             print("Nie znaleziono ścieżki do najbliższego pola dirt.")
             board.generate_board()
             initial_state = Stan(4, 4, "down")
-            tractor = Tractor(4, 4, model, feature_columns)
+            tractor = Tractor(4, 4, model, feature_columns, neuralnetwork_model)
             while board.is_rock(initial_state.row, initial_state.col):
                 board.generate_board()
             continue