dodanie podstaw sieci nuronowej

This commit is contained in:
tafit0902 2024-05-25 02:00:19 +02:00
parent 971746a0f8
commit f38a52e135
6 changed files with 157 additions and 9 deletions

21
App.py
View File

@ -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)

View File

@ -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

View File

@ -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)

14
Slot.py
View File

@ -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

BIN
model.pth Normal file

Binary file not shown.

111
neuralnetwork.py Normal file
View File

@ -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