Dodanie 'grader.py'

grader.py

@@ -0,0 +1,65 @@
+import torch
+import torchvision
+import torchvision.transforms as transforms
+import torch.nn as nn
+import torch.nn.functional as f
+import torch.optim as optim
+import numpy as np
+from matplotlib.pyplot import imshow
+import os
+import PIL
+import numpy as np
+import neural_network
+from matplotlib.pyplot import imshow
+
+# Create the model
+model = neural_network.Net()
+
+# Load state_dict
+neural_network.load_network_from_structure(model)
+
+# Create the preprocessing transformation here
+transform = transforms.Compose([neural_network.Negative(), transforms.ToTensor()])
+
+# load your image(s)
+img = PIL.Image.open('test\\0_100.jpg')
+img2 = PIL.Image.open('test\\1_100.jpg')
+img3 = PIL.Image.open('test\\4_100.jpg')
+img4 = PIL.Image.open('test\\5_100.jpg')
+
+# Transform
+input = transform(img)
+input2 = transform(img2)
+input3 = transform(img3)
+input4 = transform(img4)
+
+# unsqueeze batch dimension, in case you are dealing with a single image
+input = input.unsqueeze(0)
+input2 = input2.unsqueeze(0)
+input3 = input3.unsqueeze(0)
+input4 = input4.unsqueeze(0)
+
+# Set model to eval
+model.eval()
+
+# Get prediction
+output = model(input)
+output2 = model(input2)
+output3 = model(input3)
+output4 = model(input4)
+
+print(output)
+index = output.cpu().data.numpy().argmax()
+print(index)
+
+print(output2)
+index = output2.cpu().data.numpy().argmax()
+print(index)
+
+print(output3)
+index = output3.cpu().data.numpy().argmax()
+print(index)
+
+print(output4)
+index = output4.cpu().data.numpy().argmax()
+print(index)