app.py

import gradio as gr
import cv2
import numpy as np
import os
from skimage.metrics import structural_similarity as ssim

# Folder containing your dataset of tiles (small images)
DATASET_FOLDER = "Dataset"

def compute_features(image):
    """
    Compute a set of features for an image:
    - Average Lab color (using a Gaussian-blurred version)
    - Edge density using Canny edge detection (normalized)
    - Texture measure using the standard deviation of the grayscale image (normalized)
    - Average gradient magnitude computed via Sobel operators (normalized)
    Returns: (avg_lab, avg_edge, avg_texture, avg_grad)
    """
    # Apply Gaussian blur to reduce noise before computing Lab color
    blurred = cv2.GaussianBlur(image, (5, 5), 0)
    img_lab = cv2.cvtColor(blurred, cv2.COLOR_RGB2LAB)
    avg_lab = np.mean(img_lab, axis=(0, 1))

    # Convert to grayscale for edge and texture computations
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)

    # Edge density: apply Canny and normalize the average edge intensity
    edges = cv2.Canny(gray, 100, 200)
    avg_edge = np.mean(edges) / 255.0  # Normalized edge density

    # Texture measure: standard deviation of grayscale values (normalized)
    avg_texture = np.std(gray) / 255.0

    # Gradient magnitude: using Sobel operators in x and y directions, then average
    grad_x = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
    grad_y = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
    grad_mag = np.sqrt(grad_x**2 + grad_y**2)
    avg_grad = np.mean(grad_mag) / 255.0

    return avg_lab, avg_edge, avg_texture, avg_grad

def load_dataset_images(folder_path, tile_size):
    """
    Loads images from a folder, resizes them to tile_size, and computes a set of features:
    (RGB image, average Lab color, edge density, texture measure, gradient magnitude, image path)
    """
    dataset = []
    image_paths = [os.path.join(folder_path, img) for img in os.listdir(folder_path)
                   if img.lower().endswith(('.png', '.jpg', '.jpeg'))]

    for img_path in image_paths:
        img = cv2.imread(img_path)
        if img is None:
            continue  # Skip unreadable images
        # Resize the image to the given tile size
        img = cv2.resize(img, tile_size)
        # Convert from BGR to RGB
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        # Compute the feature vector for this dataset image
        avg_lab, avg_edge, avg_texture, avg_grad = compute_features(img)
        dataset.append((img, avg_lab, avg_edge, avg_texture, avg_grad, img_path))

    return dataset

def find_best_match(tile_features, dataset, weights=(1.0, 0.5, 0.5, 0.5)):
    """
    Finds the best matching dataset image based on a weighted combination of:
    - Color difference (in Lab space)
    - Edge density difference
    - Texture difference
    - Gradient magnitude difference

    The weights parameter is a tuple with weights for each feature in the same order.
    """
    tile_lab, tile_edge, tile_texture, tile_grad = tile_features
    min_dist = float('inf')
    best_match = None

    for data in dataset:
        ds_img, ds_lab, ds_edge, ds_texture, ds_grad, ds_path = data

        # Compute the difference for each feature:
        color_diff = np.linalg.norm(tile_lab - ds_lab)
        edge_diff = abs(tile_edge - ds_edge)
        texture_diff = abs(tile_texture - ds_texture)
        grad_diff = abs(tile_grad - ds_grad)

        # Compute a weighted distance (using a Euclidean combination)
        dist = np.sqrt(weights[0] * (color_diff ** 2) +
                       weights[1] * (edge_diff ** 2) +
                       weights[2] * (texture_diff ** 2) +
                       weights[3] * (grad_diff ** 2))

        if dist < min_dist:
            min_dist = dist
            best_match = ds_img

    return best_match

def create_photo_mosaic(input_image, dataset_folder, num_tiles_y, progress=None):
    """
    Creates an image mosaic using dataset images. For each tile of the input image,
    it computes a feature vector (color, edge, texture, gradient) and finds the best
    matching dataset image based on these features.
    """
    # Assume the uploaded image from Gradio is in RGB format
    original_image = input_image.copy()
    height, width, _ = original_image.shape

    # Compute tile height and determine tile width based on aspect ratio
    tile_height = height // num_tiles_y
    aspect_ratio = width / height
    num_tiles_x = int(num_tiles_y * aspect_ratio)
    tile_width = width // num_tiles_x
    print(f"Adjusted number of tiles: {num_tiles_x} (width) x {num_tiles_y} (height)")

    # Load the dataset images with the new feature set
    dataset = load_dataset_images(dataset_folder, (tile_width, tile_height))
    if not dataset:
        print("No images found in dataset folder!")
        return None

    # Create an empty mosaic image in RGB
    mosaic = np.zeros_like(original_image)

    # Calculate the grid ranges and total tile count for progress tracking
    rows = list(range(0, height, tile_height))
    cols = list(range(0, width, tile_width))
    total_tiles = len(rows) * len(cols)
    tile_count = 0

    # Process each tile of the input image
    for y in rows:
        for x in cols:
            y_end = min(y + tile_height, height)
            x_end = min(x + tile_width, width)
            tile = original_image[y:y_end, x:x_end]

            # Compute feature vector for the tile
            tile_features = compute_features(tile)

            # Find the best matching dataset image using the combined feature metric
            best_match = find_best_match(tile_features, dataset)

            if best_match is not None:
                # Crop the dataset image if necessary to match the tile size
                mosaic[y:y_end, x:x_end] = best_match[:y_end - y, :x_end - x]

            tile_count += 1
            if progress is not None:
                progress(tile_count / total_tiles)

    # Save the final mosaic. Since mosaic is in RGB, convert to BGR for cv2.imwrite.
    output_path = "mosaic_output.jpg"
    cv2.imwrite(output_path, cv2.cvtColor(mosaic, cv2.COLOR_RGB2BGR))

    return output_path

