#install huggingface_hub["fastai"] gradio timm
from huggingface_hub import from_pretrained_fastai
from gradio import Interface, inputs, outputs
from fastai.learner import Learner
import fastai

repo_id = "Kieranm/britishmus_plate_material_classifier"

learner = from_pretrained_fastai(repo_id)

mappings = {
    fastai.torch_core.TensorImage: {
        "type": inputs.Image(type='file', label='input'),
        "process": lambda inp : inp.name
    },
    fastai.torch_core.TensorCategory: {
        "type": outputs.Label(num_top_classes=3, label = 'output'),
        "process": lambda dls, out: {dls.vocab[i]: float(out[2][i]) for i in range(len(dls.vocab))}

    }
}

#Taken from fastgradio library

class Demo:
    def __init__(self, learner):

        self.learner = learner
        self.types = getattr(self.learner.dls, '_types')[tuple]

    def learner_predict(self, inp):
        inp = mappings[self.types[0]]["process"](inp)
        prediction = self.learner.predict(inp)
        output = mappings[self.types[1]]["process"](self.learner.dls, prediction)
        return output

    def launch(self, share=True, debug=False, auth=None, **kwargs):
        inputs = mappings[self.types[0]]["type"]

        outputs = mappings[self.types[1]]["type"]

        Interface(fn=self.learner_predict, inputs=inputs, outputs=outputs, 
                  examples = ["examples/earthen1.jpg", "examples/earthen2.png", "examples/porcelain1.png", "examples/porcelain2.png"],
                   **kwargs).launch(share=share, debug=debug, auth=auth)


Demo(learner).launch()