from fastapi import FastAPI
from fastai.vision.all import *
from huggingface_hub import hf_hub_download
import __main__

app = FastAPI()

REPO_ID = "jsagl/test-fastai"
FILENAME = "model.pkl"

def is_cat(x): return x[0].isupper()
__main__.is_cat = is_cat

learn = load_learner(hf_hub_download(repo_id=REPO_ID, filename=FILENAME))

@app.get("/")
def root():
    categories = ('Dog', 'Cat')

    dog = PILImage.create('/code/app/dog.jpg')
    dog.thumbnail((192, 192))
    dog

    def classify_images(img):
        pred, idx, probs = learn.predict(img)
        return dict(zip(categories, map(float, probs)))


    return classify_images(dog)