import gradio as gr
import numpy as np
import tensorflow as tf
import cv2

# Load your trained Keras model
model = tf.keras.models.load_model("unet_mask_segmentation.h5")

# Image preprocessing function (same as used during training)
def preprocess_image(img):
    img_resized = cv2.resize(img, (256, 256))
    img_normalized = img_resized / 255.0  # Normalize to 0-1
    return img_normalized

# Prediction and overlay function
def predict(input_img):
    # Ensure image is RGB and numpy array
    img = np.array(input_img.convert("RGB"))

    # Preprocess
    preprocessed_img = preprocess_image(img)
    input_tensor = np.expand_dims(preprocessed_img, axis=0)  # Add batch dimension

    # Model prediction
    prediction = model.predict(input_tensor)[0]  # Remove batch dim

    # Post-processing mask
    mask = (prediction > 0.5).astype(np.uint8)  # Binary mask
    mask_resized = cv2.resize(mask, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_NEAREST)

    # Create overlay
    overlay = img.astype(np.float32) / 255.0  # Normalize input image
    alpha = 0.5  # Transparency of overlay

    # Create red mask in RGB format
    red_mask = np.zeros_like(overlay)
    red_mask[:, :, 0] = mask_resized  # Red channel

    # Alpha blend original image with red mask
    blended = (1 - alpha) * overlay + alpha * red_mask
    blended = np.clip(blended * 255, 0, 255).astype(np.uint8)

    return blended

# Gradio interface
interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload Image"),
    outputs=gr.Image(type="numpy", label="Segmented Image"),
    title="Image Segmentation App",
    description="Upload an image and get the segmentation mask overlay using your trained model."
)

# Launch Gradio app (enable public link for Hugging Face Spaces)
interface.launch(share=True)