from fastai.vision.all import *
import gradio as gr
import os

# Load the learner model
learn_inf = load_learner("model/model.pkl")

# Get the categories from the model
categories = learn_inf.dls.vocab

# Comma separated string of categories with and between the last two
categories_str = ", ".join(categories[:-1]) + " or " + categories[-1]

def classify_image(img):
    pred, idx, probs = learn_inf.predict(img)
    return {f"{categories[idx]}": float(probs[idx])}


image = gr.Image(sources=["upload"], label="Image")
label = gr.Label()

# Load the images from the examples folder
examples = []
for file in os.listdir("images/examples"):
    if file.endswith(".jpg"):
        examples.append([f"images/examples/{file}"])

iface = gr.Interface(
    fn=classify_image,
    inputs=image,
    outputs=label,
    examples=examples,
    analytics_enabled=False,
    title="<h1>House Room Classifier</h1>",
    description='<p style = "font-size: 1.4rem;">This model classifies a photo of a house room as either <strong>' + categories_str + '</strong>.</p><p style = "font-size: 1.2rem;">Upload an image or use the examples below.</p>',
    flagging_options=[],
)

iface.launch()