import gradio as gr
from transformers import pipeline

# Function for image classification
def classify(image, model_name):
    try:
        # Load the pipeline with the given model name
        pipe = pipeline("image-classification", model=model_name)
        # Perform image classification
        results = pipe(image)
        return {result["label"]: round(result["score"], 2) for result in results}
    except Exception as e:
        # Handle errors gracefully, e.g., invalid model names
        return {"Error": str(e)}

# Gradio Interface
demo = gr.Interface(
    fn=classify,
    inputs=[
        gr.Image(type="pil", label="Upload an Image"),
        gr.Textbox(label="Enter timm Model Name", placeholder="e.g., timm/mobilenetv3_large_100.ra_in1k"),
    ],
    outputs=gr.Label(num_top_classes=3, label="Top Predictions"),
    title="Custom timm Model Image Classifier",
    description="Enter a timm model name from Hugging Face, upload an image, and get predictions.",
    examples=[
        ["cat.png", "timm/mobilenetv3_small_100.lamb_in1k"],
        ["cat.png", "timm/resnet50.a1_in1k"],
    ],
)

demo.launch()