from fastai.vision.core import PILImageBW, TensorImageBW
from datasets import ClassLabel
import gradio as gr
from fastai.learner import load_learner

def get_image_attr(x): return x['image']
def get_target_attr(x): return x['target']

def img2tensor(im: Image.Image):
  return TensorImageBW(array(im)).unsqueeze(0)

classLabel = ClassLabel(names=['T - shirt / top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot'], id=None)

def add_target(x:dict):
  x['target'] = classLabel.int2str(x['label'])
  return x

learn = load_learner('export.pkl', cpu=True)

def classify(inp):
  img = PILImageBW.create(inp)
  item = dict(image=img)
  pred, _, _ = learn.predict(item)
  return classLabel.int2str(int(pred))

iface = gr.Interface(
    fn=classify,
    inputs=gr.inputs.Image(),
    outputs="text",
    title="Fashion Mnist Classifier",
    description="fastai deployment in Gradio.",
).launch()