s481808/Traktor
4
2
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

added functions for loading images, model and testing

Browse Source
This commit is contained in:
MarRac 2024-05-26 19:56:18 +02:00
parent fb0ec5057c
commit b45c2e0f1f
1 changed files with 27 additions and 26 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

45
source/NN/neural_network.py
View File

@ -4,15 +4,14 @@ from torch.utils.data import DataLoader
from torchvision import datasets, transforms, utils
from torchvision.transforms import Compose, Lambda, ToTensor
import matplotlib.pyplot as plt
import numpy as np
from model import *
from .model import *
from PIL import Image
device = torch.device('cuda')
#data transform to tensors:
data_transformer = transforms.Compose([
transforms.Resize((150, 150)),
transforms.Resize((100, 100)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5 ), (0.5, 0.5, 0.5))
])
@ -24,13 +23,8 @@ test_set = datasets.ImageFolder(root='resources/test', transform=data_transforme
#to mozna nawet przerzucic do funkcji train:
#train_loader = DataLoader(train_set, batch_size=32, shuffle=True, num_workers=2)
#test_loader = DataLoader(test_set, batch_size=32, shuffle=True, num_workers=2)
#test if classes work properly:
#print(train_set.classes)
#print(train_set.class_to_idx)
#print(train_set.targets[3002])
# train_loader = DataLoader(train_set, batch_size=64, shuffle=True)
#test_loader = DataLoader(test_set, batch_size=32, shuffle=True)
#function for training model
@ -62,12 +56,10 @@ def accuracy(model, dataset):
return correct.float() / len(dataset)
model = Neural_Network_Model()
model.to(device)
#loading the already saved model:
model.load_state_dict(torch.load('model.pth'))
model.eval()
@ -78,18 +70,27 @@ model.eval()
#TEST - loading the image and getting results:
testImage_path = 'resources/images/plant_photos/pexels-polina-tankilevitch-4110456.jpg'
testImage = Image.open(testImage_path)
#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.eval()
return model
def load_image(image_path):
testImage = Image.open(image_path)
testImage = data_transformer(testImage)
testImage = testImage.unsqueeze(0)
testImage = testImage.to(device)
return testImage
model.load_state_dict(torch.load('model.pth'))
model.to(device)
model.eval()
testOutput = model(testImage)
def guess_image(model, image_tensor):
with torch.no_grad():
testOutput = model(image_tensor)
_, predicted = torch.max(testOutput, 1)
predicted_class = train_set.classes[predicted.item()]
print(f'The predicted class is: {predicted_class}')
return predicted_class