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)