whispergogogo's picture
Upload 2 files
a397049 verified
import gradio as gr
from PIL import Image, ImageDraw
import numpy as np
import math
from pathlib import Path
from itertools import cycle
from skimage.metrics import structural_similarity as ssim
from skimage.metrics import mean_squared_error
def mosaic_colour(img_array, width: int, height: int, length: int):
"""
Apply pixelated block mosaic to an RGB image array.
Divides the image into non-overlapping rectangular blocks of
size up to 'length * length',
Computes the mean RGB value of each block,
and fills that block with the averaged color.
"""
result = img_array.copy()
m, n = math.ceil(width / length), math.ceil(height / length)
for i in range(m):
for j in range(n):
left = i * length
right = min((i + 1) * length, width)
bottom = j * length
top = min((j + 1) * length, height)
rgb_avg = img_array[bottom:top, left:right, :].mean(axis=(0, 1))
result[bottom:top, left:right, 0] = round(rgb_avg[0])
result[bottom:top, left:right, 1] = round(rgb_avg[1])
result[bottom:top, left:right, 2] = round(rgb_avg[2])
return result
def adjust_average_tone(image: Image.Image, target_mean=(128, 128, 128)) -> Image.Image:
"""
Globally shift an image's average color toward a target mean.
Performs a per-channel multiplicative scaling so that the mean of the
resulting image approximately equals `target_mean`. Values are clipped
into [0, 255].
"""
arr = np.array(image)
current_mean = arr.mean(axis=(0, 1))
scale = np.array(target_mean) / (current_mean + 1e-5)
new_arr = arr * scale
new_arr = np.clip(new_arr, 0, 255).astype(np.uint8)
return new_arr
def mosaic_tile(img_array, width: int, height: int, length: int, tile_folder):
"""
Apply a tile-based mosaic using images from tile_images folder.
The image is divided into blocks of size up to 'length * length'.
For each block, this function selects the next tile image from 'tile_images' folder
(cycled), resizes it to the block size, adjusts its global average color
to match the block's average, and places it in the output.
"""
result = img_array.copy()
tile_iter = cycle(tile_folder.glob("*.jpg"))
m, n = math.ceil(width / length), math.ceil(height / length)
for i in range(m):
for j in range(n):
left = i * length
right = min((i + 1) * length, width)
bottom = j * length
top = min((j + 1) * length, height)
rgb_avg = img_array[bottom:top, left:right, :].mean(axis=(0, 1))
tile = Image.open(next(tile_iter))
new_tile_size = (right - left, top - bottom)
resized_tile = tile.resize(new_tile_size)
tile_array = adjust_average_tone(resized_tile, target_mean=rgb_avg)
result[bottom:top, left:right, :] = tile_array
return result
def draw_grid(image: Image.Image, box_size: int, color=(0,0,0), width=1) -> Image.Image:
"""
Draws grid lines on top of an image to visualize segmentation.
"""
img_with_grid = image.copy()
draw = ImageDraw.Draw(img_with_grid)
w, h = img_with_grid.size
# vertical lines
for x in range(0, w, box_size):
draw.line([(x, 0), (x, h)], fill=color, width=width)
# horizontal lines
for y in range(0, h, box_size):
draw.line([(0, y), (w, y)], fill=color, width=width)
return img_with_grid
def image_processing(input_image, Quantization: bool, Tiles: bool, resolution, box_size):
image = Image.fromarray(input_image)
# Quantization
if Quantization:
image = image.quantize()
image = image.convert("RGB")
# Resize
if resolution != "Original":
resolutions = resolution.split('×')
width, height = int(resolutions[0]), int(resolutions[1])
new_size = (width, height)
resized_image = image.resize(new_size)
else:
width, height = image.size[0], image.size[1]
resized_image = image
# Resized image with grid
segmented_image = draw_grid(resized_image, box_size)
# Mosaic
img_array = np.array(resized_image)
if Tiles:
folder = Path("tile_images")
img_array_mosaic = mosaic_tile(img_array, width, height, box_size, folder)
else:
img_array_mosaic = mosaic_colour(img_array, width, height, box_size)
# Performance Merics
metrics = calculate_metrics(img_array, img_array_mosaic)
metrics_text = f"MSE: {metrics[0]:.2f}\nSSIM: {metrics[1]:.2f}\nPSNR: {metrics[2]:.2f}\n"
return resized_image, segmented_image, Image.fromarray(img_array_mosaic), metrics_text
def calculate_metrics(resized_image, mosaic_image):
"""Compute image quality metrics between a resized and mosaic image.
Calculates:
- MSE (Mean Squared Error)
- SSIM (Structural Similarity Index)
- PSNR (Peak Signal-to-Noise Ratio, in dB)
Args:
resized_image (np.ndarray): Reference RGB image array (uint8, H×W×3).
mosaic_image (np.ndarray): Processed RGB image array (uint8, H×W×3).
Returns:
list[float, float, float]: [mse, ssim_score, psnr_db].
"""
# MSE
mse = mean_squared_error(resized_image, mosaic_image)
# SSIM
ssim_score = ssim(resized_image, mosaic_image, channel_axis=2, data_range=255)
# PSNR (Peak Signal-to-Noise Ratio)
if mse == 0:
psnr = float('inf')
else:
psnr = 20 * np.log10(255.0 / np.sqrt(mse))
return [mse, ssim_score, psnr]
# Main
demo = gr.Interface(fn=image_processing,
inputs=[
gr.Image(label="Input Image"),
gr.Checkbox(),
gr.Checkbox(),
gr.Dropdown(
[
"Original",
"640×360",
"640×480",
"480×600",
"600×480",
"800×600",
"960×540",
"720×720",
"1024×768",
"960×960",
"1280×720",
"1024×1024"
],
label="Resolution",
info="Select the resolution you wish to use."
),
gr.Slider(1, 50, step=1, value=10, label="Box Size", info="Choose between 1 and 50"),
],
outputs=[
gr.Image(label="Resized Image"),
gr.Image(label="Segmented Image"),
gr.Image(label="Mosaic Image"),
gr.Textbox(label="Performance Metrics", lines=3)
],
examples=[
["imgs/portrait_1.jpg", True, True, "Original", 10],
["imgs/portrait_2.jpg", False, False, "480×600", 15],
["imgs/portrait_3.jpg", False, True, "720×720", 20],
["imgs/landscape_1.jpg", True, True, "640×360", 8],
["imgs/landscape_2.jpg", False, False, "960×540", 10],
["imgs/animal_1.jpg", False, True, "720×720", 10],
["imgs/animal_2.jpg", True, False, "720×720", 12],
["imgs/abstract_1.jpg", True, True, "640×360", 8],
["imgs/abstract_2.jpg", False, False, "640×480", 10],
["imgs/art_1.jpg", False, True, "640×480", 7]
])
demo.launch()