def create_color_mosaic(input_image, num_tiles_y, progress=None):
    """
    Creates a simple color mosaic by dividing the image into grid cells and
    filling each cell with its average RGB color.
    """
    original_image = input_image.copy()
    height, width, _ = original_image.shape
    tile_height = height // num_tiles_y
    aspect_ratio = width / height
    num_tiles_x = int(num_tiles_y * aspect_ratio)
    tile_width = width // num_tiles_x
    print(f"Adjusted number of tiles: {num_tiles_x} (width) x {num_tiles_y} (height)")

    mosaic = np.zeros_like(original_image)
    rows = list(range(0, height, tile_height))
    cols = list(range(0, width, tile_width))
    total_tiles = len(rows) * len(cols)
    tile_count = 0

    for y in rows:
        for x in cols:
            y_end = min(y + tile_height, height)
            x_end = min(x + tile_width, width)
            tile = original_image[y:y_end, x:x_end]
            avg_color = np.mean(tile, axis=(0, 1)).astype(np.uint8)
            mosaic[y:y_end, x:x_end] = avg_color
            tile_count += 1
            if progress is not None:
                progress(tile_count / total_tiles)

    output_path = "color_mosaic_output.jpg"
    cv2.imwrite(output_path, cv2.cvtColor(mosaic, cv2.COLOR_RGB2BGR))
    return output_path

# ----------------- Performance Metrics Functions -----------------

def compute_mse(original, mosaic):
    """
    Compute Mean Squared Error (MSE) between two images.
    """
    original = original.astype("float")
    mosaic = mosaic.astype("float")
    err = np.sum((original - mosaic) ** 2)
    mse = err / float(original.shape[0] * original.shape[1] * original.shape[2])
    return mse

def compute_ssim(original, mosaic):
    """
    Compute Structural Similarity Index (SSIM) between two images.
    """
    min_dim = min(original.shape[0], original.shape[1])
    if min_dim >= 7:
        win_size = 7
    else:
        # Ensure the window size is odd.
        win_size = min_dim if min_dim % 2 == 1 else min_dim - 1
    ssim_value, _ = ssim(original, mosaic, win_size=win_size, channel_axis=-1, full=True)
    return ssim_value

def ensure_min_size(image, min_size=7):
    """
    Ensure that the image has a minimum size; if not, resize it.
    """
    h, w = image.shape[:2]
    if h < min_size or w < min_size:
        new_w = max(min_size, w)
        new_h = max(min_size, h)
        image = cv2.resize(image, (new_w, new_h))
    return image

# ----------------- Gradio Interface Function -----------------

def mosaic_gradio(input_image, num_tiles_y, mosaic_type, progress=gr.Progress()):
    """
    Gradio interface function to generate and return the mosaic image along with performance metrics.
    mosaic_type: Either "Color Mosaic" or "Image Mosaic"
    Returns: (mosaic_image_file, performance_metrics_string)
    """
    # Generate mosaic based on selected type
    if mosaic_type == "Color Mosaic":
        mosaic_path = create_color_mosaic(input_image, num_tiles_y, progress)
    else:
        mosaic_path = create_photo_mosaic(input_image, DATASET_FOLDER, num_tiles_y, progress)
    
    # Load the mosaic image from file (convert from BGR to RGB)
    mosaic_image = cv2.imread(mosaic_path)
    if mosaic_image is None:
        return None, "Error: Mosaic image could not be loaded."
    mosaic_image = cv2.cvtColor(mosaic_image, cv2.COLOR_BGR2RGB)
    
    # Ensure both images meet minimum size requirements for metric calculations
    input_for_metrics = ensure_min_size(input_image.copy())
    mosaic_for_metrics = ensure_min_size(mosaic_image.copy())
    
    # Compute performance metrics
    mse_value = compute_mse(input_for_metrics, mosaic_for_metrics)
    ssim_value = compute_ssim(input_for_metrics, mosaic_for_metrics)
    
    metrics_text = f"MSE: {mse_value:.2f}\nSSIM: {ssim_value:.4f}"
    
    return mosaic_path, metrics_text

# ----------------- Gradio App Setup -----------------

# Adding examples so that test images appear as clickable examples.
# Adjust the paths as needed.
examples = [
    ["input_images/1.jpg", 90, "Image Mosaic"],
    ["input_images/2.jpg", 90, "Image Mosaic"],
    ["input_images/3.jpg", 90, "Image Mosaic"],
    ["input_images/6.jpg", 90, "Image Mosaic"],
    ["input_images/7.jpg", 90, "Image Mosaic"],
    ["input_images/8.jpg", 90, "Image Mosaic"],
    ["input_images/9.jpg", 90, "Image Mosaic"],
    ["input_images/10.jpg", 90, "Image Mosaic"]
]

iface = gr.Interface(
    fn=mosaic_gradio,
    inputs=[
        gr.Image(type="numpy", label="Upload Image"),
        gr.Slider(10, 200, value=90, step=5, label="Number of Tiles (Height)"),
        gr.Radio(choices=["Color Mosaic", "Image Mosaic"], label="Mosaic Type", value="Image Mosaic")
    ],
    outputs=[
        gr.Image(type="filepath", label="Generated Mosaic"),
        gr.Textbox(label="Performance Metrics")
    ],
    title="Photo Mosaic Generator",
    description=("Upload an image, choose the number of tiles (height) and mosaic type. "
                 "Select 'Color Mosaic' for a mosaic using average colors, or 'Image Mosaic' to use dataset images "
                 "matched by color, edge density, texture, and gradient features. "
                 "After mosaic generation, performance metrics (MSE and SSIM) will be displayed."),
    examples=examples
)

iface.launch()