diff --git a/source/main.py b/source/main.py
index 4de86f1..5b4b62a 100644
--- a/source/main.py
+++ b/source/main.py
@@ -62,15 +62,35 @@ def main():
     while run:
         for event in pygame.event.get():
             if event.type == pygame.QUIT:
-                run = False
-
-        #small test of work_on_field method:        
+                run = False       
         time.sleep(1)
-        tile1 = tiles[0]
-        p1 = Plant('wheat', 'cereal', random.randint(1,100), random.randint(1,100), random.randint(1,100))      
+
+        # movement based on route-planning (test):
+
+        tractor.draw_tractor(WIN)
+        time.sleep(1)
+        if moves != False:
+            do_actions(tractor, WIN, moves)
+
+
+        #guessing the image under the tile:
+        goalTile = tiles[tile_index]
+        goalTile.display_photo()
+        image_path = goalTile.photo
+        image_tensor = load_image(image_path)
+        prediction = guess_image(load_model(), image_tensor)
+        print(f"The predicted image is: {prediction}")
+
+
+        p1 = Plant('wheat', 'cereal', random.randint(1,100), random.randint(1,100), random.randint(1,100))
+        goalTile.plant = p1
         d1 = Dirt(random.randint(1, 100), random.randint(1,100))
         d1.pests_and_weeds()
-        tile1.ground=d1
+        goalTile.ground=d1
+        #getting the name and type of the recognized plant:
+        p1.update_name(prediction)
+
+        #decission tree test:
         if d1.pest:
             pe = 1
         else:
@@ -117,25 +137,12 @@ def main():
 
         model = joblib.load('model.pkl')
         nowe_dane = pd.read_csv('model_data.csv')
-
         predykcje = model.predict(nowe_dane)
-
-        # movement based on route-planning (test):
-
-        tractor.draw_tractor(WIN)
-        time.sleep(1)
-        if moves != False:
-            do_actions(tractor, WIN, moves)
         print(predykcje)
-        if predykcje == 'work':
-            tractor.work_on_field(tile1, d1, p1)
 
-        #guessing the image under the tile:
-        tiles[tile_index].display_photo()
-        image_path = tiles[tile_index].photo
-        image_tensor = load_image(image_path)
-        prediction = guess_image(load_model(), image_tensor)
-        print(f"The predicted image is: {prediction}")
+        #work on field:
+        if predykcje == 'work':
+            tractor.work_on_field(goalTile, d1, p1)
         time.sleep(30)
         print("\n")
 
diff --git a/source/plant.py b/source/plant.py
index 0072b54..f7197f4 100644
--- a/source/plant.py
+++ b/source/plant.py
@@ -19,7 +19,23 @@ class Plant:
         else:
             print("Unable to grow due to bad condition of the ground")
 
-#   more properties
+    def update_name(self, predicted_class):
+        if predicted_class == "Apple":
+            self.name = "Apple"
+            self.plant_type = "fruit"
 
+        elif predicted_class == "Strawberry":
+            self.name = "Strawberry"
+            self.plant_type = "fruit"
 
-# add init, getters,setters
+        elif predicted_class == "Cucumber":
+            self.name = "Cucumber"
+            self.plant_type = "vegetable"
+
+        elif predicted_class == "Cauliflower":
+            self.name = "Cauliflower"
+            self.plant_type = "vegetable"
+            
+        elif predicted_class == "Wheat":
+            self.name = "Wheat"
+            self.plant_type = "cereal"
\ No newline at end of file