from fastapi import FastAPI, File
from fastai.vision.all import *
from typing_extensions import Annotated
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))

categories = ('Dog', 'Cat')


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


@app.post("/classify")
def classify_image(file: Annotated[bytes, File()]):
    img = PILImage.create(file)
    img.thumbnail((192, 192))

    return classify_images(img)