testing.py

import os, time, glob
import numpy as np
from PIL import Image, ImageDraw
import gradio as gr

DEFAULT_TILE_DIR = "./tiles"
EXAMPLES_DIR     = "./examples"

# ---- helpers ----
def norm_path(p: str) -> str:
    return os.path.abspath(os.path.expanduser((p or "").rstrip("/")))

def pil_to_np(img: Image.Image) -> np.ndarray:
    if img.mode != "RGB": img = img.convert("RGB")
    return np.asarray(img).astype(np.float32) / 255.0

def np_to_pil(arr: np.ndarray) -> Image.Image:
    arr = np.clip(arr * 255.0, 0, 255).astype(np.uint8)
    return Image.fromarray(arr)

def crop_to_grid(img_np, ty, tx):
    H, W, _ = img_np.shape
    ch, cw = max(1, H // ty), max(1, W // tx)
    new_h, new_w = ch * ty, cw * tx
    cropped = img_np[:new_h, :new_w, :]
    # Optional debug print:
    # print(f"[crop_to_grid] Original=({H},{W}), Cropped=({new_h},{new_w}), Cell=({ch},{cw})")
    return cropped, ch, cw

def mse(a, b): return float(np.mean((a - b) ** 2))

def simple_ssim(a, b):
    g1, g2 = a.mean(axis=2), b.mean(axis=2)
    mu1, mu2 = g1.mean(), g2.mean()
    var1, var2 = g1.var(), g2.var()
    cov = ((g1 - mu1) * (g2 - mu2)).mean()
    C1, C2 = 0.01**2, 0.03**2
    den = (mu1**2+mu2**2+C1)*(var1+var2+C2)
    return float(((2*mu1*mu2+C1)*(2*cov+C2))/den) if den else 0.0

# ---- grid stats ----
def cell_stats_vec(img_np, ty, tx):
    """Vectorized: reshape and reduce without Python loops."""
    cropped, ch, cw = crop_to_grid(img_np, ty, tx)
    grid = cropped.reshape(ty, ch, tx, cw, 3).swapaxes(1,2)  # (ty, tx, ch, cw, 3)
    return grid.mean((2,3)), np.median(grid,(2,3)), cropped, ch, cw

def cell_stats_loop(img_np, ty, tx):
    """Loop-based: compute mean/median per cell with explicit loops."""
    cropped, ch, cw = crop_to_grid(img_np, ty, tx)
    means = np.zeros((ty, tx, 3), dtype=np.float32)
    meds  = np.zeros((ty, tx, 3), dtype=np.float32)
    for i in range(ty):
        for j in range(tx):
            block = cropped[i*ch:(i+1)*ch, j*cw:(j+1)*cw, :]
            means[i, j] = block.mean(axis=(0, 1))
            meds[i, j]  = np.median(block, axis=(0, 1))
    return means, meds, cropped, ch, cw

# ---- tiles ----
def list_images(folder):
    folder = norm_path(folder)
    exts = ("jpg","jpeg","png","bmp","webp")
    files=[]
    for e in exts: files += glob.glob(os.path.join(folder, f"*.{e}"))
    return sorted(set(files))

def load_tiles(tile_dir,h,w,stat="Mean"):
    files = list_images(norm_path(tile_dir))
    if not files: raise ValueError(f"No images in {tile_dir}")
    tiles,cols=[],[]
    for f in files:
        try:
            arr = pil_to_np(Image.open(f).resize((w,h),Image.BOX))
            tiles.append(arr)
            cols.append(np.median(arr,(0,1)) if stat=="Median" else arr.mean((0,1)))
        except:
            continue
    return np.stack(tiles), np.stack(cols)

# ---- mapping ----
def nearest_indices(cell_cols,tile_cols):
    ty,tx,_=cell_cols.shape
    cm=cell_cols.reshape(-1,3)
    d=((cm[:,None,:]-tile_cols[None,:,:])**2).sum(2)
    return d.argmin(1).reshape(ty,tx)

def assemble(indices,tiles,ch,cw):
    ty,tx=indices.shape
    out=np.zeros((ty*ch,tx*cw,3), dtype=np.float32)
    for i in range(ty):
        for j in range(tx):
            out[i*ch:(i+1)*ch, j*cw:(j+1)*cw, :] = tiles[indices[i,j]]
    return out

# ---- preview overlay ----
def add_tile_grid_overlay(cropped_np, ch, cw, ty, tx, color=(255, 0, 0), width=2):
    """Draw grid lines over the cropped preview to visualize the ty×tx tiling."""
    H, W, _ = cropped_np.shape
    preview = np_to_pil(cropped_np.copy())
    draw = ImageDraw.Draw(preview)
    for i in range(1, ty):  # horizontal
        y = i * ch
        draw.line([(0, y), (W - 1, y)], fill=color, width=width)
    for j in range(1, tx):  # vertical
        x = j * cw
        draw.line([(x, 0), (x, H - 1)], fill=color, width=width)
    return preview

# ---- pipelines ----
def build_mosaic(pil_img, tiles_side, tile_dir, cell_stat="Mean", tile_stat="Mean"):
    """Vectorized default, returns (cropped+grid preview, mosaic, metrics, elapsed)."""
    t0=time.time()
    img_np=pil_to_np(pil_img)
    means,meds,cropped,ch,cw=cell_stats_vec(img_np,tiles_side,tiles_side)
    cell_cols=means if cell_stat=="Mean" else meds
    tbank,tcols=load_tiles(tile_dir,ch,cw,stat=tile_stat)
    idx=nearest_indices(cell_cols,tcols)
    mosaic=assemble(idx,tbank,ch,cw)
    m,s=mse(cropped,mosaic),simple_ssim(cropped,mosaic)
    overlay_preview = add_tile_grid_overlay(cropped, ch, cw, tiles_side, tiles_side, color=(255,0,0), width=2)
    return overlay_preview, np_to_pil(mosaic), f"MSE: {m:.6f} | SSIM*: {s:.6f}", f"{time.time()-t0:.3f}s"

def build_mosaic_mode(pil_img, tiles_side, tile_dir, cell_stat="Mean", tile_stat="Mean", use_vectorized=True):
    """Switchable vec/loop (for benchmark). Returns (mosaic, MSE, SSIM, elapsed)."""
    t0=time.time()
    img_np=pil_to_np(pil_img)
    if use_vectorized:
        means,meds,cropped,ch,cw=cell_stats_vec(img_np,tiles_side,tiles_side)
    else:
        means,meds,cropped,ch,cw=cell_stats_loop(img_np,tiles_side,tiles_side)
    cell_cols=means if cell_stat=="Mean" else meds
    tbank,tcols=load_tiles(tile_dir,ch,cw,stat=tile_stat)
    idx=nearest_indices(cell_cols,tcols)
    mosaic=assemble(idx,tbank,ch,cw)
    m,s=mse(cropped,mosaic),simple_ssim(cropped,mosaic)
    return np_to_pil(mosaic), m, s, time.time()-t0

# ---- benchmark helpers ----
def _time_run(pil_img, grid, tile_dir, cell_stat, tile_stat, vectorized: bool):
    _mos, _m, _s, elapsed = build_mosaic_mode(
        pil_img, tiles_side=grid, tile_dir=tile_dir,
        cell_stat=cell_stat, tile_stat=tile_stat, use_vectorized=vectorized
    )
    return elapsed

def run_benchmark_table(pil_img, grids=(16,32,64), tile_dir=DEFAULT_TILE_DIR, cell_stat="Mean", tile_stat="Mean"):
    rows=[]
    for g in grids:
        t_vec  = _time_run(pil_img, g, tile_dir, cell_stat, tile_stat, True)
        t_loop = _time_run(pil_img, g, tile_dir, cell_stat, tile_stat, False)
        speed  = t_loop / max(t_vec, 1e-9)
        rows.append((g, t_vec, t_loop, speed))
    header = "Grid | Vectorized Time (s) | Loop Time (s) | Speedup (loop/vec)"
    sep    = "-" * len(header)
    lines  = [header, sep] + [f"{g}×{g} | {tv:.3f} | {tl:.3f} | {sp:.2f}×" for g,tv,tl,sp in rows]
    table  = "\n".join(lines)
    analysis = (
        "Analysis:\n"
        "- Runtime increases with grid size because cells grow as (tiles_per_side)^2.\n"
        "- Matching cost ~ O(#cells × #tiles). Doubling tiles_per_side ≈ 4× more cells.\n"
        "- Vectorized NumPy is consistently faster (often ~10–20×) than loops due to optimized C.\n"
        "- Larger grids improve visual detail but cost more time; pick based on quality vs speed."
    )
    return table, analysis

# ---- gradio ui ----
with gr.Blocks(title="CS 5130 – Mosaic") as demo:
    gr.Markdown("### Image Mosaic\nUpload on the left → Mosaic result on the right.")

    # Remember last uploaded image so Benchmark can reuse it
    last_image = gr.State()

    # Build Mosaic tab
    with gr.Tab("Build Mosaic"):
        with gr.Row():
            in_img        = gr.Image(type="pil", label="Upload main image")
            orig_preview  = gr.Image(type="pil", label="Cropped-for-grid (preview)")
            out_mosa      = gr.Image(type="pil", label="Mosaic result")

        tile_dir_box = gr.Textbox(value=DEFAULT_TILE_DIR, label="Tiles folder path")
        grid         = gr.Slider(8,128,value=32,step=1,label="Tiles per side")
        cell_stat    = gr.Radio(["Mean","Median"],value="Mean",label="Cell color")
        tile_stat    = gr.Radio(["Mean","Median"],value="Mean",label="Tile color")
        btn          = gr.Button("Generate Mosaic")

        metrics = gr.Textbox(label="Similarity (MSE/SSIM*)", interactive=False)
        perf    = gr.Textbox(label="Elapsed (s)", interactive=False)

        def ui_run(img, tps, cstat, tstat, tdir):
            if img is None:
                return None, None, "Upload an image.", ""
            try:
                tdir = norm_path(tdir or DEFAULT_TILE_DIR)
                orig, mos, m, s = build_mosaic(img, int(tps), tdir, cstat, tstat)
                return orig, mos, m, s
            except Exception as e:
                return None, None, f"Error: {e}", ""

        btn.click(
            ui_run,
            inputs=[in_img, grid, cell_stat, tile_stat, tile_dir_box],
            outputs=[orig_preview, out_mosa, metrics, perf]
        )

        # Keep track of last uploaded image for Benchmark reuse
        def _remember_image(img):
            return img
        in_img.change(_remember_image, inputs=[in_img], outputs=[last_image])

        # Clickable examples (populate only the input image)
        exdir = norm_path(EXAMPLES_DIR)
        if os.path.isdir(exdir):
            ex_files = [os.path.join(exdir, f) for f in os.listdir(exdir)
                        if f.lower().endswith((".jpg",".jpeg",".png"))][:6]
            if ex_files:
                gr.Examples(ex_files, inputs=[in_img], label="Examples (Build)")

    # Benchmark tab
    with gr.Tab("Benchmark"):
        gr.Markdown(
            "Compare **Vectorized** vs **Loop** timings across grid sizes. "
            "Use an example or click **Use uploaded image** to reuse the image from the Build tab."
        )
        bench_img   = gr.Image(type="pil", label="Main image for benchmark (RGB)")
        use_current = gr.Button("Use uploaded image")  # copies from Build tab state

        def _use_uploaded(img):
            if img is None:
                return None
            return img

        use_current.click(_use_uploaded, inputs=[last_image], outputs=[bench_img])

        grids_box  = gr.Textbox(value="16, 32, 64", label="Grid sizes (comma-separated)")
        b_cell     = gr.Radio(["Mean","Median"], value="Mean", label="Cell color")
        b_tile     = gr.Radio(["Mean","Median"], value="Mean", label="Tile color")
        b_folder   = gr.Textbox(value=DEFAULT_TILE_DIR, label="Tile folder path")
        run_bench  = gr.Button("Run Benchmark")

        bench_table = gr.Code(label="Results Table")
        bench_note  = gr.Textbox(label="Brief Analysis", lines=6)

        def ui_bench(img, grids_csv, cstat, tstat, tdir):
            if img is None: return "Please upload/select an image (or click 'Use uploaded image').", ""
            try:
                grids = [int(x.strip()) for x in grids_csv.split(",") if x.strip()]
            except:
                grids = [16,32,64]
            tdir = norm_path(tdir or DEFAULT_TILE_DIR)
            table, analysis = run_benchmark_table(img, tuple(grids), tdir, cstat, tstat)
            return table, analysis

        run_bench.click(
            ui_bench,
            inputs=[bench_img, grids_box, b_cell, b_tile, b_folder],
            outputs=[bench_table, bench_note]
        )

        # Benchmark examples (optional)
        if os.path.isdir(exdir):
            ex_files_bench = [os.path.join(exdir, f) for f in os.listdir(exdir)
                              if f.lower().endswith((".jpg",".jpeg",".png"))][:6]
            if ex_files_bench:
                gr.Examples(ex_files_bench, inputs=[bench_img], label="Examples (Benchmark)")

if __name__=="__main__":
    demo.launch(allowed_paths=[norm_path(DEFAULT_TILE_DIR), norm_path(EXAMPLES_DIR)])