TOPSInfosol commited on
Commit
395598d
·
verified ·
1 Parent(s): c57960c

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +285 -0
app.py ADDED
@@ -0,0 +1,285 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import json
3
+ import pandas as pd
4
+ import gradio as gr
5
+ import numpy as np
6
+ from PIL import Image
7
+ from theme_tops import DarkTheme
8
+ from clip_base import OpenAiClipModel
9
+ import tensorflow as tf
10
+ tagged_images = {}
11
+
12
+ MODEL_PATH = os.path.join(os.getcwd(), 'clip_tflite_model.tflite')
13
+ JSON_PATH = os.path.join(os.getcwd(), 'categories.json')
14
+
15
+ def test_model(image):
16
+ """Test the TFLite model with an uploaded image"""
17
+ try:
18
+ # Check if model and JSON files exist
19
+ if not os.path.exists(MODEL_PATH):
20
+ return "Error: Model file not found. Please generate the model first."
21
+ if not os.path.exists(JSON_PATH):
22
+ return "Error: Categories file not found. Please generate the model first."
23
+
24
+ # Load and preprocess image
25
+ processed_image = load_and_preprocess_image(image)
26
+
27
+ # Load the TFLite model
28
+ interpreter = tf.lite.Interpreter(model_path=MODEL_PATH)
29
+ interpreter.allocate_tensors()
30
+
31
+ # Get input and output details
32
+ input_details = interpreter.get_input_details()
33
+ output_details = interpreter.get_output_details()
34
+
35
+ interpreter.set_tensor(input_details[0]['index'], processed_image)
36
+
37
+ interpreter.invoke()
38
+
39
+ embeddings = interpreter.get_tensor(output_details[0]['index'])
40
+
41
+ with open(JSON_PATH, 'r') as f:
42
+ categories = json.load(f)
43
+
44
+ scores_with_ids = []
45
+ for i, score in enumerate(embeddings.flatten()):
46
+ scores_with_ids.append((float(score), i))
47
+
48
+ scores_with_ids.sort(reverse=True) # Sort by score (first element of tuple)
49
+
50
+ top_results = scores_with_ids[:5]
51
+
52
+ results = []
53
+ for score, category_id in top_results:
54
+ percentage = score * 100
55
+ category = next((cat['title'] for cat in categories if cat['id'] == category_id),
56
+ f"Category {category_id}")
57
+ results.append(f"{category}: {percentage:.2f}%")
58
+
59
+ return "\n".join(results)
60
+
61
+ except Exception as e:
62
+ return f"Error processing image: {str(e)}"
63
+
64
+ def load_and_preprocess_image(image):
65
+ """Preprocess image for model input"""
66
+ if isinstance(image, str):
67
+ image = Image.open(image)
68
+ elif isinstance(image, np.ndarray):
69
+ image = Image.fromarray(image)
70
+
71
+ image = image.resize((224, 224))
72
+ image = image.convert('RGB')
73
+ image = np.array(image).astype(np.float32) / 255.0
74
+ image = np.expand_dims(image, axis=0)
75
+ return image
76
+
77
+ def process_images(payload):
78
+ tflite_model = OpenAiClipModel(payload=payload).build_model()
79
+
80
+ return tflite_model
81
+
82
+ # Function to add a new tag category
83
+ def add_tag_category(tag_category):
84
+ # Normalize and validate tag category
85
+ tag_category = tag_category.strip()
86
+ if not tag_category:
87
+ return "Please enter a valid tag category", None
88
+
89
+ # Initialize the tag category if it doesn't exist
90
+ if tag_category not in tagged_images:
91
+ tagged_images[tag_category] = []
92
+
93
+ return f"Tag Category '{tag_category}' Added", gr.File(visible=True), ""
94
+
95
+
96
+ # Function to get updated tag category choices
97
+ def get_tag_category_choices():
98
+ return gr.Dropdown(choices=list(tagged_images.keys()))
99
+
100
+
101
+ def show_category_images(tag_category):
102
+ if not tag_category:
103
+ return None, None
104
+
105
+ if tag_category in tagged_images:
106
+ return (
107
+ gr.Gallery(value=tagged_images[tag_category]),
108
+ tagged_images[tag_category]
109
+ )
110
+ return None, None
111
+
112
+
113
+ # Function to upload images for a specific tag category
114
+ def upload_images_for_tag(tag_category, image_files):
115
+ # Ensure the tag category exists
116
+ if tag_category not in tagged_images:
117
+ return "Tag category not found. Add the tag category first.", None, None
118
+
119
+ # Replace existing images with new ones for the tag category
120
+ tagged_images[tag_category] = [file.name for file in image_files] # Replace instead of append
121
+
122
+ return (
123
+ f"Added {len(image_files)} images to '{tag_category}'",
124
+ gr.Gallery(value=[file.name for file in image_files]),
125
+ tagged_images
126
+ )
127
+
128
+
129
+
130
+ # Function to export tagged images
131
+ def export_tagged_images():
132
+ return tagged_images
133
+
134
+ def clear_uploaded_images():
135
+ return None, None
136
+
137
+ # Gradio UI
138
+ with gr.Blocks(theme=DarkTheme()) as demo:
139
+ gr.Markdown("# Clip -> Tflite - TOPS Infosolutions Pvt Ltd")
140
+ gr.Markdown("Add Classification Tags")
141
+
142
+ # Tag Category Input
143
+ with gr.Row():
144
+ tag_category_input = gr.Textbox(
145
+ label="Enter Tag Category",
146
+ placeholder="e.g., Smartphone, Laptop, Tablet"
147
+ )
148
+ # add_tag_category_btn = gr.Button("Add Tag Category")
149
+ tag_category_status = gr.Textbox(label="Action Status", interactive=False)
150
+
151
+ gr.Markdown("Images")
152
+
153
+ # Image Upload for Specific Tag
154
+ with gr.Row():
155
+ tag_category_selector = gr.Dropdown(label="Select Tag Category", choices=[])
156
+ image_upload = gr.File(
157
+ file_types=["image"],
158
+ file_count="multiple",
159
+ label="Upload Images",
160
+ visible=False
161
+ )
162
+ upload_images_btn = gr.Button("Upload Images for Category")
163
+ clear_upload_btn = gr.Button("Clear Upload")
164
+
165
+ # Image Gallery with smaller previews
166
+ image_gallery = gr.Gallery(
167
+ label="Uploaded Images",
168
+ columns=[6], # Show 4 images per row
169
+ rows=[1], # Show 2 rows
170
+ height="20",
171
+ object_fit="contain", # Maintain aspect ratio
172
+ preview=False,
173
+ show_label=False,
174
+ elem_classes="small-gallery" # Custom CSS class for additional styling
175
+ )
176
+
177
+ # Export Section
178
+ with gr.Row():
179
+ # export_btn = gr.Button("Export Tagged Images")
180
+ export_output = gr.JSON(label="Exported Tagged Images")
181
+
182
+ with gr.Row():
183
+ submit_btn = gr.Button("Process Images")
184
+
185
+
186
+ with gr.Row():
187
+ download_button_tflite = gr.File(
188
+ label="Download Tflite Model",
189
+ file_count="single",
190
+ interactive=False,
191
+ type="filepath"
192
+ )
193
+
194
+ with gr.Tab("Test Model"):
195
+ with gr.Row():
196
+ with gr.Column():
197
+ test_image = gr.Image(
198
+ label="Upload Image to Test",
199
+ type="numpy"
200
+ )
201
+ test_button = gr.Button("Test Image")
202
+
203
+ with gr.Column():
204
+ output_text = gr.Textbox(
205
+ label="Prediction Results",
206
+ lines=6,
207
+ interactive=False
208
+ )
209
+
210
+ test_button.click(
211
+ fn=test_model,
212
+ inputs=[test_image],
213
+ outputs=[output_text]
214
+ )
215
+
216
+
217
+ submit_btn.click(
218
+ fn=process_images,
219
+ inputs=[export_output],
220
+ outputs=[download_button_tflite]
221
+ )
222
+
223
+ # Add custom CSS for smaller gallery images
224
+ demo.load(js="""
225
+ function() {
226
+ const style = document.createElement('style');
227
+ style.textContent = `
228
+ .small-gallery img {
229
+ max-height: 150px !important;
230
+ width: auto !important;
231
+ object-fit: contain !important;
232
+ }
233
+ .small-gallery .grid-container {
234
+ gap: 10px !important;
235
+ }
236
+ `;
237
+ document.head.appendChild(style);
238
+ }
239
+ """)
240
+
241
+ # Functionality Connections
242
+ # Add both button click and Enter key press handlers
243
+ # add_tag_category_btn.click(
244
+ # add_tag_category,
245
+ # tag_category_input,
246
+ # [tag_category_status, image_upload, tag_category_input]
247
+ # ).then(
248
+ # get_tag_category_choices,
249
+ # None,
250
+ # tag_category_selector
251
+ # )
252
+
253
+ # Add Enter key press handler
254
+ tag_category_input.submit(
255
+ add_tag_category,
256
+ tag_category_input,
257
+ [tag_category_status, image_upload, tag_category_input]
258
+ ).then(
259
+ get_tag_category_choices,
260
+ None,
261
+ tag_category_selector
262
+ )
263
+
264
+ tag_category_selector.change(
265
+ show_category_images,
266
+ tag_category_selector,
267
+ [image_gallery, image_upload]
268
+ )
269
+
270
+ upload_images_btn.click(
271
+ upload_images_for_tag,
272
+ [tag_category_selector, image_upload],
273
+ [tag_category_status, image_gallery, export_output]
274
+ )
275
+
276
+ clear_upload_btn.click(
277
+ clear_uploaded_images,
278
+ [],
279
+ [image_upload, image_gallery]
280
+ )
281
+
282
+ # export_btn.click(export_tagged_images, None, export_output)
283
+
284
+ # Launch the app
285
+ demo.launch()