Updated Lab5 modular code

app.py

@@ -1,165 +1,259 @@
 #!/usr/bin/env python3
+"""
+app.py
+
+Gradio interface for the Optimised Mosaic Generator (Lab 5).
+Loads CIFAR tiles once, then performs fast LAB-based matching using FAISS.
+
+This file connects the UI to:
+- crop_to_multiple()
+- compute_cell_means_lab()
+- TileManager
+- MosaicBuilder
+- MSE / SSIM metrics
+"""
 
 import gradio as gr
 import numpy as np
 import time
+import os
 from PIL import Image, ImageDraw
 
-from mosaic_generator.image_processor import (
-    crop_to_multiple,
-    compute_cell_means_lab
-)
+from mosaic_generator.image_processor import crop_to_multiple, compute_cell_means_lab
 from mosaic_generator.tile_manager import TileManager
 from mosaic_generator.mosaic_builder import MosaicBuilder
 from mosaic_generator.metrics import mse, ssim_rgb
 
 
 # -------------------------------------------------------------------
-# GLOBAL TILE MANAGER (load once)
+# GLOBAL TILE MANAGER → loaded ONCE for the entire Space
 # -------------------------------------------------------------------
 TM = TileManager()
-TM.load(sample_size=20000)      # Same as Lab 1
+TM.load(sample_size=20000)
 
 
 # -------------------------------------------------------------------
 # MAIN PIPELINE
 # -------------------------------------------------------------------
 def run_pipeline(
-    img,
-    grid_size,
-    tile_px,
-    tile_sample,
-    quantize_on,
-    quantize_colors,
-    show_grid
+    img, grid_size, tile_px, tile_sample,
+    quantize_on, quantize_colors, show_grid
 ):
+    """
+    Full end-to-end mosaic pipeline executed when user clicks GENERATE.
+
+    Parameters
+    ----------
+    img : PIL.Image
+    grid_size : str
+    tile_px : str
+    tile_sample : int
+    quantize_on : bool
+    quantize_colors : int
+    show_grid : bool
+
+    Returns
+    -------
+    original_img : PIL.Image
+    segmented_img : PIL.Image
+    mosaic_img : PIL.Image
+    report_str : str
+    """
+
+    # No image provided
     if img is None:
         return None, None, None, "Upload an image first."
 
     img_np = np.array(img.convert("RGB"))
-
     grid_n = int(grid_size)
-    tile_px = int(tile_px)
-    tile_sample = int(tile_sample)
 
-    # Crop image
+    # ------------------------------------------
+    # Crop image to ensure perfect cell division
+    # ------------------------------------------
     base = crop_to_multiple(img_np, grid_n)
 
+    # ------------------------------------------
     # Optional quantization
+    # ------------------------------------------
     if quantize_on:
-        pi = Image.fromarray(base).quantize(
-            colors=int(quantize_colors),
-            method=Image.MEDIANCUT,
-            dither=Image.Dither.NONE
-        ).convert("RGB")
-        base = np.array(pi)
-
-    # Compute LAB means
-    t0 = time.perf_counter()
-    cell_means, dims = compute_cell_means_lab(base, grid_n)
-    t1 = time.perf_counter()
-
-    # Prepare tiles
-    TM.prepare_scaled_tiles(dims[2], dims[3])
-
-    # Nearest tiles via FAISS
-    idxs = TM.lookup_tiles(cell_means)
-
+        try:
+            q = Image.fromarray(base).quantize(
+                colors=int(quantize_colors),
+                method=Image.MEDIANCUT,
+                dither=Image.Dither.NONE
+            ).convert("RGB")
+            base = np.array(q)
+        except Exception as e:
+            return None, None, None, f"Quantization failed: {e}"
+
+    # ------------------------------------------
+    # Compute LAB means for all grid cells
+    # ------------------------------------------
+    try:
+        t0 = time.perf_counter()
+        cell_means, dims = compute_cell_means_lab(base, grid_n)
+        t1 = time.perf_counter()
+    except Exception as e:
+        return None, None, None, f"LAB computation failed: {e}"
+
+    w, h, cell_w, cell_h = dims
+
+    # ------------------------------------------
+    # Prepare tiles (resize once per cell size)
+    # ------------------------------------------
+    TM.prepare_scaled_tiles(cell_w, cell_h)
+
+    # ------------------------------------------
+    # Find nearest tile via FAISS
+    # ------------------------------------------
+    try:
+        idxs = TM.lookup_tiles(cell_means)
+    except Exception as e:
+        return None, None, None, f"Tile lookup failed: {e}"
+
+    # ------------------------------------------
     # Build mosaic
+    # ------------------------------------------
     builder = MosaicBuilder(TM)
