added photo display next to the field

source/NN/neural_network.py

@@ -6,6 +6,7 @@ from torchvision.transforms import Compose, Lambda, ToTensor
 import matplotlib.pyplot as plt
 from NN.model import *
 from PIL import Image
+import pygame
 
 device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
 
@@ -83,6 +84,16 @@ def load_image(image_path):
     testImage = testImage.unsqueeze(0)
     return testImage
 
+def display_image(screen, image_path, position):
+    image = pygame.image.load(image_path)
+    image = pygame.transform.scale(image, (250, 250))
+    screen.blit(image, position)
+
+def display_result(screen, position, predicted_class):
+    font = pygame.font.Font(None, 30)
+    displayed_text = font.render("The predicted image is: "+str(predicted_class), 1, (255,255,255))
+    screen.blit(displayed_text, position)
+
 def guess_image(model, image_tensor):
     with torch.no_grad():
         testOutput = model(image_tensor)
@@ -92,11 +103,6 @@ def guess_image(model, image_tensor):
 
 
 
-# image_path = 'resources/images/plant_photos/pexels-dxt-73640.jpg'
-# image_tensor = load_image(image_path)
-# prediction = guess_image(load_model(), image_tensor)
-# print(f"The predicted image is: {prediction}")
-
 #TEST - loading the image and getting results:
 # testImage_path = 'resources/images/plant_photos/1c76aa4d-11f4-47d1-8bdd-2cb78deeeccf.jpg'
 # testImage = Image.open(testImage_path)

source/main.py

@@ -12,10 +12,12 @@ from ground import Dirt
 from plant import Plant
 from bfs import graphsearch, Istate, succ
 from astar import a_star
-from NN.neural_network import load_model, load_image, guess_image
+from NN.neural_network import load_model, load_image, guess_image, display_image, display_result
 from PIL import  Image
 
-WIN = pygame.display.set_mode((WIDTH, HEIGHT))
+pygame.init()
+WIN_WIDTH = WIDTH + 300
+WIN = pygame.display.set_mode((WIN_WIDTH, HEIGHT))
 pygame.display.set_caption('Intelligent tractor')
 
 
@@ -76,10 +78,15 @@ def main():
 
         #guessing the image under the tile:
         goalTile = tiles[tile_index]
-        goalTile.display_photo()
         image_path = goalTile.photo
+        display_image(WIN, goalTile.photo, (WIDTH-20 , 300))   #displays photo next to the field
+        pygame.display.update()
+
         image_tensor = load_image(image_path)
         prediction = guess_image(load_model(), image_tensor)
+        
+        display_result(WIN, (WIDTH - 50 , 600), prediction)      #display text under the photo
+        pygame.display.update()
         print(f"The predicted image is: {prediction}")
 
 
@@ -144,7 +151,7 @@ def main():
         #work on field:
         if predykcje == 'work':
             tractor.work_on_field(goalTile, d1, p1)
-        time.sleep(30)
+        time.sleep(50)
         print("\n")
 
 

source/tile.py

@@ -52,12 +52,3 @@ class Tile:
             self.image = "resources/images/rock_dirt.png"
 
 
-    def display_photo(self):
-        image_path = self.photo
-        img = Image.open(image_path)
-        plt.ion()
-        plt.imshow(img)
-        plt.axis('off')
-        plt.show()
-        time.sleep(5)
-        plt.close()