app.py

# app.py
import os
import shutil
import zipfile
import pathlib
import tempfile
import mimetypes

import gradio as gr
import pandas as pd
from PIL import Image, UnidentifiedImageError
import huggingface_hub
from autogluon.multimodal import MultiModalPredictor


# -----------------------------
# Config / constants
# -----------------------------
MODEL_REPO_ID = "ddecosmo/nn_automl_model"           # Hugging Face repo id (model)
ZIP_FILENAME  = "autogluon_image_predictor_dir.zip"  # the zipped predictor file in the repo
HF_TOKEN      = os.getenv("HF_TOKEN", None)          # pulled from HF Space secret (optional)

# Local cache/extract dirs inside the Space
CACHE_DIR   = pathlib.Path("hf_assets")
EXTRACT_DIR = CACHE_DIR / "predictor_native"

# App limits / UI config
ALLOWED_MIME = {"image/png", "image/jpeg"}
ALLOWED_EXTS = {".png", ".jpg", ".jpeg"}
MAX_BYTES    = 5 * 1024 * 1024
TARGET       = 224  # model input size


# -----------------------------
# Model setup
# -----------------------------
def _prepare_predictor_dir() -> str:
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    local_zip = huggingface_hub.hf_hub_download(
        repo_id=MODEL_REPO_ID,
        filename=ZIP_FILENAME,
        repo_type="model",
        token=HF_TOKEN,                     # works if model is private; otherwise ignored
        local_dir=str(CACHE_DIR),
        local_dir_use_symlinks=False,
    )

    if EXTRACT_DIR.exists():
        shutil.rmtree(EXTRACT_DIR)
    EXTRACT_DIR.mkdir(parents=True, exist_ok=True)

    with zipfile.ZipFile(local_zip, "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 str(predictor_root)


PREDICTOR_DIR = _prepare_predictor_dir()
PREDICTOR = MultiModalPredictor.load(PREDICTOR_DIR)

# If you ever need explicit mapping, keep here (not strictly required if predictor already knows classes)
CLASS_LABELS = {0: "Western", 1: "Asian"}


# -----------------------------
# Helpers
# -----------------------------
def _human_filesize(n: int) -> str:
    for unit in ["B","KB","MB","GB","TB"]:
        if n < 1024.0:
            return f"{n:.1f} {unit}"
        n /= 1024.0
    return f"{n:.1f} PB"


def _preprocess(img: Image.Image) -> Image.Image:
    w, h = img.size
    if min(w, h) == 0:
        return Image.new("RGB", (TARGET, TARGET), (0, 0, 0))
    scale = TARGET / min(w, h)
    nw, nh = int(round(w * scale)), int(round(h * scale))
    img = img.resize((nw, nh), Image.BILINEAR)
    left = (img.width  - TARGET) // 2
    top  = (img.height - TARGET) // 2
    return img.crop((left, top, left + TARGET, top + TARGET))


# -----------------------------
# Inference function (filepath input)
# -----------------------------
def process_and_predict(file_path: str):
    # basic validation
    if not file_path:
        gr.Warning("Please upload an image.")
        return "Please upload an image.", None, None, {}

    try:
        size_bytes = os.path.getsize(file_path)
        if size_bytes > MAX_BYTES:
            gr.Error(f"File too large (max {_human_filesize(MAX_BYTES)}).")
            return "File too large.", None, None, {}
    except Exception:
        # ok to continue (webcam/temp paths may not resolve size cleanly)
        pass

    ext = os.path.splitext(file_path)[1].lower()
    mime = (mimetypes.guess_type(file_path)[0] or "").lower()
    if (ext not in ALLOWED_EXTS) or (mime not in ALLOWED_MIME):
        gr.Error("Unsupported file type. Please use PNG or JPG/JPEG.")
        return "Unsupported file type.", None, None, {}

    # open image safely
    try:
        with Image.open(file_path) as im:
            im.verify()
        original = Image.open(file_path).convert("RGB")
    except UnidentifiedImageError:
        return "Invalid image file. Please upload a real image.", None, None, {}
    except Exception:
        return "Could not process the file. Is it a valid image?", None, None, {}

    # preprocess (what the model sees)
    processed = _preprocess(original)

    # predict using processed image
    tmpdir = pathlib.Path(tempfile.mkdtemp())
    proc_path = tmpdir / "model_input.png"
    processed.save(proc_path)

    df = pd.DataFrame({"image": [str(proc_path)]})
    proba_df = PREDICTOR.predict_proba(df).rename(columns={0: "Western (0)", 1: "Asian (1)"})
    row = proba_df.iloc[0]
    pretty = {
        "Western": float(row.get("Western (0)", 0.0)),
        "Asian":   float(row.get("Asian (1)", 0.0)),
    }

    gr.Info("Image accepted.")
    return "Image loaded successfully.", original, processed, pretty


# -----------------------------
# UI
# -----------------------------
EXAMPLES = [
    ["examples/paella.jpg"],
    ["examples/mapo_tofu.jpg"],
    ["examples/kimchi.jpg"],
]

with gr.Blocks() as demo:
    gr.Markdown("# Asian Food or Western Food?")
    gr.Markdown(
        "This Space demonstrates serving an AutoGluon image classifier in Gradio. "
        "Upload a **PNG/JPG** (≤ 5 MB) or use your **webcam**. "
        "You’ll see the original image and the preprocessed input the model sees."
    )
    gr.Markdown(
        """
        **Model:** [ddecosmo/nn_automl_model](https://huggingface.co/ddecosmo/nn_automl_model)
        
        **Library:** [AutoGluon TabularPredictor docs](https://auto.gluon.ai/dev/api/autogluon.tabular.TabularPredictor.html)
        **Dataset:** [maryzhang/hw1-24679-image-dataset](https://huggingface.co/datasets/maryzhang/hw1-24679-image-dataset)
        """
    )

    with gr.Row():
        with gr.Column():
            image_in = gr.Image(
                type="filepath",
                label="Input image",
                sources=["upload", "webcam"],
                show_label=True,
                height=300,
            )
            gr.Examples(
                examples=[e[0] for e in EXAMPLES],
                inputs=image_in,
                label="Representative examples",
                examples_per_page=8,
                cache_examples=False,
            )

        with gr.Column():
            status = gr.Markdown()
            with gr.Row():
                original_out  = gr.Image(label="Original", height=300)
                processed_out = gr.Image(label=f"Preprocessed ({TARGET}x{TARGET})", height=300)
            proba_pretty = gr.Label(num_top_classes=2, label="Class probabilities")

    image_in.change(
        fn=process_and_predict,
        inputs=[image_in],
        outputs=[status, original_out, processed_out, proba_pretty],
        show_progress="minimal",
    )

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