import gradio
import pandas as pd
import PIL.Image
import tempfile, pathlib, requests, io

from huggingface_hub import snapshot_download
from autogluon.multimodal import MultiModalPredictor

# Hugging Face model repo
MODEL_REPO_ID = "scottymcgee/image-classifier"

# Download full repo snapshot locally
repo_dir = snapshot_download(repo_id=MODEL_REPO_ID)

# Load predictor directly
PREDICTOR = MultiModalPredictor.load(repo_dir)
TMP_PATH = pathlib.Path("tmp_input.png")

# Explicit class labels
CLASS_LABELS = {0: "🚫 No Stop Sign", 1: "🛑 Stop Sign"}

def do_predict(pil_img: PIL.Image.Image):
    if pil_img is None:
        return {"Error": 1.0}
        
    pil_resized = pil_img.convert("RGB").resize((256, 256))
    
    # Save to tmp for AutoGluon input
    pil_resized.save(TMP_PATH)

    df = pd.DataFrame({"image": [str(TMP_PATH)]})

    # Predict probabilities
    proba_df = PREDICTOR.predict_proba(df)

    # Map columns to names
    pretty_dict = {}
    for idx, col in enumerate(proba_df.columns):
        label = CLASS_LABELS.get(idx, str(col))
        pretty_dict[label] = float(proba_df.iloc[0][col])
    return pil_resized, pretty_dict

EXAMPLES = [
    ["examples/stop1.jpg"],
    ["examples/no_stop1.jpg"],
    ["examples/stop2.jpg"],
]

with gradio.Blocks() as demo:
    gradio.Markdown("# Stop Sign Classifier")
    gradio.Markdown("Upload a traffic image. The model shows the **original**, the **preprocessed**, and the prediction.")

    with gradio.Row():
        threshold = gradio.Slider(
            0.0, 1.0, value=0.2, step=0.05, label="Confidence threshold"
        )

    with gradio.Row():
        image_in = gradio.Image(type="pil", label="Upload image", sources=["upload", "webcam"])
        proc_out = gradio.Image(type="pil", label="Preprocessed (256x256)")
        proba_pretty = gradio.Label(num_top_classes=2, label="Class probabilities")

    image_in.change(fn=do_predict, inputs=[image_in], outputs=[proc_out, proba_pretty])

    gradio.Examples(
        examples=EXAMPLES,
        inputs=[image_in],
        outputs=[proc_out, proba_pretty],
        fn=do_predict,
        cache_examples=False,
    )

if __name__ == "__main__":
    demo.launch()