import gradio as gr
import numpy as np
import cv2
from PIL import Image
import tensorflow as tf
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input, decode_predictions
import json

# Load MobileNetV2 model
model = MobileNetV2(weights='imagenet')

def add_hsv_noise(image, hue_noise=0, saturation_noise=0, value_noise=0):
    """Add HSV noise to an image"""
    if image is None:
        return None

    # Convert PIL to numpy array
    img_array = np.array(image)

    # Convert RGB to HSV
    hsv = cv2.cvtColor(img_array, cv2.COLOR_RGB2HSV).astype(np.float32)

    # Add noise to each channel
    hsv[:, :, 0] = np.clip(hsv[:, :, 0] + hue_noise, 0, 179)  # Hue: 0-179
    hsv[:, :, 1] = np.clip(hsv[:, :, 1] + saturation_noise, 0, 255)  # Saturation: 0-255
    hsv[:, :, 2] = np.clip(hsv[:, :, 2] + value_noise, 0, 255)  # Value: 0-255

    # Convert back to RGB
    rgb = cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2RGB)

    return Image.fromarray(rgb)

def predict_image(image, top_n, hue_noise, saturation_noise, value_noise):
    """Predict image classes with noise applied"""
    if image is None:
        return None, "Please upload an image first."

    # Apply HSV noise
    noisy_image = add_hsv_noise(image, hue_noise, saturation_noise, value_noise)

    # Preprocess for MobileNet
    img_resized = noisy_image.resize((224, 224))
    img_array = np.array(img_resized)
    img_array = np.expand_dims(img_array, axis=0)
    img_array = preprocess_input(img_array)

    # Make prediction
    predictions = model.predict(img_array)
    decoded_predictions = decode_predictions(predictions, top=top_n)[0]

    # Format results
    results = []
    for i, (class_id, class_name, probability) in enumerate(decoded_predictions):
        results.append(f"{i+1}. {class_name}: {probability:.4f} ({probability*100:.2f}%)")

    results_text = "\n".join(results)

    return noisy_image, results_text

# Create Gradio interface
with gr.Blocks(title="MobileNet HSV Noise Analysis") as demo:
    gr.Markdown("# MobileNet Classification with HSV Noise")
    gr.Markdown("Upload an image and adjust HSV noise sliders to see how it affects MobileNet predictions.")

    with gr.Row():
        with gr.Column():
            # Input controls
            input_image = gr.Image(type="pil", label="Upload Image")
            top_n = gr.Slider(minimum=1, maximum=10, value=5, step=1, label="Top N Classes")

            gr.Markdown("### HSV Noise Controls")
            hue_noise = gr.Slider(minimum=-50, maximum=50, value=0, step=1, label="Hue Noise (-50 to 50)")
            saturation_noise = gr.Slider(minimum=-100, maximum=100, value=0, step=5, label="Saturation Noise (-100 to 100)")
            value_noise = gr.Slider(minimum=-100, maximum=100, value=0, step=5, label="Value/Brightness Noise (-100 to 100)")

        with gr.Column():
            # Output displays
            output_image = gr.Image(label="Image with Noise Applied")
            predictions_text = gr.Textbox(label="Top Predictions", lines=10, max_lines=15)

    # Set up real-time updates
    inputs = [input_image, top_n, hue_noise, saturation_noise, value_noise]
    outputs = [output_image, predictions_text]

    # Update predictions when any input changes
    for input_component in inputs:
        input_component.change(
            fn=predict_image,
            inputs=inputs,
            outputs=outputs
        )

# Launch the app
if __name__ == "__main__":
    demo.launch(share=True, debug=True)