from fastcore.all import *
from fastai.vision.all import *
import gradio as gr
from huggingface_hub import hf_hub_download
from PIL import Image

CATEGORIES = ["human", "robot"]
IMG_SIZE = 226

model = os.path.join("models","human_or_robot_model.pkl")

# if os.path.exists(model):
#     learn = load_learner(model)
# else:
hf_model = hf_hub_download(
    repo_id="bengid/human_or_robot_model",
    filename="human_or_robot_model.pkl"
)
learn = load_learner(hf_model)

# hf_weights = hf_hub_download(
#     repo_id="bengid/human_or_robot_model",
#     filename="model_weights.pth"
# )

# learn = vision_learner()

# def build_learner():
#     dls = ImageDataLoaders.from_folder()
    
def classify_image(img):
    pil = Image.fromarray(img).resize((IMG_SIZE, IMG_SIZE))

    _, _, probs = learn.predict(pil)
    return dict(zip(CATEGORIES, map(float, probs)))


def main():
    image = gr.Image(type="numpy")
    label = gr.Label()
    examples = "examples"
    intf = gr.Interface(fn=classify_image, inputs=image, outputs=label, examples=examples)
    intf.launch(inline=False)

if __name__ == "__main__":
    main()