import gradio as gr
import numpy as np
import cv2
import tensorflow as tf
from PIL import Image

def apply_random_mask(img):
    mask_size = 50
    y1, x1 = np.random.randint(0, img.shape[0] - mask_size, 1)[0], np.random.randint(0, img.shape[1] - mask_size, 1)[0]
    y2, x2 = y1 + mask_size, x1 + mask_size

    binary_mask = np.zeros((img.shape[0], img.shape[1]), dtype=np.uint8)
    binary_mask[y1:y2, x1:x2] = 1

    masked_img = img.copy()
    masked_img[y1:y2, x1:x2] = 0

    return masked_img

def preprocess_image(image):
    image = np.array(image)
    image = cv2.resize(image, (256, 256))
    image = apply_random_mask(image)
    image = image / 255.0
    return image

def predict(image):
    model = tf.keras.models.load_model("initial_generator_model2.keras")
    processed_image = preprocess_image(image)
    input_image = np.expand_dims(processed_image, axis=0)
    output_image = model.predict(input_image)[0]
    output_image = (output_image * 255).astype(np.uint8)
    filtered_image = cv2.medianBlur(output_image, 3)
    output_masked = (processed_image * 255).astype(np.uint8)
    return Image.fromarray(output_masked), Image.fromarray(filtered_image)

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=[gr.Image(type="pil"), gr.Image(type="pil")],
    title="Image Inpainting Model",
    description="Upload an image. A random mask will be applied, and the model will attempt to fill it in."
)

demo.launch()