diff --git a/.gitignore b/.gitignore index e037cab..0abd250 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ __pycache__/ .idea/ -tree.png \ No newline at end of file +tree.png +dataset/ +dataset.zip \ No newline at end of file diff --git a/App.py b/App.py index 68aa702..591e937 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,20 @@ 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, 'model.pth') + print_to_console("sieć nuronowa nauczona") + 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)) + traktor.snake_move_predict_plant(pole, model) start_flag=False # demo_move() old_info=get_info(old_info) diff --git a/Image.py b/Image.py index a755f0a..90fefca 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,28 @@ class Image: def return_gasStation(self): return self.gasStation_image + +# losowanie zdjęcia z testowego datasetu bez powtórzeń +imagePathList = [] +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) + + while imgPath in imagePathList: + for event in pygame.event.get(): + if event.type == pygame.QUIT: + quit() + 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) + + imagePathList.append(imgPath) + + image = pygame.image.load(imgPath) + image=pygame.transform.scale(image,(dCon.CUBE_SIZE,dCon.CUBE_SIZE)) + return image, 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..69a971a 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,14 @@ 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() + # print(self.plant_image) + self.plant=Roslina.Roslina(self.label) self.set_image() def set_image(self): @@ -66,6 +73,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/Tractor.py b/Tractor.py index 71a0ee9..b135b08 100644 --- a/Tractor.py +++ b/Tractor.py @@ -9,11 +9,13 @@ import Osprzet import Node import Condition import Drzewo +import neuralnetwork as nn condition=Condition.Condition() drzewo=Drzewo.Drzewo() format_string = "{:<25}{:<25}{:<25}{:<10}{:<10}{:<10}{:<25}{:<15}{:<20}{:<10}{:<15}" +format_string_nn="{:<10}{:<20}{:<20}{:<15}{:<20}" tab = [-1, 0, 0, 0, 0, 1, 1, 1, 1, 1, @@ -191,6 +193,35 @@ class Tractor: self.turn_left() print("podlanych slotów: ", str(counter)) + def snake_move_predict_plant(self, pole, model): + headers=['Coords','Real plant','Predicted plant','Result','Fertilizer'] + print(format_string_nn.format(*headers)) + initPos = (self.slot.x_axis, self.slot.y_axis) + count = 0 + for i in range(initPos[1], dCon.NUM_Y): + for j in range(initPos[0], dCon.NUM_X): + for event in pygame.event.get(): + if event.type == pygame.QUIT: + quit() + if self.slot.imagePath != None: + predictedLabel = nn.predictLabel(self.slot.imagePath, model) + + #print(str("Coords: ({:02d}, {:02d})").format(self.slot.x_axis, self.slot.y_axis), "real:", self.slot.label, "predicted:", predictedLabel, "correct" if (self.slot.label == predictedLabel) else "incorrect", 'nawożę za pomocą:', nn.fertilizer[predictedLabel]) + print(format_string_nn.format(f"{self.slot.x_axis,self.slot.y_axis}",self.slot.label,predictedLabel,"correct" if (self.slot.label == predictedLabel) else "incorrect",nn.fertilizer[predictedLabel])) + if self.slot.label != predictedLabel: + self.slot.mark_visited() + count += 1 + self.move_forward(pole, False) + if i % 2 == 0 and i != dCon.NUM_Y - 1: + self.turn_right() + self.move_forward(pole, False) + self.turn_right() + elif i != dCon.NUM_Y - 1: + self.turn_left() + self.move_forward(pole, False) + self.turn_left() + print(f"Dobrze nawiezionych roślin: {20*12-count}, źle nawiezionych roślin: {count}") + def snake_move(self,pole,x,y): next_slot_coordinates=(x,y) if(self.do_move_if_valid(pole,next_slot_coordinates)): diff --git a/model_2_crops.pth b/model_2_crops.pth new file mode 100644 index 0000000..6954cc1 Binary files /dev/null and b/model_2_crops.pth differ diff --git a/model_500_hidden.pth b/model_500_hidden.pth new file mode 100644 index 0000000..30627b5 Binary files /dev/null and b/model_500_hidden.pth differ diff --git a/neuralnetwork.py b/neuralnetwork.py new file mode 100644 index 0000000..568abaf --- /dev/null +++ b/neuralnetwork.py @@ -0,0 +1,114 @@ +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 = ['carrot','corn', 'potato', 'tomato'] # musi być w kolejności alfabetycznej +fertilizer = {labels[0]: 'kompost', labels[1]: 'saletra amonowa', labels[2]: 'superfosfat', labels[3]:'obornik kurzy'} +#labels = ['corn','tomato'] #uncomment this two lines for 2 crops only +#fertilizer = {labels[0]: 'kompost', labels[1]: 'saletra amonowa'} +torch.manual_seed(42) + +#device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +device = torch.device("cpu") +# device = torch.device("mps") if torch.backends.mps.is_available() else torch.device('cpu') +# print(device) + +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 = 500 + 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, path): + print("Saving model") + torch.save(model.state_dict(), path) + +def loadModel(path): + print("Loading model") + 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) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model.to(device) + 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] + + # Znalezienie indeksu klasy o największej wartości prawdopodobieństwa + predicted_class = torch.argmax(output).item() + return labels[predicted_class] + + +def preprocess_image(image): + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + transform = getTransformation() + image = transform(image).unsqueeze(0) # Add batch dimension + image = image.to(device) # Move the image tensor to the same device as the model + return image + + diff --git a/testModels/modelCPUdataset2.pth b/testModels/modelCPUdataset2.pth new file mode 100644 index 0000000..c3d40b0 Binary files /dev/null and b/testModels/modelCPUdataset2.pth differ diff --git a/testModels/modelCPUdataset2_500.pth b/testModels/modelCPUdataset2_500.pth new file mode 100644 index 0000000..e9aa218 Binary files /dev/null and b/testModels/modelCPUdataset2_500.pth differ diff --git a/testModels/modelMPS.pth b/testModels/modelMPS.pth new file mode 100644 index 0000000..554a424 Binary files /dev/null and b/testModels/modelMPS.pth differ diff --git a/testModels/modelMPS650.pth b/testModels/modelMPS650.pth new file mode 100644 index 0000000..9504a6b Binary files /dev/null and b/testModels/modelMPS650.pth differ diff --git a/testModels/modelMPS_AL.pth b/testModels/modelMPS_AL.pth new file mode 100644 index 0000000..298d3db Binary files /dev/null and b/testModels/modelMPS_AL.pth differ