import gradio as gr
import torch
from transformers import AutoImageProcessor, AutoModelForImageClassification
from PIL import Image
import numpy as np

# Load a logo recognition model from Hugging Face
# Using a model fine-tuned for logo detection
model_name = "google/mobilenet_v2_1.0_224"  # Fallback general purpose model
processor_name = "google/mobilenet_v2_1.0_224"

try:
    # Try to load a specialized logo model if available
    # Alternative: "facebook/dino-vits16" for better image understanding
    image_processor = AutoImageProcessor.from_pretrained(processor_name)
    model = AutoModelForImageClassification.from_pretrained(model_name)
except Exception as e:
    print(f"Error loading model: {e}")
    image_processor = AutoImageProcessor.from_pretrained("google/mobilenet_v2_1.0_224")
    model = AutoModelForImageClassification.from_pretrained("google/mobilenet_v2_1.0_224")

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
model.eval()

def recognize_logo(image):
    """
    Recognize a logo from an uploaded image.
    
    Args:
        image: PIL Image object or numpy array
        
    Returns:
        Dictionary with predictions and confidence scores
    """
    if image is None:
        return "Please upload an image first."
    
    try:
        # Convert to PIL Image if necessary
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        elif not isinstance(image, Image.Image):
            image = Image.fromarray(image)
        
        # Process the image
        inputs = image_processor(images=image, return_tensors="pt").to(device)
        
        # Get predictions
        with torch.no_grad():
            outputs = model(**inputs)
        
        # Get logits and convert to probabilities
        logits = outputs.logits
        probabilities = torch.nn.functional.softmax(logits, dim=-1)
        
        # Get top predictions
        top_k = 5
        top_probs, top_indices = torch.topk(probabilities, top_k)
        
        # Format results
        results = {}
        for i, (prob, idx) in enumerate(zip(top_probs[0], top_indices[0])):
            class_name = model.config.id2label.get(idx.item(), f"Class {idx.item()}")
            confidence = float(prob.item()) * 100
            results[class_name] = f"{confidence:.2f}%"
        
        return results
    
    except Exception as e:
        return f"Error processing image: {str(e)}"

# Create Gradio interface
def create_interface():
    with gr.Blocks(title="Logo Recognition AI") as demo:
        gr.Markdown("""
        # 🎯 Logo Recognition AI
        
        Upload a logo image and let our AI identify it!
        
        This application uses state-of-the-art image recognition models from Hugging Face
        to analyze and identify logos from your images.
        """)
        
        with gr.Row():
            with gr.Column():
                gr.Markdown("### Upload Your Logo")
                image_input = gr.Image(
                    type="pil",
                    label="Logo Image",
                    show_label=True,
                    sources=["upload", "webcam"],
                    interactive=True
                )
                submit_btn = gr.Button("🔍 Recognize Logo", variant="primary", size="lg")
            
            with gr.Column():
                gr.Markdown("### Recognition Results")
                output = gr.JSON(label="Predictions")
        
        submit_btn.click(
            fn=recognize_logo,
            inputs=image_input,
            outputs=output
        )
        
        # Add examples
        gr.Markdown("### Example Logos")
        gr.Markdown("""
        Try uploading images of well-known logos such as:
        - 🍎 Apple
        - Ⓜ️ Microsoft
        - 🅶 Google
        - 📘 Facebook
        - 🐦 Twitter
        """)
    
    return demo

if __name__ == "__main__":
    interface = create_interface()
    interface.launch(share=False)