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

Neural_network #4

Merged
s481894 merged 12 commits from Neural_network into master 2024-06-04 16:59:03 +02:00
Conversation 0 Commits 12 Files Changed 4847 +27 -26
1 changed files with 27 additions and 26 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files
Showing only changes of commit b45c2e0f1f - Show all commits

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