Add preprocess files

app.py

@@ -2,165 +2,127 @@ import gradio as gr
 import cv2
 import numpy as np
 import os
+import pickle
+import math
 from skimage.metrics import structural_similarity as ssim
 
-# Folder containing your dataset of tiles (small images)
-DATASET_FOLDER = "Dataset"
+# ----------------- Constants -----------------
+DATASET_FOLDER = "Dataset"  # (Not used directly now)
+KD_TREE_PATH = "kdtree_dataset.pkl"  # Path to the precomputed KDTree file
+KD_TILE_SIZE = (50, 50)  # Must match the tile size used when building the KDTree
 
+# ----------------- Feature Extraction Functions -----------------
 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)
+      - 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_edge = np.mean(edges) / 255.0
     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):
+def compute_weighted_features(image):
     """
-    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)
+    Compute the weighted feature vector for KDTree search.
+    The image should be resized to KD_TILE_SIZE before feature extraction.
+    Weights: for Lab channels use 1.0; for edge, texture, and gradient use 0.5
+    (implemented as multiplying by sqrt(0.5)).
     """
-    dataset = []
-    image_paths = [os.path.join(folder_path, img) for img in os.listdir(folder_path)
-                   if img.lower().endswith(('.png', '.jpg', '.jpeg'))]
+    scale = np.array([1.0, 1.0, 1.0, math.sqrt(0.5), math.sqrt(0.5), math.sqrt(0.5)])
+    avg_lab, avg_edge, avg_texture, avg_grad = compute_features(image)
+    raw_feature = np.concatenate([avg_lab, [avg_edge, avg_texture, avg_grad]])
+    weighted_feature = raw_feature * scale
+    return weighted_feature
 
-    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)):
+# ----------------- Utility Function -----------------
+def ensure_min_size(image, min_size=7):
     """
-    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.
+    Ensure that the image has at least a minimum size; if not, resize it.
     """
-    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
+    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
 
