Neural_network #4

Merged
s481894 merged 12 commits from Neural_network into master 2024-06-04 16:59:03 +02:00
16 changed files with 60 additions and 23 deletions
Showing only changes of commit 9667655a2a - Show all commits

Binary file not shown.

View File

@ -3,16 +3,26 @@ import torch
import torch.nn.functional as F
class Neural_Network_Model(nn.Module):
def __init__(self, num_classes=5,hidden_layer1 = 100,hidden_layer2 = 100):
super(Neural_Network_Model, self).__init__()
self.fc1 = nn.Linear(150*150*3,hidden_layer1)
class Conv_Neural_Network_Model(nn.Module):
def __init__(self, num_classes=5,hidden_layer1 = 512,hidden_layer2 = 256):
super(Conv_Neural_Network_Model, self).__init__()
self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, stride=1, padding=1)
self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1)
self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
self.conv3 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1)
self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
self.fc1 = nn.Linear(128*12*12,hidden_layer1)
self.fc2 = nn.Linear(hidden_layer1,hidden_layer2)
self.out = nn.Linear(hidden_layer2,num_classes)
# two hidden layers
def forward(self, x):
x = x.view(-1, 150*150*3)
x = self.pool1(F.relu(self.conv1(x)))
x = self.pool2(F.relu(self.conv2(x)))
x = self.pool3(F.relu(self.conv3(x)))
x = x.view(-1, 128*12*12) #<----flattening the image
x = self.fc1(x)
x = torch.relu(x)
x = self.fc2(x)

View File

@ -4,7 +4,7 @@ from torch.utils.data import DataLoader
from torchvision import datasets, transforms, utils
from torchvision.transforms import Compose, Lambda, ToTensor
import matplotlib.pyplot as plt
from .model import *
from model import *
from PIL import Image
device = torch.device('cuda')
@ -29,7 +29,7 @@ test_set = datasets.ImageFolder(root='resources/test', transform=data_transforme
#function for training model
def train(model, dataset, iter=100, batch_size=64):
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
criterion = nn.NLLLoss()
train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
model.train()
@ -56,12 +56,12 @@ def accuracy(model, dataset):
return correct.float() / len(dataset)
model = Neural_Network_Model()
model = Conv_Neural_Network_Model()
model.to(device)
#loading the already saved model:
model.load_state_dict(torch.load('model.pth'))
model.eval()
# model.load_state_dict(torch.load('model.pth'))
# model.eval()
#training the model:
# train(model, train_set)
@ -69,12 +69,10 @@ model.eval()
# torch.save(model.state_dict(), 'model.pth')
#TEST - loading the image and getting results:
#testImage_path = 'resources/images/plant_photos/pexels-dxt-73640.jpg'
def load_model():
model = Neural_Network_Model()
model.load_state_dict(torch.load('model.pth'))
model = Conv_Neural_Network_Model()
model.load_state_dict(torch.load('CNN_model.pth'))
model.eval()
return model
@ -94,3 +92,23 @@ 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/pexels-justus-menke-3490295-5213970.jpg'
testImage = Image.open(testImage_path)
testImage = data_transformer(testImage)
testImage = testImage.unsqueeze(0)
testImage = testImage.to(device)
model.load_state_dict(torch.load('CNN_model.pth'))
model.to(device)
model.eval()
testOutput = model(testImage)
_, predicted = torch.max(testOutput, 1)
predicted_class = train_set.classes[predicted.item()]
print(f'The predicted class is: {predicted_class}')

Binary file not shown.

View File

@ -3,16 +3,17 @@ import time
import random
import pandas as pd
import joblib
from area.constants import WIDTH, HEIGHT, TILE_SIZE
from area.field import drawWindow
from area.tractor import Tractor, do_actions
from area.field import tiles, fieldX, fieldY
from area.field import get_tile_coordinates
from area.field import get_tile_coordinates, get_tile_index
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
WIN = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption('Intelligent tractor')
@ -23,7 +24,7 @@ def main():
pygame.display.update()
#getting coordinates of our "goal tile":
tile_index=127
tile_index = get_tile_index()
tile_x, tile_y = get_tile_coordinates(tile_index)
if tile_x is not None and tile_y is not None:
print(f"Coordinates of tile {tile_index} are: ({tile_x}, {tile_y})")
@ -128,6 +129,13 @@ def main():
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}")
time.sleep(30)
print("\n")

Binary file not shown.

View File

@ -2,6 +2,7 @@ import random
import os
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
# path to plant images folder (used in randomize_photo function)
folder_path = "resources/images/plant_photos"
@ -50,9 +51,9 @@ class Tile:
self.image = "resources/images/rock_dirt.png"
def displayPhoto(self):
img = self.photo
img = img / 2 + 0.5 # unnormalize
npimg = img.numpy()
plt.imshow(np.transpose(npimg, (1, 2, 0)))
def display_photo(self):
image_path = self.photo
img = Image.open(image_path)
plt.imshow(img)
plt.axis('off')
plt.show()