File size: 2,631 Bytes
62c6e9b 5d3b587 db56122 62c6e9b 5d3b587 0bbb029 5d3b587 db56122 5d3b587 0bbb029 db56122 0bbb029 db56122 0bbb029 db56122 0bbb029 db56122 5d3b587 0bbb029 5d3b587 0bbb029 5d3b587 0bbb029 5d3b587 db56122 5d3b587 db56122 0bbb029 db56122 0bbb029 5d3b587 db56122 5d3b587 |
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 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
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() |