-def create_photo_mosaic(input_image, dataset_folder, num_tiles_y, progress=None):
+# ----------------- Mosaic Generation Functions -----------------
+def create_photo_mosaic(input_image, kdtree_path, 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.
+    Create an image mosaic using a precomputed KDTree.
+    For each mosaic tile in the input image, the tile is resized to KD_TILE_SIZE,
+    its weighted features are computed, and the KDTree is queried to find the best match.
+    The matched dataset image is then resized to the tile’s actual size before placement.
     """
-    # Assume the uploaded image from Gradio is in RGB format
+    # Load the precomputed KDTree and dataset images
+    with open(kdtree_path, "rb") as f:
+        tree_data = pickle.load(f)
+    tree = tree_data['tree']
+    dataset_images = tree_data['images']
+
     original_image = input_image.copy()
     height, width, _ = original_image.shape
 
-    # Compute tile height and determine tile width based on aspect ratio
+    # Determine mosaic grid dimensions
     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)
+            # Resize the tile to the KDTree tile size for feature extraction
+            tile_resized = cv2.resize(tile, KD_TILE_SIZE)
+            query_feature = compute_weighted_features(tile_resized)
 
-            # Find the best matching dataset image using the combined feature metric
-            best_match = find_best_match(tile_features, dataset)
+            # Query the KDTree for the best match (returns index)
+            dist, ind = tree.query([query_feature], k=1)
+            best_index = ind[0][0]
+            best_match = dataset_images[best_index]
 
-            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]
+            # Resize the best match image to the current tile size and place it into the mosaic
+            best_match_resized = cv2.resize(best_match, (x_end - x, y_end - y))
+            mosaic[y:y_end, x:x_end] = best_match_resized
 
             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.
+    # Save the final mosaic (convert from RGB to BGR for saving with cv2)
     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.
+    Create 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
@@ -191,8 +153,7 @@ def create_color_mosaic(input_image, num_tiles_y, progress=None):
     cv2.imwrite(output_path, cv2.cvtColor(mosaic, cv2.COLOR_RGB2BGR))
     return output_path
 
-# ----------------- Performance Metrics Functions -----------------
-
+# ----------------- Performance Metrics -----------------
 def compute_mse(original, mosaic):
     """
     Compute Mean Squared Error (MSE) between two images.
@@ -203,66 +164,42 @@ def compute_mse(original, mosaic):
     mse = err / float(original.shape[0] * original.shape[1] * original.shape[2])
     return mse
 
-def compute_ssim(original, mosaic):
+def compute_ssim_metric(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
+    win_size = 7 if min_dim >= 7 else (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 -----------------
-
+# ----------------- Gradio Interface -----------------
 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"
+    mosaic_type: "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_path = create_photo_mosaic(input_image, KD_TREE_PATH, num_tiles_y, progress)
+
     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)
-    
+    ssim_value = compute_ssim_metric(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"],

kdtree_dataset.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:82d2da18cad1927be1094563f598f33668595a1c62ca336fd86eadcfeda75c59
+size 44070974

preprocess.py

@@ -0,0 +1,104 @@
+# build_kdtree.py
+
+import os
+import cv2
+import numpy as np
+import pickle
+import math
+from sklearn.neighbors import KDTree
+
+# ----------------- Constants -----------------
+DATASET_FOLDER = "Dataset"     # Folder containing your dataset images
+KD_TILE_SIZE = (50, 50)        # Fixed size to which each dataset image will be resized
+KD_TREE_PATH = "kdtree_dataset.pkl"  # Output pickle file
+
+# ----------------- Feature Extraction -----------------
+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)
+    """
+    # 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
+    edges = cv2.Canny(gray, 100, 200)
+    avg_edge = np.mean(edges) / 255.0
+
+    # Texture: standard deviation (normalized)
+    avg_texture = np.std(gray) / 255.0
+
+    # Gradient magnitude using Sobel operators
+    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 build_kdtree():
+    """
+    Build a KDTree from dataset images. Each image is resized to KD_TILE_SIZE,
+    its features are computed and then weighted (using weights: 1.0 for Lab channels,
+    0.5 for edge, texture, and gradient differences).
+    The KDTree along with the list of dataset images is stored in a pickle file.
+    """
+    # Weights: for the Lab channels, weight = 1.0 (so sqrt(1.0)=1),
+    # for the other features, weight = 0.5 (so multiply by sqrt(0.5)).
+    scale = np.array([1.0, 1.0, 1.0, math.sqrt(0.5), math.sqrt(0.5), math.sqrt(0.5)])
+    
+    feature_list = []
+    images_list = []
+
+    # Get full paths for images in the dataset folder
+    image_paths = [os.path.join(DATASET_FOLDER, img) for img in os.listdir(DATASET_FOLDER)
+                   if img.lower().endswith(('.png', '.jpg', '.jpeg'))]
+
+    for img_path in image_paths:
+        img = cv2.imread(img_path)
+        if img is None:
+            continue
+        # Resize image to KD_TILE_SIZE and convert BGR -> RGB
+        img = cv2.resize(img, KD_TILE_SIZE)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+        # Compute features for the image
+        avg_lab, avg_edge, avg_texture, avg_grad = compute_features(img)
+        # Concatenate the features into a 6-dimensional vector:
+        raw_feature = np.concatenate([avg_lab, [avg_edge, avg_texture, avg_grad]])
+        # Apply weighting: multiply each element by the square-root of its weight
+        weighted_feature = raw_feature * scale
+        feature_list.append(weighted_feature)
+        images_list.append(img)
+
+    if not feature_list:
+        print("No images found in dataset folder!")
+        return
+
+    features = np.array(feature_list)
+    # Build the KDTree using the weighted features
+    tree = KDTree(features)
+
+    tree_data = {
+        'tree': tree,
+        'images': images_list,
+        'features': features  # optional: may be used for debugging
+    }
+
+    # Save the KDTree and dataset images to a pickle file
+    with open(KD_TREE_PATH, "wb") as f:
+        pickle.dump(tree_data, f)
+
+    print(f"KDTree built and saved to {KD_TREE_PATH}. Total images: {len(images_list)}")
+
+if __name__ == "__main__":
+    build_kdtree()