diff --git a/App.py b/App.py index 68aa702..b437928 100644 --- a/App.py +++ b/App.py @@ -9,6 +9,7 @@ import Osprzet import Ui import BFS import AStar +import neuralnetwork bfs1_flag=False @@ -18,7 +19,9 @@ Astar = False Astar2 = False if bfs3_flag or Astar or Astar2: Pole.stoneFlag = True -TreeFlag=True +TreeFlag=False +nnFlag=True +newModel=False pygame.init() show_console=True @@ -29,7 +32,7 @@ image_loader=Image.Image() image_loader.load_images() goalTreasure = AStar.getRandomGoalTreasure() # nie wiem czy to najlepsze miejsce, obecnie pole zawiera pole gasStation, które służy do renderowania odpowiedniego zdjęcia pole=Pole.Pole(screen,image_loader, goalTreasure) -pole.draw_grid() #musi byc tutaj wywołane ponieważ inicjalizuje sloty do slownika +pole.draw_grid(nnFlag) #musi byc tutaj wywołane ponieważ inicjalizuje sloty do slownika ui=Ui.Ui(screen) #Tractor creation traktor_slot = pole.get_slot_from_cord((0, 0)) @@ -40,7 +43,7 @@ def init_demo(): #Demo purpose old_info="" traktor.draw_tractor() time.sleep(2) - pole.randomize_colors() + pole.randomize_colors(nnFlag) traktor.draw_tractor() start_flag=True while True: @@ -116,6 +119,18 @@ def init_demo(): #Demo purpose if(TreeFlag): traktor.move_forward(pole) traktor.tree_move(pole) + if(nnFlag): + global model + if (newModel): + print_to_console("uczenie sieci neuronowej") + model = neuralnetwork.trainNewModel() + neuralnetwork.saveModel(model) + print('model został wygenerowany') + else: + model = neuralnetwork.loadModel('model.pth') + print_to_console("model został załądowny") + testset = neuralnetwork.getDataset(False) + print(neuralnetwork.accuracy(model, testset)) start_flag=False # demo_move() old_info=get_info(old_info) diff --git a/Image.py b/Image.py index a755f0a..732d622 100644 --- a/Image.py +++ b/Image.py @@ -1,6 +1,8 @@ import pygame import displayControler as dCon import random +import neuralnetwork +import os class Image: def __init__(self): @@ -53,3 +55,13 @@ class Image: def return_gasStation(self): return self.gasStation_image + +def getRandomImageFromDataBase(): + label = random.choice(neuralnetwork.labels) + + folderPath = f"dataset/test/{label}" + files = os.listdir(folderPath) + random_image = random.choice(files) + imgPath = os.path.join(folderPath, random_image) + + return pygame.image.load(imgPath), label, imgPath diff --git a/Pole.py b/Pole.py index 96544ee..f8e6aa3 100644 --- a/Pole.py +++ b/Pole.py @@ -6,6 +6,8 @@ import time import Ui import math import random +import neuralnetwork +import Image stoneList = [(3,3), (3,4), (3,5), (3,6), (4,6), (5,6), (6,6), (7,6), (8,6), (9,6), (10,6), (11,6), (12,6), (13,6), (14,6), (15,6), (16,6), (16,7), (16,8), (16,9)] stoneFlag = False @@ -30,7 +32,7 @@ class Pole: return self.slot_dict #Draw grid and tractor (new one) - def draw_grid(self): + def draw_grid(self, nn=False): for x in range(0,dCon.NUM_X): #Draw all cubes in X axis for y in range(0,dCon.NUM_Y): #Draw all cubes in Y axis new_slot=Slot.Slot(x,y,Colors.BROWN,self.screen,self.image_loader) #Creation of empty slot @@ -48,7 +50,7 @@ class Pole: st=self.slot_dict[self.gasStation] st.set_gasStation_image() - def randomize_colors(self): + def randomize_colors(self, nn = False): pygame.display.update() time.sleep(3) #self.ui.render_text("Randomizing Crops") @@ -59,7 +61,7 @@ class Pole: if(coordinates==(0,0)): continue else: - self.slot_dict[coordinates].set_random_plant() + self.slot_dict[coordinates].set_random_plant(nn) def change_color_of_slot(self,coordinates,color): #Coordinates must be tuple (x,y) (left top slot has cord (0,0) ), color has to be from defined in Colors.py or custom in RGB value (R,G,B) self.get_slot_from_cord(coordinates).color_change(color) diff --git a/Slot.py b/Slot.py index 30ad717..6f8e30c 100644 --- a/Slot.py +++ b/Slot.py @@ -16,6 +16,8 @@ class Slot: self.field=pygame.Rect(self.x_axis*dCon.CUBE_SIZE,self.y_axis*dCon.CUBE_SIZE,dCon.CUBE_SIZE,dCon.CUBE_SIZE) self.image_loader=image_loader self.garage_image=None + self.label = None + self.imagePath = None def draw(self): pygame.draw.rect(self.screen,Colors.BROWN,self.field,0) #Draw field @@ -38,9 +40,13 @@ class Slot: self.plant=color self.draw() - def set_random_plant(self): - (plant_name,self.plant_image)=self.random_plant() - self.plant=Roslina.Roslina(plant_name) + def set_random_plant(self, nn=False): + if not nn: + (plant_name,self.plant_image)=self.random_plant() + self.plant=Roslina.Roslina(plant_name) + else: + self.plant_image, self.label, self.imagePath = self.random_plant_dataset() + self.plant=Roslina.Roslina(self.label) self.set_image() def set_image(self): @@ -66,6 +72,8 @@ class Slot: def random_plant(self): #Probably will not be used later only for demo purpouse return self.image_loader.return_random_plant() + def random_plant_dataset(self): + return Image.getRandomImageFromDataBase() def return_plant(self): return self.plant diff --git a/model.pth b/model.pth new file mode 100644 index 0000000..90ece27 Binary files /dev/null and b/model.pth differ diff --git a/neuralnetwork.py b/neuralnetwork.py new file mode 100644 index 0000000..352f54a --- /dev/null +++ b/neuralnetwork.py @@ -0,0 +1,111 @@ +import torch +import torch.nn as nn +from torch.utils.data import DataLoader +from torchvision import datasets +from torchvision.transforms import Compose, Lambda, ToTensor +import torchvision.transforms as transforms +import matplotlib.pyplot as plt +from PIL import Image +import random + +imageSize = (128, 128) +labels = ['beetroot', 'potato', 'carrot'] +labels.sort() + +torch.manual_seed(42) + +device = torch.device('cuda') if torch.cuda.is_available () else torch.device('cpu') + +def getTransformation(): + transform=transforms.Compose([ + transforms.ToTensor(), + transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), + transforms.Resize(imageSize), + Lambda(lambda x: x.flatten())]) + return transform + +def getDataset(train=True): + transform = getTransformation() + if (train): + trainset = datasets.ImageFolder(root='dataset/train', transform=transform) + return trainset + else: + testset = datasets.ImageFolder(root='dataset/test', transform=transform) + return testset + + +def train(model, dataset, n_iter=100, batch_size=256): + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + criterion = nn.NLLLoss() + dl = DataLoader(dataset, batch_size=batch_size) + model.train() + for epoch in range(n_iter): + for images, targets in dl: + optimizer.zero_grad() + out = model(images.to(device)) + loss = criterion(out, targets.to(device)) + loss.backward() + optimizer.step() + if epoch % 10 == 0: + print('epoch: %3d loss: %.4f' % (epoch, loss)) + return model + +def accuracy(model, dataset): + model.eval() + correct = sum([(model(images.to(device)).argmax(dim=1) == targets.to(device)).sum() + for images, targets in DataLoader(dataset, batch_size=256)]) + return correct.float() / len(dataset) + +def getModel(): + hidden_size = 300 + model = nn.Sequential( + nn.Linear(imageSize[0] * imageSize[1] * 3, hidden_size), + nn.ReLU(), + nn.Linear(hidden_size, len(labels)), + nn.LogSoftmax(dim=-1) + ).to(device) + return model + +def saveModel(model): + torch.save(model.state_dict(), 'model.pth') + +def loadModel(path): + model = getModel() + model.load_state_dict(torch.load(path)) + return model + +def trainNewModel(n_iter=100, batch_size=256): + trainset = getDataset(True) + model = getModel() + model = train(model, trainset) + return model + +def predictLabel(imagePath, model): + image = Image.open(imagePath).convert("RGB") + image = preprocess_image(image) + with torch.no_grad(): + model.eval() # Ustawienie modelu w tryb ewaluacji + output = model(image) + + # Znalezienie indeksu klasy o największej wartości prawdopodobieństwa + predicted_class = torch.argmax(output).item() + return labels[predicted_class] + +def predictLabel(image, model): + image = preprocess_image(image) + with torch.no_grad(): + model.eval() # Ustawienie modelu w tryb ewaluacji + output = model(image) + + # Znalezienie indeksu klasy o największej wartości prawdopodobieństwa + predicted_class = torch.argmax(output).item() + return labels[predicted_class] + + +def preprocess_image(image): + transform = getTransformation() + image = transform(image).unsqueeze(0) # Dodanie wymiaru batch_size + + return image + +