import gradio as gr

from fastcore.all import *
from fastai.vision.all import *
from fastai.vision.widgets import *
import pickle
import torch


def label_func (f):
    return f[:2] !="1 "

Pkl_Filename = "learner.pkl"
with open(Pkl_Filename, 'rb') as file:  
    learn = torch.load ( file , map_location=torch.device('cpu') )
    
category = ("Waifu" , "Human")
def infer_im (im ) :
    pre , idx , prob = learn.predict (im)
    flo_prob = list ( map (float , prob) )[0]
    return dict ( zip (category , ( flo_prob , 1 - flo_prob ) ) )


examples = ["1 (1004).jpg",
           "498970d2f45988eae9d6a2eb59bc450f.jpg" ,
           "10010892_result.jpg" , 
           "10019581_result.jpg"]

gr.Interface(fn= infer_im, 
             inputs=gr.inputs.Image(shape=(128 , 128)),
             outputs=gr.outputs.Label(num_top_classes=2) ,
            examples = examples ).launch(inline = False)