import os import json import pandas as pd import gradio as gr import numpy as np from PIL import Image from theme_tops import DarkTheme from clip_base import OpenAiClipModel import tensorflow as tf tagged_images = {} MODEL_PATH = os.path.join(os.getcwd(), 'clip_tflite_model.tflite') JSON_PATH = os.path.join(os.getcwd(), 'categories.json') def test_model(image): """Test the TFLite model with an uploaded image""" try: # Check if model and JSON files exist if not os.path.exists(MODEL_PATH): return "Error: Model file not found. Please generate the model first." if not os.path.exists(JSON_PATH): return "Error: Categories file not found. Please generate the model first." # Load and preprocess image processed_image = load_and_preprocess_image(image) # Load the TFLite model interpreter = tf.lite.Interpreter(model_path=MODEL_PATH) interpreter.allocate_tensors() # Get input and output details input_details = interpreter.get_input_details() output_details = interpreter.get_output_details() interpreter.set_tensor(input_details[0]['index'], processed_image) interpreter.invoke() embeddings = interpreter.get_tensor(output_details[0]['index']) with open(JSON_PATH, 'r') as f: categories = json.load(f) scores_with_ids = [] for i, score in enumerate(embeddings.flatten()): scores_with_ids.append((float(score), i)) scores_with_ids.sort(reverse=True) # Sort by score (first element of tuple) top_results = scores_with_ids[:5] results = [] for score, category_id in top_results: percentage = score * 100 category = next((cat['title'] for cat in categories if cat['id'] == category_id), f"Category {category_id}") results.append(f"{category}: {percentage:.2f}%") return "\n".join(results) except Exception as e: return f"Error processing image: {str(e)}" def load_and_preprocess_image(image): """Preprocess image for model input""" if isinstance(image, str): image = Image.open(image) elif isinstance(image, np.ndarray): image = Image.fromarray(image) image = image.resize((224, 224)) image = image.convert('RGB') image = np.array(image).astype(np.float32) / 255.0 image = np.expand_dims(image, axis=0) return image def process_images(payload): tflite_model = OpenAiClipModel(payload=payload).build_model() return tflite_model # Function to add a new tag category def add_tag_category(tag_category): # Normalize and validate tag category tag_category = tag_category.strip() if not tag_category: return "Please enter a valid tag category", None # Initialize the tag category if it doesn't exist if tag_category not in tagged_images: tagged_images[tag_category] = [] return f"Tag Category '{tag_category}' Added", gr.File(visible=True), "" # Function to get updated tag category choices def get_tag_category_choices(): return gr.Dropdown(choices=list(tagged_images.keys())) def show_category_images(tag_category): if not tag_category: return None, None if tag_category in tagged_images: return ( gr.Gallery(value=tagged_images[tag_category]), tagged_images[tag_category] ) return None, None # Function to upload images for a specific tag category def upload_images_for_tag(tag_category, image_files): # Ensure the tag category exists if tag_category not in tagged_images: return "Tag category not found. Add the tag category first.", None, None # Replace existing images with new ones for the tag category tagged_images[tag_category] = [file.name for file in image_files] # Replace instead of append return ( f"Added {len(image_files)} images to '{tag_category}'", gr.Gallery(value=[file.name for file in image_files]), tagged_images ) # Function to export tagged images def export_tagged_images(): return tagged_images def clear_uploaded_images(): return None, None # Gradio UI with gr.Blocks(theme=DarkTheme()) as demo: gr.Markdown("# Clip -> Tflite - TOPS Infosolutions Pvt Ltd") gr.Markdown("Add Classification Tags") # Tag Category Input with gr.Row(): tag_category_input = gr.Textbox( label="Enter Tag Category", placeholder="e.g., Smartphone, Laptop, Tablet" ) # add_tag_category_btn = gr.Button("Add Tag Category") tag_category_status = gr.Textbox(label="Action Status", interactive=False) gr.Markdown("Images") # Image Upload for Specific Tag with gr.Row(): tag_category_selector = gr.Dropdown(label="Select Tag Category", choices=[]) image_upload = gr.File( file_types=["image"], file_count="multiple", label="Upload Images", visible=False ) upload_images_btn = gr.Button("Upload Images for Category") clear_upload_btn = gr.Button("Clear Upload") # Image Gallery with smaller previews image_gallery = gr.Gallery( label="Uploaded Images", columns=[6], # Show 4 images per row rows=[1], # Show 2 rows height="20", object_fit="contain", # Maintain aspect ratio preview=False, show_label=False, elem_classes="small-gallery" # Custom CSS class for additional styling ) # Export Section with gr.Row(): # export_btn = gr.Button("Export Tagged Images") export_output = gr.JSON(label="Exported Tagged Images") with gr.Row(): submit_btn = gr.Button("Process Images") with gr.Row(): download_button_tflite = gr.File( label="Download Tflite Model", file_count="single", interactive=False, type="filepath" ) with gr.Tab("Test Model"): with gr.Row(): with gr.Column(): test_image = gr.Image( label="Upload Image to Test", type="numpy" ) test_button = gr.Button("Test Image") with gr.Column(): output_text = gr.Textbox( label="Prediction Results", lines=6, interactive=False ) test_button.click( fn=test_model, inputs=[test_image], outputs=[output_text] ) submit_btn.click( fn=process_images, inputs=[export_output], outputs=[download_button_tflite] ) # Add custom CSS for smaller gallery images demo.load(js=""" function() { const style = document.createElement('style'); style.textContent = ` .small-gallery img { max-height: 150px !important; width: auto !important; object-fit: contain !important; } .small-gallery .grid-container { gap: 10px !important; } `; document.head.appendChild(style); } """) # Functionality Connections # Add both button click and Enter key press handlers # add_tag_category_btn.click( # add_tag_category, # tag_category_input, # [tag_category_status, image_upload, tag_category_input] # ).then( # get_tag_category_choices, # None, # tag_category_selector # ) # Add Enter key press handler tag_category_input.submit( add_tag_category, tag_category_input, [tag_category_status, image_upload, tag_category_input] ).then( get_tag_category_choices, None, tag_category_selector ) tag_category_selector.change( show_category_images, tag_category_selector, [image_gallery, image_upload] ) upload_images_btn.click( upload_images_for_tag, [tag_category_selector, image_upload], [tag_category_status, image_gallery, export_output] ) clear_upload_btn.click( clear_uploaded_images, [], [image_upload, image_gallery] ) # export_btn.click(export_tagged_images, None, export_output) # Launch the app demo.launch()