import random
import gradio as gr
import numpy as np
from PIL import Image
from datasets import load_dataset
from tensorflow.keras.models import load_model

# Load pre-trained U-Net model-
model = load_model("unet_model.h5", compile=False)

# Load Hugging Face dataset
dataset = load_dataset("AIOmarRehan/Cropped_Yale_Faces")

# Preprocess function
def preprocess_image(image, target_size=(192, 176)):
    image = image.resize((target_size[1], target_size[0]))  # width, height
    image = np.array(image) / 255.0
    if image.ndim == 2:
        image = np.expand_dims(image, axis=-1)
    return np.expand_dims(image, axis=0)

# Salt-and-pepper noise function
def add_salt_and_pepper_noise(image, amount=0.05):
    """
    image: PIL Image in grayscale ('L') or RGB
    amount: fraction of pixels to corrupt
    """
    img_array = np.array(image)
    # Salt noise
    num_salt = np.ceil(amount * img_array.size * 0.5)
    coords = [np.random.randint(0, i - 1, int(num_salt)) for i in img_array.shape]
    img_array[tuple(coords)] = 255
    # Pepper noise
    num_pepper = np.ceil(amount * img_array.size * 0.5)
    coords = [np.random.randint(0, i - 1, int(num_pepper)) for i in img_array.shape]
    img_array[tuple(coords)] = 0
    return Image.fromarray(img_array)

# Prediction function
def predict(img=None, use_dataset=False, add_noise=False):
    if use_dataset:
        # Pick random image from dataset
        example = random.choice(dataset["train"])
        img = example["image"]

    if img is None:
        return None, None

    img = img.convert("L")

    noisy_img = img
    if add_noise:
        noisy_img = add_salt_and_pepper_noise(img)

    input_data = preprocess_image(noisy_img)

    pred = model.predict(input_data)[0]

    if pred.ndim == 3 and pred.shape[-1] == 1:
        pred = np.squeeze(pred, axis=-1)

    denoised_img = (pred * 255).astype(np.uint8)
    denoised_img = Image.fromarray(denoised_img)

    return noisy_img, denoised_img

# Gradio Interface
interface = gr.Interface(
    fn=predict,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Checkbox(label="Use Random Dataset Image"),
        gr.Checkbox(label="Add Salt-and-Pepper Noise")
    ],
    outputs=[
        gr.Image(type="pil", label="Noisy Input Image"),
        gr.Image(type="pil", label="Denoised Output Image")
    ],
    title="U-Net Image Denoising with Salt-and-Pepper Noise",
    description="Upload an image or pick a random image from the Cropped Yale Faces dataset. "
                "Optionally add salt-and-pepper noise to the image before denoising."
)

# Launch the app
interface.launch()