app.py

import gradio as gr
import cv2
import numpy as np
from PIL import Image
import os
import time
from logic.imgPreprocess import resize_img, color_quantize
from logic.imgGrid import segment_image_grid
from logic.tileMapping import load_tile_images, map_tiles_to_grid
from logic.perfMetric import mse, ssim_metric, timed, create_performance_report, run_performance_benchmark
from logic.validation import (
    ValidationError,
    ensure_grid_divisibility,
    ensure_image,
    ensure_positive_int,
)

DEFAULT_IMAGE_SIZE = 512
BENCHMARK_IMAGE_PATH = "data/test_images/marten.jpg"
BENCHMARK_IMAGE_SIZES = [256, 512, 1024]
BENCHMARK_GRID_SIZES = [16, 32, 64]

def preprocess_and_mosaic(
    image: Image.Image,
    grid_size: int,
    n_colors: int,
    image_size: int = DEFAULT_IMAGE_SIZE,
):
    """Generate a mosaic image and accompanying diagnostic data.

    Args:
        image (PIL.Image.Image): User supplied image.
        grid_size (int): Requested grid dimension (NxN).
        n_colors (int): Number of quantized colors.
        image_size (int, optional): Target square size. Defaults to ``DEFAULT_IMAGE_SIZE``.

    Returns:
        tuple[PIL.Image.Image, PIL.Image.Image, PIL.Image.Image, str]:
            Original display, segmentation preview, mosaic result, and textual metrics.
    """
    start_total = time.time()
    try:
        validated_image = ensure_image(image)
        ensure_positive_int(n_colors, "Color palette size", minimum=2, maximum=64)
        ensure_positive_int(grid_size, "Grid size", minimum=2, maximum=512)
        ensure_positive_int(image_size, "Image size", minimum=64, maximum=2048)
        divisible = ensure_grid_divisibility(image_size, grid_size)
    except (ValidationError, ValueError) as exc:
        raise gr.Error(str(exc)) from exc
    if not divisible:
        gr.Warning(
            "Image size does not divide evenly by the grid. The app will crop to the nearest valid region."
        )

    # Step 1: Image preprocessing
    start_preprocess = time.time()
    resized_img = resize_img(validated_image, image_size)
    quantized_img, color_centers = color_quantize(resized_img, n_colors)
    preprocess_time = time.time() - start_preprocess

    # Step 2: Image grid processing
    resized_np = np.array(quantized_img)
    if resized_np.shape[-1] == 3:
        resized_np = cv2.cvtColor(resized_np, cv2.COLOR_RGB2BGR)

    cell_size = max(1, image_size // grid_size)
    actual_grid_size = image_size // cell_size
    actual_image_size = actual_grid_size * cell_size
    
    # Crop the image to match the actual grid dimensions
    resized_np_cropped = resized_np[:actual_image_size, :actual_image_size]
    
    # Step 2: Grid segmentation (timed)
    grid_labels, seg_time = timed(segment_image_grid)(resized_np_cropped, actual_grid_size, color_centers)
    
    # Step 3: Tile mapping
    # Step 3a: Load tiles (timed)
    start_tile_load = time.time()
    try:
        tiles, tile_colors = load_tile_images('data/tiles', n_colors, cell_size)
    except (FileNotFoundError, ValidationError) as exc:
        raise gr.Error(str(exc)) from exc
    tile_load_time = time.time() - start_tile_load
    
    # Step 3b: Map tiles to grid (timed)
    mosaic_img, mosaic_time = timed(map_tiles_to_grid)(grid_labels, tiles, cell_size, color_centers, tile_colors)
    
    total_time = time.time() - start_total
    
    # Step 4: Performance Metrics
    start_metrics = time.time()
    mse_val = mse(resized_np_cropped, mosaic_img)
    ssim_val = ssim_metric(resized_np_cropped, mosaic_img)
    metrics_time = time.time() - start_metrics
    
    # Convert for display - all outputs as PIL Images
    orig_disp = Image.fromarray(cv2.cvtColor(resized_np_cropped, cv2.COLOR_BGR2RGB))
    mosaic_disp = Image.fromarray(cv2.cvtColor(mosaic_img, cv2.COLOR_BGR2RGB))
    
    # Segmented image: show each cell as its cluster color
    seg_img = np.zeros_like(resized_np_cropped)
    for i in range(actual_grid_size):
        for j in range(actual_grid_size):
            seg_img[i*cell_size:(i+1)*cell_size, j*cell_size:(j+1)*cell_size] = color_centers[grid_labels[i,j]]
    seg_disp = Image.fromarray(cv2.cvtColor(seg_img, cv2.COLOR_BGR2RGB))
    
    # Enhanced performance info using perfMetric module
    performance_info = create_performance_report(
        actual_grid_size, cell_size, n_colors, actual_image_size,
        preprocess_time, seg_time, tile_load_time, mosaic_time, metrics_time,
        total_time, mse_val, ssim_val
    )
    
    return orig_disp, seg_disp, mosaic_disp, performance_info

def performance_benchmark(image: Image.Image, n_colors: int = 16):
    """Run the automated benchmark on either the uploaded or default image.

    Args:
        image (PIL.Image.Image): Optional user image.
        n_colors (int, optional): Color palette size. Defaults to 16.

    Returns:
        tuple[str, str] | tuple[None, str]: Benchmark plots and textual summary.
    """
    preprocess_funcs = {
        'resize': resize_img,
        'quantize': color_quantize
    }

    if image is None:
        try:
            benchmark_image = Image.open(BENCHMARK_IMAGE_PATH).convert("RGB")
        except FileNotFoundError:
            return None, f"Benchmark image not found at {BENCHMARK_IMAGE_PATH}"
    else:
        try:
            benchmark_image = ensure_image(image)
        except ValidationError as exc:
            raise gr.Error(str(exc)) from exc

    ensure_positive_int(n_colors, "Color palette size", minimum=2, maximum=64)

    return run_performance_benchmark(
        preprocess_funcs,
        segment_image_grid,
        load_tile_images,
        map_tiles_to_grid,
        benchmark_image,
        n_colors,
        image_sizes=BENCHMARK_IMAGE_SIZES,
        grid_sizes=BENCHMARK_GRID_SIZES
    )

example_data = [
    ("data/test_images/bird_couple.avif", 16, 256),
    ("data/test_images/flamingo.avif", 32, 512),
    ("data/test_images/crow.jpg", 64, 1024),
    ("data/test_images/marten.jpg", 32, 512),
]

example_image_sizes = [item[2] for item in example_data]
examples = [[path, grid, 16, size] for path, grid, size in example_data]

example_urls = [path for path, _, _ in example_data]

benchmark_examples = [[BENCHMARK_IMAGE_PATH, 16]]

with gr.Blocks(title="Mosaic Generator - Performance Analysis") as demo:
    gr.Markdown("# 🌸 Mosaic Generator by Adrien Mery")
    gr.Markdown("Transform your images into beautiful mosaics using colorful photo tiles. Includes comprehensive performance analysis!")
    with gr.Accordion("Usage Tips", open=False):
        gr.Markdown(
            "- Choose matching image/grid sizes so cells remain square.\n"
            "- Upload high-resolution images for better mosaics.\n"
            "- Benchmark tab caches tiles automatically; rerun after the first pass for fastest results."
        )
    
    with gr.Tabs():
        with gr.Tab("🎨 Mosaic Generator"):
            with gr.Row():
                with gr.Column(scale=2):
                    image_input = gr.Image(type="pil", label="Upload Image or Select Example", height=300)
                with gr.Column(scale=1):
                    grid_slider = gr.Slider(8, 128, value=64, step=1, label="Mosaic Grid Resolution (NxN cells)")
                    color_slider = gr.Slider(2, 32, value=16, step=1, label="Color Palette Size (quantization)")
                    image_size_input = gr.Dropdown(
                        choices=example_image_sizes,
                        value=DEFAULT_IMAGE_SIZE,
                        label="Target Image Size (pixels)"
                    )
                    mosaic_btn = gr.Button("🎨 Generate Mosaic", variant="primary")
            
            with gr.Row():
                orig_out = gr.Image(label="Original (Resized & Quantized)")
                seg_out = gr.Image(label="Grid Segmentation Preview")
                mosaic_out = gr.Image(label="Final Mosaic Result")
            
            performance_out = gr.Textbox(label="📊 Detailed Performance Metrics", interactive=False, lines=15)
            
            gr.Examples(
                examples=examples,
                inputs=[image_input, grid_slider, color_slider, image_size_input],
                label="Example Images"
            )
        
        with gr.Tab("📈 Performance Benchmark"):
            gr.Markdown("""
            ### Performance Analysis Tool
            This tool benchmarks either your uploaded image or the default marten test image across multiple grid and image sizes to analyze:
            - **Processing time scaling** with grid complexity
            - **Quality metrics** (MSE & SSIM) across different resolutions  
            - **Performance bottlenecks** in the pipeline
            - **Optimal settings** for speed vs quality trade-offs
            """)
            
            with gr.Row():
                benchmark_input = gr.Image(type="pil", label="Upload Image for Benchmarking", height=300)
                with gr.Column(): 
                    benchmark_colors = gr.Slider(2, 32, value=16, step=1, label="Color Palette Size")
                    benchmark_btn = gr.Button("🔬 Run Performance Benchmark", variant="secondary")

            gr.Examples(examples=benchmark_examples, inputs=[benchmark_input, benchmark_colors], label="Example Images for Benchmarking")
            
            benchmark_plot = gr.Image(label="📊 Performance Analysis Charts")
            benchmark_results = gr.Textbox(label="📋 Detailed Benchmark Results", interactive=False, lines=20)
    
    mosaic_btn.click(
        preprocess_and_mosaic, 
        inputs=[image_input, grid_slider, color_slider, image_size_input], 
        outputs=[orig_out, seg_out, mosaic_out, performance_out]
    )
    
    benchmark_btn.click(
        performance_benchmark,
        inputs=[benchmark_input, benchmark_colors],
        outputs=[benchmark_plot, benchmark_results]
    )

demo.launch()