-    mosaic_np = builder.build(idxs, dims, grid_n)
-    t2 = time.perf_counter()
+    try:
+        mosaic_np = builder.build(idxs, dims, grid_n)
+    except Exception as e:
+        return None, None, None, f"Mosaic build failed: {e}"
 
-    # Metrics
-    mse_val = mse(base, mosaic_np)
-    ssim_val = ssim_rgb(base, mosaic_np)
+    t2 = time.perf_counter()
 
-    # Segmented grid overlay
+    # ------------------------------------------
+    # Compute metrics
+    # ------------------------------------------
+    try:
+        mse_val = mse(base, mosaic_np)
+        ssim_val = ssim_rgb(base, mosaic_np)
+    except Exception as e:
+        mse_val, ssim_val = -1, -1
+
+    # ------------------------------------------
+    # Grid overlay (optional)
+    # ------------------------------------------
     segmented = Image.fromarray(base)
     if show_grid:
         seg = segmented.copy()
         draw = ImageDraw.Draw(seg)
-        w, h, cw, ch = dims
-        for x in range(0, w, cw):
-            draw.line([(x, 0), (x, h)], fill=(255, 0, 0), width=1)
-        for y in range(0, h, ch):
-            draw.line([(0, y), (w, y)], fill=(255, 0, 0), width=1)
+        for x in range(0, w, cell_w):
+            draw.line([(x, 0), (x, h)], fill="red", width=1)
+        for y in range(0, h, cell_h):
+            draw.line([(0, y), (w, y)], fill="red", width=1)
         segmented = seg
 
-    # Build report
+    # ------------------------------------------
+    # Text report
+    # ------------------------------------------
     report = (
         f"MSE: {mse_val:.2f}\n"
         f"SSIM: {ssim_val:.4f}\n\n"
         f"Preprocessing Time: {t1 - t0:.3f}s\n"
-        f"Mosaic Build Time: {t2 - t1:.3f}s\n"
-        f"Total Time: {t2 - t0:.3f}s\n"
+        f"Mosaic Build Time:  {t2 - t1:.3f}s\n"
+        f"Total Time:          {t2 - t0:.3f}s\n"
     )
 
     return (
-        Image.fromarray(base),         # Original cropped
-        segmented,                     # Grid segmented
-        Image.fromarray(mosaic_np),    # Mosaic output
-        report                         # Timing & metrics
+        Image.fromarray(base),
+        segmented,
+        Image.fromarray(mosaic_np),
+        report
     )
 
 
 # -------------------------------------------------------------------
-# GRADIO UI
+# BUILD GRADIO UI
 # -------------------------------------------------------------------
 def build_demo():
     with gr.Blocks(title="High-Performance Mosaic Generator") as demo:
+
         gr.Markdown("# ⚡ High-Performance Mosaic Generator (Lab 5)")
-        gr.Markdown("20×–100× faster than Lab 1 using FAISS, Numba-LAB, and OpenCV.")
+        gr.Markdown("Ultra-fast FAISS-powered image mosaic generator.\n")
 
         with gr.Row():
+
+            # ----------------------------------------------------
+            # LEFT COLUMN — INPUTS
+            # ----------------------------------------------------
             with gr.Column(scale=1):
+
                 img_in = gr.Image(type="pil", label="Upload Image")
 
                 grid_size = gr.Radio(
-                    ["16", "32", "64", "128"], value="32",
-                    label="Grid Size (cells per side)"
+                    ["16", "32", "64", "128"],
+                    value="32",
+                    label="Grid Size"
                 )
                 tile_px = gr.Radio(
-                    ["8", "16", "24", "32"], value="16",
+                    ["8", "16", "24", "32"],
+                    value="16",
                     label="Tile Resolution (px)"
                 )
 
                 tile_sample = gr.Slider(
                     512, 20000, step=256, value=2048,
-                    label="Number of CIFAR-100 Tiles to Sample"
+                    label="Tile Sample Size"
                 )
 
-                quantize_on = gr.Checkbox(False, label="Enable color quantization")
-                quantize_colors = gr.Slider(8, 128, value=32, step=8,
-                                            label="Quantization Palette Size")
+                quantize_on = gr.Checkbox(True, label="Enable Color Quantization")
+                quantize_colors = gr.Slider(
+                    8, 128, value=32, step=8,
+                    label="Quantization Palette Size"
+                )
 
-                show_grid = gr.Checkbox(True, label="Show Grid Overlay")
+                show_grid = gr.Checkbox(True, label="Show Grid")
 
                 run_btn = gr.Button("Generate Mosaic", variant="primary")
 
