Replace Space with latest LAB1 code

app.py

@@ -1,4 +1,4 @@
-import io, time, zipfile, math
+import time
 from pathlib import Path
 from typing import List, Tuple, Optional, Dict
 
@@ -9,17 +9,17 @@ from skimage.metrics import structural_similarity as ssim
 from skimage.color import rgb2lab
 from sklearn.cluster import KMeans
 
-# ---- Hugging Face dataset: hard-wired ----
+# ---- Hugging Face dataset: hard-wired to CIFAR-100 ----
 from datasets import load_dataset
 
-HF_DATASET = "benjamin-paine/imagenet-1k-32x32"  # always use this
+HF_DATASET = "uoft-cs/cifar100"  # tiles come from here
 HF_SPLIT = "train"
-TILE_LIMIT = 1500            # cap tiles to keep mapping fast; raise if you want
-BASE_TILE_SIZE = 32          # dataset images are 32x32
+TILE_LIMIT = 0                # cap to keep mapping fast (adjust as needed)
+BASE_TILE_SIZE = 32              # CIFAR-100 images are 32x32
 
 # Global caches
-_TILES_RAW_32: Optional[List[np.ndarray]] = None     # list of 32x32 RGB uint8 arrays
-_TILE_CACHE_BY_SIZE: Dict[int, Tuple[List[np.ndarray], np.ndarray]] = {}  # cell_size -> (tiles_resized, tiles_lab_means)
+_TILES_RAW_32: Optional[List[np.ndarray]] = None  # list of 32x32 RGB uint8 arrays
+_TILE_CACHE_BY_SIZE: Dict[int, Tuple[List[np.ndarray], np.ndarray]] = {}  # size -> (tiles_resized, lab_means)
 
 # =======================
 # Image utils
@@ -53,10 +53,10 @@ def mse(a: np.ndarray, b: np.ndarray) -> float:
     return float(np.mean((a.astype(np.float32) - b.astype(np.float32))**2))
 
 # =======================
-# Load & cache tiles from HF dataset (once)
+# CIFAR-100: load & cache base tiles (32x32)
 # =======================
 def _load_tiles_raw_32(limit: int = TILE_LIMIT) -> List[np.ndarray]:
-    """Load 32x32 tiles (RGB uint8) from benjamin-paine/imagenet-1k-32x32."""
+    """Load 32x32 tiles (RGB uint8) from CIFAR-100 (uoft-cs/cifar100)."""
     global _TILES_RAW_32
     if _TILES_RAW_32 is not None:
         return _TILES_RAW_32
@@ -64,15 +64,16 @@ def _load_tiles_raw_32(limit: int = TILE_LIMIT) -> List[np.ndarray]:
     ds = load_dataset(HF_DATASET, split=HF_SPLIT)
     tiles = []
     for i, ex in enumerate(ds):
-        if "image" not in ex:
+        # CIFAR-100 column name is "img"
+        if "img" not in ex:
             continue
-        img: Image.Image = ex["image"].convert("RGB")
-        # dataset already 32x32; enforce in case
+        img: Image.Image = ex["img"].convert("RGB")
         if img.size != (BASE_TILE_SIZE, BASE_TILE_SIZE):
             img = img.resize((BASE_TILE_SIZE, BASE_TILE_SIZE), Image.BILINEAR)
         tiles.append(np.asarray(img))
-        if limit and len(tiles) >= limit:
+        if limit > 0 and len(tiles) >= limit:
             break
+
     if len(tiles) == 0:
         raise gr.Error(f"No tiles loaded from {HF_DATASET}.")
     _TILES_RAW_32 = tiles
@@ -85,22 +86,20 @@ def _average_color_lab(tile: np.ndarray) -> np.ndarray:
 def _tiles_for_cell_size(cell_size: int) -> Tuple[List[np.ndarray], np.ndarray]:
     """
     Return (tiles_resized, tiles_lab_means) for the requested cell size.
-    Caches results to avoid recompute on every click.
+    Caches results to avoid recomputation on each run.
     """
     if cell_size in _TILE_CACHE_BY_SIZE:
         return _TILE_CACHE_BY_SIZE[cell_size]
 
     raw_tiles = _load_tiles_raw_32()
