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

from transformers import pipeline

checkpoint = "openai/clip-vit-large-patch14"
detector = pipeline(model=checkpoint, task="zero-shot-image-classification")
# Function to predict dog category
def predict_dog_category(image):
    # List of dog categories
    dog_category = [
        'Siberian Husky', 'Boxer',  # Working Dogs
        'Border Collie', 'Australian Shepherd',  # Herding Dogs
        'Chihuahua', 'Pomeranian',  # Toy Dogs
        'Labrador Retriever', 'Golden Retriever',  # Sporting Dogs
        'Yorkshire Terrier', 'Bull Terrier',  # Terriers
        'Bulldog', 'Poodle'  # Non-Sporting Dogs
    ]
    
    # Use CLIP model to predict dog category
    predictions = detector(image, candidate_labels=dog_category)
    return {predictions[i]['label']: float(predictions[i]['score']) for i in range(len(predictions))}

# Create Gradio interface
iface = gr.Interface(
    fn=predict_dog_category,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=12)
)
iface.launch(share=True)