import spaces
import gradio as gr
import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
from huggingface_hub import hf_hub_download

# Order matches full_train.classes from the training notebook (ImageFolder's
# alphabetical folder sort) — do not reorder without re-checking the notebook.
CLASS_NAMES = [
    "Black_Beans", "Black_EyedPea", "Brown_Chickpea", "Desi_chickpea",
    "Horse_Gram", "Husked_BlackGram", "Husked_RedLentil", "Kidney_Bean",
    "Red_Beans", "Red_KidneyBean", "Red_chickpea", "Soya_Bean",
    "White_Chickpea", "White_beans", "White_peabeans", "Whole_BlackGram",
    "Whole_GreenGram", "Whole_TurkishGram", "chana1", "chana2",
]

# Weights live in a separate Model repo (not bundled in this Space), so the
# Space repo stays small and never needs Git LFS. Downloaded once per
# container start and cached.
WEIGHTS_PATH = hf_hub_download(
    repo_id="silversur4/indian-lentils-image-classifier",
    filename="resnet_lentils.pth",
)

# ZeroGPU only grants GPU access inside a function decorated with
# @spaces.GPU — no CUDA context exists at import time, so the model is
# loaded on CPU here and moved to CUDA per-call inside predict().
model = models.resnet18(weights=None)
model.fc = nn.Linear(model.fc.in_features, len(CLASS_NAMES))
model.load_state_dict(torch.load(WEIGHTS_PATH, map_location="cpu"))
model.eval()

# Same eval_transform used during evaluation in the training notebook.
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])


@spaces.GPU
def predict(image: Image.Image):
    if image is None:
        return {}
    device = torch.device("cuda")
    model.to(device)
    x = transform(image.convert("RGB")).unsqueeze(0).to(device)
    with torch.no_grad():
        probs = torch.softmax(model(x), dim=1)[0]
    return {CLASS_NAMES[i]: float(probs[i]) for i in range(len(CLASS_NAMES))}


demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", sources=["upload", "webcam"], label="Photo of a lentil / bean / chickpea"),
    outputs=gr.Label(num_top_classes=5, label="Prediction"),
    title="Indian Lentils Classifier",
    description=(
        "ResNet18, fine-tuned on the Indian Lentils dataset (20 classes). "
        "Upload a photo or use your camera to classify a lentil, bean, or chickpea variety."
    ),
)

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