Lewy
328b171e5a
- changed name of files to represent inside - added main to seperate processes - changed raw code to functions
65 lines
1.6 KiB
Python
65 lines
1.6 KiB
Python
import PIL
|
|
import torchvision.transforms as transforms
|
|
|
|
from AI import neural_network
|
|
|
|
|
|
# wcześniej grinder.py
|
|
# Get accuracy for neural_network model 'network_model.pth'
|
|
def NN_accuracy():
|
|
# 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('../src/test/0_100.jpg')
|
|
img2 = PIL.Image.open('../src/test/1_100.jpg')
|
|
img3 = PIL.Image.open('../src/test/4_100.jpg')
|
|
img4 = PIL.Image.open('../src/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)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
NN_accuracy()
|