import gradio as gr
from transformers import AutoImageProcessor, SiglipForImageClassification
from PIL import Image
import torch

model_name = "prithivMLmods/Dog-Breed-120"
processor = AutoImageProcessor.from_pretrained(model_name)
model = SiglipForImageClassification.from_pretrained(model_name)
model.eval()

def predict(image):
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    logits = outputs.logits
    probs = torch.nn.functional.softmax(logits, dim=1).squeeze()
    top_idx = probs.argmax().item()
    label = model.config.id2label[top_idx]
    confidence = probs[top_idx].item()
    return f"{label} ({round(confidence*100, 1)}%)"

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="Dog Breed Classifier 🐶",
    description="Upload a dog photo to identify its breed using the prithivMLmods/Dog-Breed-120 model."
)

demo.launch()