Spaces:
Sleeping
Sleeping
File size: 4,125 Bytes
4366827 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | 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)
|