jonloporto commited on
Commit
4366827
Β·
verified Β·
1 Parent(s): 55ea16f

Upload 2 files

Browse files
Files changed (2) hide show
  1. app.py +122 -0
  2. requirements.txt +7 -0
app.py ADDED
@@ -0,0 +1,122 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from transformers import AutoImageProcessor, AutoModelForImageClassification
4
+ from PIL import Image
5
+ import numpy as np
6
+
7
+ # Load a logo recognition model from Hugging Face
8
+ # Using a model fine-tuned for logo detection
9
+ model_name = "google/mobilenet_v2_1.0_224" # Fallback general purpose model
10
+ processor_name = "google/mobilenet_v2_1.0_224"
11
+
12
+ try:
13
+ # Try to load a specialized logo model if available
14
+ # Alternative: "facebook/dino-vits16" for better image understanding
15
+ image_processor = AutoImageProcessor.from_pretrained(processor_name)
16
+ model = AutoModelForImageClassification.from_pretrained(model_name)
17
+ except Exception as e:
18
+ print(f"Error loading model: {e}")
19
+ image_processor = AutoImageProcessor.from_pretrained("google/mobilenet_v2_1.0_224")
20
+ model = AutoModelForImageClassification.from_pretrained("google/mobilenet_v2_1.0_224")
21
+
22
+ device = "cuda" if torch.cuda.is_available() else "cpu"
23
+ model.to(device)
24
+ model.eval()
25
+
26
+ def recognize_logo(image):
27
+ """
28
+ Recognize a logo from an uploaded image.
29
+
30
+ Args:
31
+ image: PIL Image object or numpy array
32
+
33
+ Returns:
34
+ Dictionary with predictions and confidence scores
35
+ """
36
+ if image is None:
37
+ return "Please upload an image first."
38
+
39
+ try:
40
+ # Convert to PIL Image if necessary
41
+ if isinstance(image, np.ndarray):
42
+ image = Image.fromarray(image)
43
+ elif not isinstance(image, Image.Image):
44
+ image = Image.fromarray(image)
45
+
46
+ # Process the image
47
+ inputs = image_processor(images=image, return_tensors="pt").to(device)
48
+
49
+ # Get predictions
50
+ with torch.no_grad():
51
+ outputs = model(**inputs)
52
+
53
+ # Get logits and convert to probabilities
54
+ logits = outputs.logits
55
+ probabilities = torch.nn.functional.softmax(logits, dim=-1)
56
+
57
+ # Get top predictions
58
+ top_k = 5
59
+ top_probs, top_indices = torch.topk(probabilities, top_k)
60
+
61
+ # Format results
62
+ results = {}
63
+ for i, (prob, idx) in enumerate(zip(top_probs[0], top_indices[0])):
64
+ class_name = model.config.id2label.get(idx.item(), f"Class {idx.item()}")
65
+ confidence = float(prob.item()) * 100
66
+ results[class_name] = f"{confidence:.2f}%"
67
+
68
+ return results
69
+
70
+ except Exception as e:
71
+ return f"Error processing image: {str(e)}"
72
+
73
+ # Create Gradio interface
74
+ def create_interface():
75
+ with gr.Blocks(title="Logo Recognition AI") as demo:
76
+ gr.Markdown("""
77
+ # 🎯 Logo Recognition AI
78
+
79
+ Upload a logo image and let our AI identify it!
80
+
81
+ This application uses state-of-the-art image recognition models from Hugging Face
82
+ to analyze and identify logos from your images.
83
+ """)
84
+
85
+ with gr.Row():
86
+ with gr.Column():
87
+ gr.Markdown("### Upload Your Logo")
88
+ image_input = gr.Image(
89
+ type="pil",
90
+ label="Logo Image",
91
+ show_label=True,
92
+ sources=["upload", "webcam"],
93
+ interactive=True
94
+ )
95
+ submit_btn = gr.Button("πŸ” Recognize Logo", variant="primary", size="lg")
96
+
97
+ with gr.Column():
98
+ gr.Markdown("### Recognition Results")
99
+ output = gr.JSON(label="Predictions")
100
+
101
+ submit_btn.click(
102
+ fn=recognize_logo,
103
+ inputs=image_input,
104
+ outputs=output
105
+ )
106
+
107
+ # Add examples
108
+ gr.Markdown("### Example Logos")
109
+ gr.Markdown("""
110
+ Try uploading images of well-known logos such as:
111
+ - 🍎 Apple
112
+ - Ⓜ️ Microsoft
113
+ - πŸ…Ά Google
114
+ - πŸ“˜ Facebook
115
+ - 🐦 Twitter
116
+ """)
117
+
118
+ return demo
119
+
120
+ if __name__ == "__main__":
121
+ interface = create_interface()
122
+ interface.launch(share=False)
requirements.txt ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ gradio==4.26.0
2
+ torch==2.1.2
3
+ torchvision==0.16.2
4
+ transformers==4.41.2
5
+ Pillow==10.1.0
6
+ numpy==1.24.3
7
+ huggingface-hub==1.3.1