# app.py
import os, shutil, zipfile, pathlib, tempfile
import pandas as pd
import gradio as gr
from PIL import Image
from autogluon.multimodal import MultiModalPredictor

ZIP_FILENAME = "autogluon_image_predictor_dir.zip"
EXTRACT_DIR = pathlib.Path("predictor_native")

def _load_predictor() -> MultiModalPredictor:
    if not os.path.exists(ZIP_FILENAME):
        raise FileNotFoundError(f"Missing {ZIP_FILENAME} in repo root.")
    if EXTRACT_DIR.exists():
        shutil.rmtree(EXTRACT_DIR)
    EXTRACT_DIR.mkdir(parents=True, exist_ok=True)
    with zipfile.ZipFile(ZIP_FILENAME, "r") as zf:
        zf.extractall(str(EXTRACT_DIR))
    contents = list(EXTRACT_DIR.iterdir())
    predictor_root = contents[0] if (len(contents) == 1 and contents[0].is_dir()) else EXTRACT_DIR
    return MultiModalPredictor.load(str(predictor_root))  # ← fixed

PREDICTOR = _load_predictor()
CLASS_LABELS = {i: chr(65 + i) for i in range(26)}  # 0->A ... 25->Z

def predict(pil_img: Image.Image):
    if pil_img is None:
        return {}, None
    if pil_img.mode != "RGB":
        pil_img = pil_img.convert("RGB")
    processed = pil_img.resize((224, 224))
    tmpdir = pathlib.Path(tempfile.mkdtemp())
    img_path = tmpdir / "input.png"
    processed.save(img_path)

    df = pd.DataFrame({"image": [str(img_path)]})
    proba_df = PREDICTOR.predict_proba(df)

    pretty = {}
    for col in proba_df.columns:
        label = f"Letter {CLASS_LABELS[col]}" if isinstance(col, int) and 0 <= col < 26 else str(col)
        pretty[label] = float(proba_df[col].iloc[0])
    pretty = dict(sorted(pretty.items(), key=lambda kv: kv[1], reverse=True))
    return pretty, processed

EXAMPLES = [
    ["https://www.signingsavvy.com/images/words/alphabet/2/a1.jpg"],
    ["https://www.signingsavvy.com/images/words/alphabet/2/b1.jpg"],
    ["https://www.signingsavvy.com/images/words/alphabet/2/c1.jpg"],
]

with gr.Blocks(title="Sign Language Recognition") as demo:
    gr.Markdown("# Sign Language Recognition (AutoGluon Image)")
    with gr.Row():
        with gr.Column():
            img_in = gr.Image(type="pil", label="Upload image", sources=["upload", "webcam"])
        with gr.Column():
            processed = gr.Image(type="pil", label="Preprocessed (224×224)")
    topk = gr.Label(num_top_classes=5, label="Top-5 predictions")
    img_in.change(predict, inputs=img_in, outputs=[topk, processed])
    gr.Examples(examples=EXAMPLES, inputs=[img_in], cache_examples=False)

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