65 lines
1.4 KiB
Python
65 lines
1.4 KiB
Python
|
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)
|