source/NN/neural_network.py

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