import gradio as gr
import requests
from PIL import Image
from io import BytesIO
from transformers import pipeline

# Adjust these if your model's order is actually different.
# For example, if your dataset folders are named (alphabetically): 
#   bumblebee, honeybee, vespidae,
# then 0 => bumblebee, 1 => honeybee, 2 => vespidae (the default PyTorch order).
# Verify your label indices by printing `test_dataset.classes` in your training script.

label_map = {
    "LABEL_0": "bumblebee",
    "LABEL_1": "honeybee",
    "LABEL_2": "vespidae"
}

model_id = "Honey-Bee-Society/honeybee_bumblebee_vespidae_resnet50"
classifier = pipeline("image-classification", model=model_id)

def classify_image_from_url(image_url: str):
    """
    Downloads an image from a public URL and runs it through
    the ResNet-50 image-classification pipeline, returning the top predictions.
    """
    try:
        # Fetch the image
        response = requests.get(image_url)
        response.raise_for_status()
        image = Image.open(BytesIO(response.content)).convert("RGB")
        
        # Run inference
        results = classifier(image)

        # 1) Post-process labels
        # 2) Format scores to remove scientific notation
        for r in results:
            # Map from "LABEL_x" to your real class name
            if r["label"] in label_map:
                r["label"] = label_map[r["label"]]
            # Format score with, e.g., 8 decimal places to avoid scientific notation
            r["score"] = float(f"{r['score']:.8f}") 

        return results
    
    except Exception as e:
        return {"error": str(e)}

demo = gr.Interface(
    fn=classify_image_from_url,
    inputs=gr.Textbox(lines=1, label="Image URL"),
    outputs="json",
    title="ResNet-50 Image Classifier",
    description="Enter a public image URL to get top predictions with custom labels."
)

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