-            # --- Outputs ---
+                # ------------------------------------------------
+                # EXAMPLE IMAGES
+                # ------------------------------------------------
+                gr.Markdown("### Example Images")
+                TEST_DIR = "test"
+
+                example_files = [
+                    os.path.join(TEST_DIR, f) for f in os.listdir(TEST_DIR)
+                    if f.lower().endswith((".png", ".jpg", ".jpeg"))
+                ]
+
+                gr.Examples(
+                    examples=[[f] for f in example_files],
+                    inputs=[img_in],
+                    label="",
+                    cache_examples=False
+                )
+
+            # ----------------------------------------------------
+            # RIGHT COLUMN — OUTPUTS
+            # ----------------------------------------------------
             with gr.Column(scale=2):
-                with gr.Tab("Original (Cropped)"):
+
+                with gr.Tab("Original"):
                     img_orig = gr.Image()
 
-                with gr.Tab("Segmented (Grid)"):
+                with gr.Tab("Grid View"):
                     img_seg = gr.Image()
 
-                with gr.Tab("Mosaic Output"):
+                with gr.Tab("Mosaic"):
                     img_mosaic = gr.Image()
 
-                report = gr.Textbox(label="Metrics & Timing", lines=10)
+                report = gr.Textbox(label="Timing & Metrics", lines=10)
 
-        # FIXED — No None in outputs
         run_btn.click(
             fn=run_pipeline,
-            inputs=[img_in, grid_size, tile_px, tile_sample, quantize_on, quantize_colors, show_grid],
+            inputs=[img_in, grid_size, tile_px, tile_sample,
+                    quantize_on, quantize_colors, show_grid],
             outputs=[img_orig, img_seg, img_mosaic, report]
         )
 
     return demo
 
 
+# -------------------------------------------------------------------
+# LAUNCH
+# -------------------------------------------------------------------
 if __name__ == "__main__":
     demo = build_demo()
     demo.launch()

mosaic_generator/config.py

@@ -1,5 +1,93 @@
-# config.py
+"""
+config.py
 
+Central configuration module for the Optimised Mosaic Generator (Lab 5).
+
+Defines default parameters for:
+- tile sampling
+- grid size
+- tile resolution
+
+This helps maintain consistency across the project and allows the UI
+and benchmark scripts to share the same defaults.
+"""
+
+# -------------------------------------------------------------------
+# DEFAULT PARAMETERS
+# -------------------------------------------------------------------
+
+# Number of CIFAR-100 tiles to sample by default
 DEFAULT_TILE_COUNT = 2048
-DEFAULT_TILE_SIZE = 32
+
+# Pixel resolution of each tile before scaling (e.g., 8, 16, 24, 32)
+DEFAULT_TILE_SIZE = 16
+
+# Mosaic grid dimension (32x32 → 1024 total cells)
 DEFAULT_GRID = 32
+
+
+# -------------------------------------------------------------------
+# OPTIONAL VALIDATION HELPERS
+# These are simple checks that can be used by app.py or benchmarks.
+# -------------------------------------------------------------------
+
+def validate_grid_size(n):
+    """
+    Validate grid size (must be divisible into the image cleanly).
+
+    Parameters
+    ----------
+    n : int
+        Desired grid dimension per side.
+
+    Returns
+    -------
+    int
+        Validated grid size.
+
+    Raises
+    ------
+    ValueError
+        If the grid size is invalid.
+    """
+    if not isinstance(n, int) or n <= 0:
+        raise ValueError(f"Grid size must be a positive integer. Got: {n}")
+
+    if n not in [8, 16, 32, 64, 128]:
+        raise ValueError(
+            f"Unsupported grid size {n}. Choose from [8, 16, 32, 64, 128]."
+        )
+
+    return n
+
+
+def validate_tile_sample(k):
+    """
+    Validate number of sampled CIFAR tiles.
+
+    Ensures the number is within a reasonable bound for performance.
+
+    Parameters
+    ----------
+    k : int
+        Requested tile sample size.
+
+    Returns
+    -------
+    int
+        Validated tile count.
+
+    Raises
+    ------
+    ValueError
+        If the tile count is invalid.
+    """
+    if not isinstance(k, int) or k <= 0:
+        raise ValueError(f"Tile sample must be a positive integer. Got: {k}")
+
+    if k > 20000:
+        raise ValueError(
+            f"Tile sample {k} is too large. Max allowed: 20000."
+        )
+
+    return k

mosaic_generator/image_processor.py

@@ -1,28 +1,128 @@
-# image_processor.py
+"""
+image_processor.py
+
+Utility functions for image preprocessing used in the mosaic generator:
+- Cropping an image so it's divisible by the grid
+- Computing LAB cell means for FAISS-based tile matching
+"""
 
 import numpy as np
 import cv2
 
 from .utils import fast_rgb2lab
 