-    # Resize to cell_size if needed
     if cell_size == BASE_TILE_SIZE:
         tiles_resized = raw_tiles
     else:
-        tiles_resized = [np.asarray(Image.fromarray(t).resize((cell_size, cell_size), Image.BILINEAR))
-                         for t in raw_tiles]
-
-    # LAB means (size does not matter much for mean, but compute on resized set)
+        tiles_resized = [
+            np.asarray(Image.fromarray(t).resize((cell_size, cell_size), Image.BILINEAR))
+            for t in raw_tiles
+        ]
     tiles_lab = np.array([_average_color_lab(t) for t in tiles_resized], dtype=np.float32)
-
     _TILE_CACHE_BY_SIZE[cell_size] = (tiles_resized, tiles_lab)
     return tiles_resized, tiles_lab
 
@@ -167,11 +166,11 @@ def segment_preview(src: np.ndarray, cell: int) -> np.ndarray:
     return out.astype(np.uint8)
 
 # =======================
-# Full pipeline (tiles always from HF dataset)
+# Full pipeline (tiles always from CIFAR-100)
 # =======================
 def build_mosaic(
     input_image: Image.Image,
-    cell_size: int = 32,           # default 32 to match dataset; you can change
+    cell_size: int = 32,
     use_vectorized: bool = True,
     quant_k: int = 0,
     similarity_metric: str = "SSIM",
@@ -196,7 +195,7 @@ def build_mosaic(
         mean_rgb, R, C = grid_mean_colors_loop(src, cell_size)
     t_grid = time.perf_counter() - t0
 
-    # 4) Tiles from HF dataset (cached & resized to cell_size)
+    # 4) Tiles from CIFAR-100 (cached & resized to cell_size)
     tiles, tiles_lab = _tiles_for_cell_size(cell_size)
 
     # 5) Map & build mosaic
@@ -259,12 +258,12 @@ if not (EXAMPLES_DIR / "gradient1.png").exists():
     grad1 = np.dstack([g1, np.flipud(g1).copy(), np.roll(g1, 160, axis=1)])
     Image.fromarray(grad1).save(EXAMPLES_DIR/"gradient1.png")
 
-with gr.Blocks(title="Image Mosaic (ImageNet32 tiles)", css="footer {visibility: hidden}") as demo:
+with gr.Blocks(title="Image Mosaic (CIFAR-100 tiles)", css="footer {visibility: hidden}") as demo:
     gr.Markdown(
         f"""
         # 🧩 Image Mosaic Generator (tiles from `{HF_DATASET}`)
-        - Tiles are auto-loaded from **Hugging Face** dataset: `{HF_DATASET}` (split `{HF_SPLIT}`, limit {TILE_LIMIT}).  
-        - Upload an image and generate a mosaic **immediately** — no extra tile setup.  
+        - Tiles auto-loaded from **Hugging Face** dataset: `{HF_DATASET}` (split `{HF_SPLIT}`, limit {TILE_LIMIT}).  
+        - Upload an image and generate a mosaic immediately.
         """
     )
     with gr.Row():
@@ -275,7 +274,8 @@ with gr.Blocks(title="Image Mosaic (ImageNet32 tiles)", css="footer {visibility:
                 inputs=[inp],
                 label="Example"
             )
-            cell = gr.Slider(16, 64, value=32, step=2, label="Grid cell size (px)")
+            cell = gr.Slider(16, 128, value=32, step=2, label="Grid cell size (px)")
+
             quant_k = gr.Slider(0, 24, value=0, step=1, label="Optional color quantization (k-means K)")
             similarity = gr.Radio(choices=["SSIM", "MSE"], value="SSIM", label="Similarity metric")
             vec = gr.Checkbox(value=True, label="Use vectorized NumPy (uncheck for loop baseline)")
@@ -307,6 +307,5 @@ if __name__ == "__main__":
     try:
         _load_tiles_raw_32(TILE_LIMIT)
     except Exception as e:
-        # Gradio will still start; you'll see an error if tiles can't be loaded
         print("Warning: failed to preload tiles:", e)
     demo.launch()