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

# Load model
model = load_model("unet_model.h5", compile=False)

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

# Prediction function for Gradio
def predict(img):
    # Convert to grayscale
    img = img.convert("L")

    # Preprocess
    input_data = preprocess_image(img)

    # Model prediction
    pred = model.predict(input_data)[0]

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

    # Convert to image
    pred_img = (pred * 255).astype(np.uint8)
    pred_img = Image.fromarray(pred_img)

    return pred_img

# Gradio UI
interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=gr.Image(type="pil", label="Denoised Output"),
    title="U-Net Image Denoising",
    description="Upload a grayscale image and get the denoised result using a U-Net model."
)

interface.launch()