+
 def crop_to_multiple(img, grid_n):
+    """
+    Crop an RGB image so that its width and height are perfectly divisible
+    by the chosen grid size.
+
+    Parameters
+    ----------
+    img : np.ndarray
+        RGB image array of shape (H, W, 3).
+    grid_n : int
+        Number of cells per side in the mosaic grid.
+
+    Returns
+    -------
+    np.ndarray
+        Cropped RGB image whose dimensions are multiples of `grid_n`.
+
+    Raises
+    ------
+    ValueError
+        If `img` is not a valid image array or grid size is invalid.
+
+    Notes
+    -----
+    This does NOT resize the image — it simply trims extra pixels so that
+    (H % grid_n == 0) and (W % grid_n == 0).
+    """
+    if img is None or not isinstance(img, np.ndarray):
+        raise ValueError("Input image must be a valid NumPy RGB array.")
+
+    if img.ndim != 3 or img.shape[2] != 3:
+        raise ValueError(f"Expected image shape (H, W, 3), got {img.shape}.")
+
+    if not isinstance(grid_n, int) or grid_n <= 0:
+        raise ValueError("grid_n must be a positive integer.")
+
     h, w = img.shape[:2]
+
+    if h < grid_n or w < grid_n:
+        raise ValueError(
+            f"Image too small for grid size {grid_n}. "
+            f"Received image of size {w}x{h}."
+        )
+
     new_w = (w // grid_n) * grid_n
     new_h = (h // grid_n) * grid_n
+
     return img[:new_h, :new_w]
 
+
 def compute_cell_means_lab(img, grid_n):
-    """Convert FULL image to LAB once, then extract grid cell means."""
+    """
+    Compute LAB mean color for each grid cell in the image.
+
+    Parameters
+    ----------
+    img : np.ndarray
+        Cropped RGB image array (H, W, 3).
+    grid_n : int
+        Grid size — number of cells per side.
+
+    Returns
+    -------
+    means : np.ndarray
+        Array of shape (grid_n * grid_n, 3). LAB mean per grid cell.
+    dims : tuple
+        (W, H, cell_w, cell_h)
+
+        - W, H  : final image dimensions
+        - cell_w, cell_h : size of each grid cell in pixels
+
+    Raises
+    ------
+    ValueError
+        If the image is not divisible by grid_n, or has unexpected shape.
+
+    Notes
+    -----
+    The function converts the full image to LAB **once**, then extracts
+    block means efficiently without redundant conversions.
+    """
+    if img is None or not isinstance(img, np.ndarray):
+        raise ValueError("Input image must be a valid NumPy RGB array.")
+
+    if img.ndim != 3 or img.shape[2] != 3:
+        raise ValueError(f"Expected RGB image with 3 channels, got {img.shape}.")
+
+    if not isinstance(grid_n, int) or grid_n <= 0:
+        raise ValueError("grid_n must be a positive integer.")
+
     h, w = img.shape[:2]
+
+    if h % grid_n != 0 or w % grid_n != 0:
+        raise ValueError(
+            f"Image size ({w}x{h}) is not divisible by grid size {grid_n}. "
+            "Call crop_to_multiple() first."
+        )
+
     cell_h, cell_w = h // grid_n, w // grid_n
 
+    # Single conversion for full image
     lab = fast_rgb2lab(img)
 
+    # Output: N cells × 3 channels
     means = np.zeros((grid_n * grid_n, 3), dtype=np.float32)
     k = 0
+
     for gy in range(grid_n):
         for gx in range(grid_n):
             block = lab[gy*cell_h:(gy+1)*cell_h, gx*cell_w:(gx+1)*cell_w]
+            # Safe flatten + mean
             means[k] = block.reshape(-1, 3).mean(axis=0)
             k += 1
 

mosaic_generator/metrics.py

@@ -1,11 +1,101 @@
-# metrics.py
+"""
+metrics.py
+
+Provides image similarity/quality metrics for the mosaic generator:
+- Mean Squared Error (MSE)
+- Structural Similarity Index (SSIM) averaged over RGB channels
+"""
 
 import numpy as np
 from skimage.metrics import structural_similarity as ssim
 
+
 def mse(a, b):
-    return float(np.mean((a.astype(np.float32) - b.astype(np.float32)) ** 2))
+    """
+    Compute Mean Squared Error between two RGB images.
+
+    Parameters
+    ----------
+    a : np.ndarray
+        First RGB image array (H, W, 3).
+    b : np.ndarray
+        Second RGB image array (H, W, 3).
+
+    Returns
+    -------
+    float
+        Scalar MSE value.
+
+    Raises
+    ------
+    ValueError
+        If the input images are not the same shape or not valid RGB arrays.
+    """
+    if a is None or b is None:
+        raise ValueError("mse(): both input images must be provided.")
+
+    if not isinstance(a, np.ndarray) or not isinstance(b, np.ndarray):
+        raise ValueError("mse(): inputs must be NumPy arrays.")
+
+    if a.shape != b.shape:
+        raise ValueError(
+            f"mse(): image size mismatch. Got {a.shape} vs {b.shape}."
+        )
+
+    if a.ndim != 3 or a.shape[2] != 3:
+        raise ValueError(f"mse(): expected RGB images, got shape {a.shape}.")
+
+    diff = a.astype(np.float32) - b.astype(np.float32)
+    return float(np.mean(diff ** 2))
+
 
 def ssim_rgb(a, b):
-    vals = [ssim(a[...,c], b[...,c], data_range=255) for c in range(3)]
-    return float(sum(vals)/3)
+    """
+    Compute SSIM (Structural Similarity Index) for RGB images.
+
+    SSIM is computed per-channel and then averaged to produce a single score.
+
+    Parameters
+    ----------
+    a : np.ndarray
+        First RGB image array (H, W, 3).
+    b : np.ndarray
+        Second RGB image array (H, W, 3).
+
+    Returns
+    -------
+    float
+        Mean SSIM across the 3 RGB channels.
+
+    Raises
+    ------
+    ValueError
+        If input images are mismatched or invalid.
+    """
+    if a is None or b is None:
+        raise ValueError("ssim_rgb(): both input images must be provided.")
+
+    if not isinstance(a, np.ndarray) or not isinstance(b, np.ndarray):
+        raise ValueError("ssim_rgb(): inputs must be NumPy arrays.")
+
+    if a.shape != b.shape:
+        raise ValueError(
+            f"ssim_rgb(): image size mismatch. Got {a.shape} vs {b.shape}."
+        )
+
+    if a.ndim != 3 or a.shape[2] != 3:
+        raise ValueError(f"ssim_rgb(): expected RGB images, got shape {a.shape}.")
+
+    # Compute SSIM per channel
+    vals = [
+        ssim(
+            a[..., c],
+            b[..., c],
+            data_range=255,
+            win_size=7,      # helps stability for small tiles
+            gaussian_weights=True
+        )
+        for c in range(3)
+    ]
+
+    return float(sum(vals) / 3)

mosaic_generator/mosaic_builder.py

@@ -1,19 +1,103 @@
+"""
+mosaic_builder.py
+
+Reconstructs the final mosaic image by placing pre-scaled tiles into their
+corresponding grid-cell positions.
+"""
+
 import numpy as np
 
+
 class MosaicBuilder:
+    """
+    MosaicBuilder assembles the output mosaic using:
+    - FAISS-selected tile indices
+    - Pre-resized tiles (from TileManager)
+    - Grid/cell dimensions
+    """
+
     def __init__(self, tm):
+        """
+        Parameters
+        ----------
+        tm : TileManager
+            A TileManager instance containing pre-scaled tiles and FAISS index.
+        """
         self.tm = tm
 
+    # -------------------------------------------------------------
+    # MAIN MOSAIC RECONSTRUCTION
+    # -------------------------------------------------------------
     def build(self, tile_indices, dims, grid_n):
+        """
+        Construct final mosaic image using selected tile indices.
+
+        Parameters
+        ----------
+        tile_indices : np.ndarray
+            Flattened array of selected tile indices (length = grid_n * grid_n).
+        dims : tuple
+            (W, H, cell_w, cell_h):
+                W, H     → final image width & height
+                cell_w   → width of each grid cell
+                cell_h   → height of each grid cell
+        grid_n : int
+            Number of cells per side in the mosaic.
+
+        Returns
+        -------
+        np.ndarray
+            Final mosaic as an RGB array of shape (H, W, 3).
+
+        Raises
+        ------
+        ValueError
+            If tile indices, dims, or pre-scaled tiles are invalid.
+        RuntimeError
+            If tiles have not been pre-resized by TileManager.
+        """
+
+        # ------------------ VALIDATION ------------------
+        if tile_indices is None or not isinstance(tile_indices, np.ndarray):
+            raise ValueError("tile_indices must be a NumPy array.")
+
+        expected_len = grid_n * grid_n
+        if tile_indices.size != expected_len:
+            raise ValueError(
+                f"Expected {expected_len} tile indices, got {tile_indices.size}."
+            )
+
+        if self.tm.pre_scaled_tiles is None:
+            raise RuntimeError(
+                "Tiles have not been resized. Call TileManager.prepare_scaled_tiles() first."
+            )
+
+        if not isinstance(dims, tuple) or len(dims) != 4:
+            raise ValueError("dims must be a tuple of (W, H, cell_w, cell_h).")
+
         w, h, cell_w, cell_h = dims
+        if any(x <= 0 for x in [w, h, cell_w, cell_h]):
+            raise ValueError(f"Invalid dims values: {dims}")
+
+        # ------------------ OUTPUT CANVAS ------------------
         out = np.zeros((h, w, 3), dtype=np.uint8)
 
+        # ------------------ PLACE TILES ------------------
         k = 0
         for gy in range(grid_n):
             for gx in range(grid_n):
-                tile = self.tm.pre_scaled_tiles[tile_indices[k]]
-                out[gy * cell_h:(gy + 1) * cell_h,
-                    gx * cell_w:(gx + 1) * cell_w] = tile
+                idx = tile_indices[k]
+
+                if idx < 0 or idx >= len(self.tm.pre_scaled_tiles):
+                    raise ValueError(f"Tile index {idx} out of range.")
+
+                tile = self.tm.pre_scaled_tiles[idx]
+
+                out[
+                    gy * cell_h:(gy + 1) * cell_h,
+                    gx * cell_w:(gx + 1) * cell_w
+                ] = tile
+
                 k += 1
 
         return out

mosaic_generator/tile_manager.py

@@ -1,15 +1,34 @@
+"""
+tile_manager.py
+
+Manages loading, caching, and preprocessing of CIFAR-100 tiles.
+Handles:
+- Extracting RGB tiles
+- Computing LAB means for FAISS
+- Building FAISS index for fast NN search
+- Efficient tile resizing (cached)
+"""
+
 import os
 import pickle
 import numpy as np
 import cv2
+import faiss
 from datasets import load_dataset
+
 from .utils import fast_rgb2lab
-import faiss
+
 
 CACHE_DIR = "tile_cache"
 os.makedirs(CACHE_DIR, exist_ok=True)
 
+
 class TileManager:
+    """
+    TileManager handles loading CIFAR tiles, computing LAB mean features,
+    building a FAISS index, and caching pre-resized tiles for mosaic creation.
+    """
+
     def __init__(self):
         self.tiles_rgb = None
         self.tiles_lab_mean = None
@@ -18,51 +37,78 @@ class TileManager:
         self.loaded_sample_size = None
 
     # -------------------------------------------------------
-    # LOAD WITH DISK CACHING FOR 10K–20K TILES
+    # LOAD TILES (WITH CACHING)
     # -------------------------------------------------------
     def load(self, sample_size=2048):
         """
-        Loads CIFAR tiles.
-        Uses disk cache to avoid recomputing every launch.
+        Load CIFAR-100 tiles (RGB + LAB means).
+
+        Parameters
+        ----------
+        sample_size : int
+            Number of tiles to load (recommended: 2k–20k).
+
+        Notes
+        -----
+        - The first load may take ~20–50 seconds depending on size.
+        - Subsequent runs load instantly from disk cache.
         """
+        if not isinstance(sample_size, int) or sample_size <= 0:
+            raise ValueError("sample_size must be a positive integer.")
+
         self.loaded_sample_size = sample_size
         cache_file = f"{CACHE_DIR}/tiles_{sample_size}.pkl"
 
         # ------------------------------
-        # 1. LOAD FROM CACHE IF EXISTS
+        # 1. LOAD FROM CACHE
         # ------------------------------
         if os.path.exists(cache_file):
             print(f"✓ Loading cached tiles: {cache_file}")
-            with open(cache_file, "rb") as f:
-                data = pickle.load(f)
+            try:
+                with open(cache_file, "rb") as f:
+                    data = pickle.load(f)
 
-            self.tiles_rgb = data["tiles_rgb"]
-            self.tiles_lab_mean = data["tiles_lab_mean"]
-            self.index = self._build_faiss(self.tiles_lab_mean)
-            return
+                self.tiles_rgb = data["tiles_rgb"]
+                self.tiles_lab_mean = data["tiles_lab_mean"]
+                self.index = self._build_faiss(self.tiles_lab_mean)
+                return
+            except Exception as e:
+                print(f"⚠ Cache load failed — rebuilding. Reason: {e}")
 
         # ------------------------------
-        # 2. CACHE DOESN’T EXIST → BUILD
+        # 2. CACHE MISSING → BUILD
         # ------------------------------
-        print("⚠ No tile cache found — extracting tiles from CIFAR-100 (one-time cost)")
+        print("⚠ No valid tile cache found — extracting CIFAR-100 tiles (one-time cost)")
 
-        ds = load_dataset("cifar100", split="train")
+        try:
+            ds = load_dataset("cifar100", split="train")
+        except Exception as e:
+            raise RuntimeError(f"Failed to load CIFAR-100 dataset: {e}")
+
+        if sample_size > len(ds):
+            raise ValueError(f"Requested {sample_size} tiles but CIFAR-100 only has {len(ds)} images.")
 
         tiles = []
         means = []
 
         for i in range(sample_size):
-            arr = np.array(ds[i]["img"].convert("RGB"), dtype=np.uint8)
+            img = ds[i]["img"]
+            if img is None:
+                raise RuntimeError(f"Dataset returned a None image at index {i}")
+
+            arr = np.array(img.convert("RGB"), dtype=np.uint8)
 
             # Compute LAB means
-            lab = fast_rgb2lab(arr)
-            mean_lab = lab.mean(axis=(0, 1))
+            try:
+                lab = fast_rgb2lab(arr)
+            except Exception:
+                raise RuntimeError(f"fast_rgb2lab failed on tile index {i}")
 
             tiles.append(arr)
-            means.append(mean_lab)
+            means.append(lab.mean(axis=(0, 1)))
 
-            # Progress indicator
-            if (i + 1) % 2000 == 0:
+            # Optional progress printing
+            if (i + 1) % 2000 == 0 or (i + 1) == sample_size:
                 print(f"  → processed {i+1}/{sample_size} tiles")
 
         tiles = np.array(tiles)
@@ -71,14 +117,20 @@ class TileManager:
         # Build FAISS index
         index = self._build_faiss(means)
 
-        # Save cache
-        with open(cache_file, "wb") as f:
-            pickle.dump({
-                "tiles_rgb": tiles,
-                "tiles_lab_mean": means,
-            }, f)
-
-        print(f"✓ Saved tile cache → {cache_file}")
+        # Save cache safely
+        try:
+            with open(cache_file, "wb") as f:
+                pickle.dump(
+                    {
+                        "tiles_rgb": tiles,
+                        "tiles_lab_mean": means,
+                    },
+                    f,
+                    protocol=pickle.HIGHEST_PROTOCOL,
+                )
+            print(f"✓ Saved tile cache → {cache_file}")
+        except Exception as e:
+            print(f"⚠ Failed to save tile cache: {e}")
 
         self.tiles_rgb = tiles
         self.tiles_lab_mean = means
@@ -88,16 +140,53 @@ class TileManager:
     # BUILD FAISS INDEX
     # -------------------------------------------------------
     def _build_faiss(self, vectors):
+        """
+        Create a FAISS L2 index from N×3 LAB feature vectors.
+
+        Parameters
+        ----------
+        vectors : np.ndarray
+            Array of shape (N, 3), LAB means for each tile.
+
+        Returns
+        -------
+        faiss.IndexFlatL2
+        """
+        if vectors is None or vectors.ndim != 2 or vectors.shape[1] != 3:
+            raise ValueError(f"Invalid feature vector shape for FAISS: {vectors.shape}")
+
         dim = vectors.shape[1]
         index = faiss.IndexFlatL2(dim)
         index.add(vectors.astype("float32"))
         return index
 
     # -------------------------------------------------------
-    # LOOKUP TILE USING FAISS
+    # NEAREST TILE LOOKUP
     # -------------------------------------------------------
     def lookup_tiles(self, cell_means):
+        """
+        Search FAISS index for the nearest tile for each grid cell.
+
+        Parameters
+        ----------
+        cell_means : np.ndarray
+            LAB mean values for each grid cell.
+
+        Returns
+        -------
+        np.ndarray
+            Flattened tile indices (one per grid cell).
+        """
+        if self.index is None:
+            raise RuntimeError("FAISS index not built. Call load() first.")
+
         cell_means = np.asarray(cell_means, dtype="float32")
+
+        if cell_means.ndim != 2 or cell_means.shape[1] != 3:
+            raise ValueError(
+                f"Expected cell_means shape (N, 3), got {cell_means.shape}"
+            )
+
         _, idxs = self.index.search(cell_means, 1)
         return idxs.flatten()
 
@@ -106,20 +195,35 @@ class TileManager:
     # -------------------------------------------------------
     def prepare_scaled_tiles(self, cell_w, cell_h):
         """
-        Resize all tiles only when required.
+        Resize all tiles to match a grid cell size.
+        This is cached — resizing happens only when dimensions change.
+
+        Parameters
+        ----------
+        cell_w : int
+            Target cell width.
+        cell_h : int
+            Target cell height.
         """
+
+        if self.tiles_rgb is None:
+            raise RuntimeError("Tiles not loaded. Call load() first.")
+
         if (
             self.pre_scaled_tiles is not None
             and self.pre_scaled_tiles.shape[1] == cell_h
             and self.pre_scaled_tiles.shape[2] == cell_w
         ):
-            return  # already cached for this size
+            return  # Already resized
 
         print(f"Resizing {len(self.tiles_rgb)} tiles → {cell_w}×{cell_h}")
 
         out = []
-        for tile in self.tiles_rgb:
-            resized = cv2.resize(tile, (cell_w, cell_h), interpolation=cv2.INTER_NEAREST)
-            out.append(resized)
+        for i, tile in enumerate(self.tiles_rgb):
+            try:
+                resized = cv2.resize(tile, (cell_w, cell_h), interpolation=cv2.INTER_NEAREST)
+                out.append(resized)
+            except Exception:
+                raise RuntimeError(f"Tile resize failed at index {i}")
 
         self.pre_scaled_tiles = np.array(out)

mosaic_generator/utils.py

@@ -1,32 +1,57 @@
-# utils.py
+"""
+utils.py
+
+Low-level utility functions used across the mosaic generator.
+Includes:
+- Numba-accelerated RGB → LAB conversion
+- A safe wrapper for ensuring correct image dtype and shape
+"""
+
 import numpy as np
 import cv2
 from numba import njit
 
+
+# ----------------------------------------------------------------------
+# NUMBA RGB → LAB (HIGH SPEED)
+# ----------------------------------------------------------------------
 @njit
 def fast_rgb2lab_numba(rgb):
-    """Fast Numba-based approximate RGB→LAB conversion."""
+    """
+    Fast approximate RGB → LAB conversion using Numba JIT.
+
+    Parameters
+    ----------
+    rgb : np.ndarray
+        Float32 array of shape (H, W, 3) in [0, 255].
+
+    Returns
+    -------
+    np.ndarray
+        LAB array of shape (H, W, 3) (float32).
+    """
     R = rgb[..., 0] / 255.0
     G = rgb[..., 1] / 255.0
     B = rgb[..., 2] / 255.0
 
-    # sRGB to XYZ
+    # sRGB → linear RGB
     def f(c):
         return np.where(c > 0.04045, ((c + 0.055) / 1.055) ** 2.4, c / 12.92)
 
     R = f(R); G = f(G); B = f(B)
 
-    X = 0.4124*R + 0.3576*G + 0.1805*B
-    Y = 0.2126*R + 0.7152*G + 0.0722*B
-    Z = 0.0193*R + 0.1192*G + 0.9505*B
+    # Linear RGB → XYZ
+    X = 0.4124 * R + 0.3576 * G + 0.1805 * B
+    Y = 0.2126 * R + 0.7152 * G + 0.0722 * B
+    Z = 0.0193 * R + 0.1192 * G + 0.9505 * B
 
-    # Normalize by D65 white point
+    # Normalize by D65
     X /= 0.95047
     Z /= 1.08883
 
-    # XYZ → LAB
+    # XYZ → LAB helper
     def g(t):
-        return np.where(t > 0.008856, t ** (1/3), 7.787*t + 16/116)
+        return np.where(t > 0.008856, t ** (1/3), 7.787 * t + 16/116)
 
     fx = g(X); fy = g(Y); fz = g(Z)
 
@@ -40,6 +65,41 @@ def fast_rgb2lab_numba(rgb):
     out[..., 2] = b
     return out
 
+
+# ----------------------------------------------------------------------
+# SAFE WRAPPER
+# ----------------------------------------------------------------------
 def fast_rgb2lab(img_rgb):
-    """Wrapper to ensure correct format."""
+    """
+    Safe wrapper for Numba LAB conversion.
+
+    Parameters
+    ----------
+    img_rgb : np.ndarray
+        RGB image, shape (H, W, 3), dtype uint8 or float32.
+
+    Returns
+    -------
+    np.ndarray
+        LAB image of shape (H, W, 3), dtype float32.
+
+    Raises
+    ------
+    ValueError
+        If the input image is not a valid RGB array.
+
+    Notes
+    -----
+    - Numba does NOT allow Python exceptions inside the JIT function.
+      Therefore, validation happens here before calling Numba.
+    """
+    if img_rgb is None or not isinstance(img_rgb, np.ndarray):
+        raise ValueError("fast_rgb2lab(): expected a NumPy array.")
+
+    if img_rgb.ndim != 3 or img_rgb.shape[2] != 3:
+        raise ValueError(
+            f"fast_rgb2lab(): expected image shape (H, W, 3), got {img_rgb.shape}"
+        )
+
+    # Ensure float32 for Numba kernel
     return fast_rgb2lab_numba(img_rgb.astype(np.float32))