def classify(self, show_output=False):
"""Send the preprocessed images to the NN classifier"""
print('{0} Numbers to be classified'.format(len(self.cropped_images)))
return_list = []
self.apply_cropping(show_output=show_output)
net = Neural_Network()
net.load_state_dict(torch.load(TENSOR_LOCATION))
net.eval()
for image in self.cropped_images:
image = Image.fromarray(image)
# Resizes the number and adds a 10 px border
transfrom = transforms.Compose([
transforms.Grayscale(),
transforms.Resize(self.output_size - self.border_size),
transforms.CenterCrop(self.output_size),
transforms.ToTensor(),
])
img_tensor = transfrom(image)
if show_output:
plt.imshow(np.array(img_tensor)[0, :, :],
cmap=plt.cm.gray_r,
interpolation='nearest')
plt.title('Image used for classification')
plt.show()
img_tensor.unsqueeze_(0)
outputs = net.forward(Variable(img_tensor))
dummy, predicted_labels = torch.max(outputs.data, 1)
return_list.append(int(predicted_labels.numpy().max()))
print('Classified: {0}'.format(predicted_labels.numpy().max()))
return return_list
