import gradio as gr
from fastai.vision.all import *
from PIL import Image as pilIm
#
#learn = load_learner('export.pkl')
#learn = torch.load('digit_classifier.pth')
#learn.eval() #switch to eval mode
model_dict=torch.load('my_model.pt')
W1,B1,W2,B2,W3,B3=model_dict['W1'],model_dict['B1'],model_dict['W2'],model_dict['B2'],model_dict['W3'],model_dict['B3']
def mdlV2(xb):
    res = xb@W1+B1
    res = res.max(tensor(0.))
    res = res@W2+B2  # returns 10 features for each input
    res = res.max(tensor(0.))
    res = res@W3+B3  # returns 10 features for each input
    return res


labels = [str(x) for x in range(10)]
# #################################
# #Define class for importing Model
# class DigitClassifier(torch.nn.Module):
#     def __init__(self):
#         super().__init__()
#         self.fc1 = torch.nn.Linear(64, 32)
#         self.fc2 = torch.nn.Linear(32, 16)
#         self.fc3 = torch.nn.Linear(16, 10)

#     def forward(self, x):
#         x = x.view(-1, 64)
#         x = torch.relu(self.fc1(x))
#         x = torch.relu(self.fc2(x))
#         x = self.fc3(x)
#         return x
#########################################
#Define function to reduce image of arbitrary size to 8x8 per model requirements.
def reduce_image_count(image):
    output_size = (8, 8)
    block_size = (image.shape[0] // output_size[0], image.shape[1] // output_size[1])
    output = np.zeros(output_size)
    
    for i in range(output_size[0]):
        for j in range(output_size[1]):
            block = image[i*block_size[0]:(i+1)*block_size[0], j*block_size[1]:(j+1)*block_size[1]]
            count = np.count_nonzero(block)
            output[i, j] = count
    normalizer=np.amax(output)
    output=output*16/normalizer
    return output

#########################################

def predict(img):
    #First take input and reduce it to 8x8 px as the dataset was
    pil_image = pilIm.open(img) #get image
    gray_img = pil_image.convert('L')#grayscale
    pic = np.array(gray_img) #convert to array
    inp_img=reduce_image_count(pic)#Reduce image to required input size
    z=Tensor(inp_img)
    y=z.view(-1,64)
    x=mdlV2(y)
    w=F.softmax(x,dim=-1)
    v=w[0]
    u=v.data
    otpt=u
    #pred,pred_idx,probs = learn.predict(img)
    
    return dict([[labels[i], float(otpt[i])] for i in range(len(labels))]),inp_img/16

gr.Interface(fn=predict, inputs=gr.inputs.Image(type='filepath'),   outputs=[gr.outputs.Label(num_top_classes=10), gr.outputs.Image()]).launch()