app.py

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()