import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
class ConvolutionalNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        #The first layer has one input channel (the grayscalle color channel). We'll use 6 output channels for feature extraction and
        #a kernel size of 3 to make a 3x3 filter. The step size will be 1.
        self.conv1 = nn.Conv2d(1, 6, 3, 1)
        #The second layer will take our 6 input channels and deliver 16 output channels.
        self.conv2 = nn.Conv2d(6, 16, 3, 1)


        #The input size of (5x5x16) is determined by the effect of our kernels on the input image size.
        #A 3x3 filter applied to a 28x28 image leaves a 1-pixel edge on all four sides.
        #In one layer the size changes from 28x28 to 26x26. We could address this with zero-padding, but since an MNIST image is mostly black
        #at the edges, we should be safe ignoring these pixels. We'll apply the kernel twice, and apply pooling layers twice,
        #so our resulting output will be $\;(((28-2)/2)-2)/2 = 5.5\;$ which rounds down to 5 pixels per side.
        self.fc1 = nn.Linear(5*5*16, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84,10)

    def forward(self, X):
        X = F.relu(self.conv1(X))
        X = F.max_pool2d(X, 2, 2)
        X = F.relu(self.conv2(X))
        X = F.max_pool2d(X, 2, 2)

        #Flatten the data for the fully connected layers
        X = X.view(-1, 5*5*16)
        X = F.relu(self.fc1(X))
        X = F.relu(self.fc2(X))
        X = self.fc3(X)
        return F.log_softmax(X, dim=1)

model2 = ConvolutionalNetwork()
model2.load_state_dict(torch.load('mnist2.pt'));
model2.eval()
cap = cv2.VideoCapture(0)


width = 640
height = 480

while True:

    ret, frame = cap.read()
    img = cv2.cvtColor(frame,cv2.COLOR_BGR2GRAY)

    img = img/255
    img = cv2.resize(img,(28,28))
    img = img.reshape(1,1,28,28)
    img = torch.from_numpy(img)
    with torch.no_grad():
        # view(1,1,28,28) batch size of one, one channel, 28,28 pixels
        label = model2(img.float()).argmax()
        log_prob =  model2(img.float())
        probs = torch.exp(log_prob)*100
        max_probs = torch.max(probs[0])
        #max_probs = torch.max(probs[0])
        print("Predicted value:",label.item())
        print("Probability:", max_probs.item())
    cv2.putText(frame,str(label) + "   "+str(max_probs),
                    (50,50),cv2.FONT_HERSHEY_COMPLEX,
                    1,(0,0,255),1)
    cv2.imshow('frame',frame)
    print(img.shape)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()