import gradio as gr
from fastai.vision.all import *
from PIL import Image
#
#learn = load_learner('export.pkl')
learn = torch.load('digit_classifier.pth')
learn.eval() #switch to eval mode
labels = [str(x) for x in range(10)]

#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] = 16 - ((count / (block_size[0] * block_size[1])) * 16)
            
    return output

def predict(img):
    #First take input and reduce it to 8x8 px as the dataset was
    pil_image = Image.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
    
    otpt=F.softmax(learn.forward(inp_img.view(-1,64)))
    #pred,pred_idx,probs = learn.predict(img)
    
    return {labels[i]: float(otpt[0].data[i]) for i in range(len(labels)),'image': inp_img}

gr.Interface(fn=predict, inputs=gr.inputs.Image(shape=(512, 512)),   outputs=[gr.outputs.Label(num_top_classes=3), gr.outputs.Image()]).launch(share=True)