import gradio as gr
from fastai.vision.all import *
import timm
import random
import pickle


infile = open('gamename_dict.pkl','rb')
game_name_dict = pickle.load(infile)
infile.close()

learn=load_learner('convnext_nano_4freeze_5epochs_on_jpg_26error.pkl') 
categories=learn.dls.vocab
categories=[game_name_dict[gn] for gn in categories]

def recognize_game(img):
    img= PILImage.create(img)
    pred,idx,probs=learn.predict(img)
    return pred,str(100*float(dict(zip(categories, probs))[pred]))[:4]+"%"

def recognize_game_all(img):
    img= PILImage.create(img)
    pred,idx,probs=learn.predict(img)
    img.save(pred+str(random.randrange(1,10000))+'.jpg')
    return dict(zip(categories, map(float,probs)))


image = gr.inputs.Image(shape=(224,224))
label = gr.outputs.Label(num_top_classes=48)
title= 'A Game For That - Gamification Using Snapshot Images'
examples = ['1.1.jpg','1.2.jpg','2.jpg','3.jpg','4.jpg', '5.jpg' ]

iface = gr.Interface(
    fn=recognize_game_all, 
    inputs=image, 
    outputs=label,
    title=title, 
    examples=examples)

#iface.launch(inline=False)
iface.launch(debug=True)