Spaces:
Sleeping
Sleeping
File size: 1,464 Bytes
b25c611 6ee7a0c 5541eee 6ee7a0c b25c611 02fa2c5 b25c611 5541eee b25c611 dc81c9f 329be7c 3dd0570 b25c611 329be7c b25c611 6ee7a0c b25c611 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
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()
|