Add view toggle: Nearest Neighbor (baseline + 10 NN) vs Results (3 baselines + Ours, top 3 prompts)

app.py

@@ -6,103 +6,113 @@ import os
 
 import gradio as gr
 
-from dataset_loader import list_samples, get_noise_demo_paths
+from dataset_loader import list_samples, get_nn_demo_paths, get_results_demo_paths
 
 
 SAMPLE_IDS = list_samples()
 TOP1_ID = SAMPLE_IDS[0] if SAMPLE_IDS else None
 
 
-def build_noise_demo(sample_id: str | None):
-    """Returns for each of 3 blocks: prompt_md, baseline (spec, bg, fg, m), method (spec, bg, fg, m)."""
+def build_nn_view(sample_id: str | None):
+    """NN view: [Baseline] [NN1] [NN2] ... [NN10]. Each NN has spec + m_wav."""
     if not sample_id:
-        return (None,) * 27
-    data = get_noise_demo_paths(sample_id)
+        return (None,) * (4 + 10 * 2)
+    data = get_nn_demo_paths(sample_id, top_k=10)
+    out = []
+    bl = data.get("baseline", {})
+    out.extend([bl.get("spec"), bl.get("bg_wav"), bl.get("fg_wav"), bl.get("m_wav")])
+    for nn in data.get("nn_list", [])[:10]:
+        out.extend([nn.get("spec"), nn.get("m_wav")])
+    while len(out) < 4 + 20:
+        out.append(None)
+    return tuple(out[: 4 + 20])
+
+
+def build_results_view(sample_id: str | None):
+    """Results view: 3 blocks, each with Baseline (orig), Gaussian, Youtube-noise, Ours."""
+    if not sample_id:
+        return (None,) * (3 * (1 + 4 * 4))
+    data = get_results_demo_paths(sample_id)
     out = []
     for i in range(1, 4):
         block = data.get(f"block{i}", {})
         prompt = block.get("prompt", "") or ""
         out.append(f"**Prompt:** {prompt}" if prompt else "")
-        bl = block.get("baseline", {})
-        out.extend([bl.get("spec"), bl.get("bg_wav"), bl.get("fg_wav"), bl.get("m_wav")])
-        nn = block.get("nn", {})
-        out.extend([nn.get("spec"), nn.get("bg_wav"), nn.get("fg_wav"), nn.get("m_wav")])
+        for key in ("baseline_original", "baseline_gaussian", "baseline_youtube", "ours"):
+            b = block.get(key, {})
+            out.extend([b.get("spec"), b.get("bg_wav"), b.get("fg_wav"), b.get("m_wav")])
     return tuple(out)
 
 
-with gr.Blocks(title="NearestNeighbor Audio Demo", css=".gradio-container { max-width: 1200px; }") as app:
+with gr.Blocks(title="NearestNeighbor Audio Demo", css=".gradio-container { max-width: 1400px; }") as app:
     gr.Markdown("# NearestNeighbor Audio Demo")
     gr.Markdown("Data from [AE-W/batch_outputs](https://huggingface.co/datasets/AE-W/batch_outputs)")
 
+    view_radio = gr.Radio(
+        choices=["Nearest Neighbor", "Results"],
+        value="Nearest Neighbor",
+        label="View",
+    )
+    noise_dd = gr.Dropdown(choices=SAMPLE_IDS, value=TOP1_ID, label="Noise (ID)")
+
     gr.Markdown("""
 **How to read the IDs**
 - **Numeric IDs** (e.g. `00_000357`) come from the **SONYC** dataset.
 - **IDs starting with `fold`** come from the **UrbanSound8k** dataset.
 
-**Audio labels**
-- **BG** = background noise  
-- **FG** = generated foreground sound  
-- **Mix** = BG + FG (mixed)
+**Audio labels**: **BG** = background noise | **FG** = generated foreground | **Mix** = BG + FG
 """)
 
-    noise_dd = gr.Dropdown(
-        choices=SAMPLE_IDS,
-        value=TOP1_ID,
-        label="Noise (ID)",
-    )
-
-    def prompt_block_ui(block_label: str):
-        """One block: prompt text, then Baseline and Ours side by side (image each), audios below each."""
-        with gr.Group():
-            gr.Markdown(f"### {block_label}")
-            prompt_md = gr.Markdown(value="")
-            with gr.Row():
-                with gr.Column(scale=1):
-                    gr.Markdown("**Baseline**")
-                    bl_img = gr.Image(label="Baseline", show_label=True, height=320)
-                    with gr.Row():
-                        bl_bg = gr.Audio(label="BG", show_label=True)
-                        bl_fg = gr.Audio(label="FG", show_label=True)
-                        bl_m = gr.Audio(label="Mix", show_label=True)
-                with gr.Column(scale=1):
-                    gr.Markdown("**Ours**")
-                    nn_img = gr.Image(label="Ours", show_label=True, height=320)
-                    with gr.Row():
-                        nn_bg = gr.Audio(label="BG", show_label=True)
-                        nn_fg = gr.Audio(label="FG", show_label=True)
-                        nn_m = gr.Audio(label="Mix", show_label=True)
-        return prompt_md, bl_img, bl_bg, bl_fg, bl_m, nn_img, nn_bg, nn_fg, nn_m
-
-    # Block 1: Prompt 1 -> Baseline -> Ours
-    p1_md, bl1_img, bl1_bg, bl1_fg, bl1_m, nn1_img, nn1_bg, nn1_fg, nn1_m = prompt_block_ui("Prompt 1")
-    gr.Markdown("---")
-    # Block 2
-    p2_md, bl2_img, bl2_bg, bl2_fg, bl2_m, nn2_img, nn2_bg, nn2_fg, nn2_m = prompt_block_ui("Prompt 2")
-    gr.Markdown("---")
-    # Block 3
-    p3_md, bl3_img, bl3_bg, bl3_fg, bl3_m, nn3_img, nn3_bg, nn3_fg, nn3_m = prompt_block_ui("Prompt 3")
-
-    all_outputs = [
-        p1_md, bl1_img, bl1_bg, bl1_fg, bl1_m, nn1_img, nn1_bg, nn1_fg, nn1_m,
-        p2_md, bl2_img, bl2_bg, bl2_fg, bl2_m, nn2_img, nn2_bg, nn2_fg, nn2_m,
-        p3_md, bl3_img, bl3_bg, bl3_fg, bl3_m, nn3_img, nn3_bg, nn3_fg, nn3_m,
-    ]
-
-    def on_noise_select(sid):
-        return build_noise_demo(sid)
-
-    noise_dd.change(
-        on_noise_select,
-        inputs=[noise_dd],
-        outputs=all_outputs,
-    )
-
-    def init():
-        return build_noise_demo(TOP1_ID)
-
-    app.load(init, outputs=all_outputs)
+    # ---- NN View: Baseline + 10 NN ----
+    with gr.Column(visible=True) as nn_col:
+        gr.Markdown("### Nearest Neighbor: Baseline + top 10 NN")
+        with gr.Row():
+            with gr.Column(min_width=180):
+                gr.Markdown("**Baseline**")
+                nn_bl_img = gr.Image(label="Spec", show_label=False, height=220)
+                nn_bl_bg = gr.Audio(label="BG", show_label=True)
+                nn_bl_fg = gr.Audio(label="FG", show_label=True)
+                nn_bl_m = gr.Audio(label="Mix", show_label=True)
+        gr.Markdown("**NN1–NN10**")
+        nn_items = []
+        with gr.Row():
+            for i in range(10):
+                with gr.Column(min_width=120):
+                    nn_items.append(gr.Image(label=f"NN{i+1}", show_label=True, height=140))
+                    nn_items.append(gr.Audio(label="Mix", show_label=True))
+        nn_outputs = [nn_bl_img, nn_bl_bg, nn_bl_fg, nn_bl_m] + nn_items
+
+    # ---- Results View: 3 prompts × 4 methods ----
+    with gr.Column(visible=False) as res_col:
+        gr.Markdown("### Results: 3 baselines + Ours (top 3 prompts)")
+        res_outputs = []
+        for i in range(1, 4):
+            with gr.Group():
+                res_p_md = gr.Markdown(value="")
+                res_outputs.append(res_p_md)
+                with gr.Row():
+                    for _ in ["Original", "Gaussian", "Youtube", "Ours"]:
+                        res_outputs.append(gr.Image(height=180))
+                        res_outputs.append(gr.Audio(label="BG"))
+                        res_outputs.append(gr.Audio(label="FG"))
+                        res_outputs.append(gr.Audio(label="Mix"))
+
+    def on_change(sid, view):
+        nn_vals = build_nn_view(sid)
+        res_vals = build_results_view(sid)
+        is_nn = view == "Nearest Neighbor"
+        return (
+            list(nn_vals) + list(res_vals) +
+            [gr.update(visible=is_nn), gr.update(visible=not is_nn)]
+        )
+
+    all_outputs = nn_outputs + res_outputs + [nn_col, res_col]
+
+    noise_dd.change(on_change, inputs=[noise_dd, view_radio], outputs=all_outputs)
+    view_radio.change(on_change, inputs=[noise_dd, view_radio], outputs=all_outputs)
+
+    app.load(lambda: on_change(TOP1_ID, "Nearest Neighbor"), outputs=all_outputs)
 
-# Allow serving files from HF dataset cache (required on Spaces)
 _hf_hub_cache = os.environ.get(
     "HUGGINGFACE_HUB_CACHE",
     os.path.join(os.path.expanduser("~"), ".cache", "huggingface", "hub"),

dataset_loader.py

@@ -59,25 +59,60 @@ def _find_files(inner: str) -> list[str]:
     return [f for f in files if f.startswith(inner + "/")]
 
 
-def get_nn_demo_paths(bid: str) -> dict:
+def _collect_block(file_list: list, folder_prefix: str) -> dict:
+    """From files under folder_prefix, get spec + bg_wav, fg_wav, m_wav."""
+    spec = bg = fg = m = None
+    for f in file_list:
+        if folder_prefix not in f:
+            continue
+        name = f.split("/")[-1]
+        if name.endswith(".png"):
+            spec = f
+        elif name.endswith("_bg.wav"):
+            bg = f
+        elif name.endswith("_fg.wav"):
+            fg = f
+        elif name.endswith("_m.wav"):
+            m = f
+    return {
+        "spec": _download_file(spec) if spec else None,
+        "bg_wav": _download_file(bg) if bg else None,
+        "fg_wav": _download_file(fg) if fg else None,
+        "m_wav": _download_file(m) if m else None,
+    }
+
+
+def get_nn_demo_paths(bid: str, top_k: int = 10) -> dict:
     """
-    For Section 1: [bg] [NN1] [NN2] [NN3].
-    Returns {bg_wav, bg_spec, nn_list: [{fg_wav, spec, prompt, similarity}, ...]}.
+    For NN view: [Baseline] [NN1] [NN2] ... [NN10].
+    Returns {baseline: {spec, bg, fg, m}, nn_list: [{fg_wav, spec, bg_wav, prompt, similarity}, ...]}.
     """
     inner = f"{ROOT_PREFIX}{bid}/{bid}"
     prompts = _load_json_from_repo(f"{inner}/temp_retrieval.json")
     if not prompts:
         prompts = _load_json_from_repo(f"{inner}/natural_prompts.json")
     if not prompts:
-        return {"bg_wav": None, "bg_spec": None, "nn_list": []}
+        return {"baseline": {"spec": None, "bg_wav": None, "fg_wav": None, "m_wav": None}, "nn_list": []}
 
     files = _find_files(inner)
+    baseline_inner = f"{inner}/baseline"
+    baseline_files = _find_files(baseline_inner) if any(f.startswith(baseline_inner) for f in files) else []
+
+    # Baseline: first baseline folder (generated_baseline_01_*)
+    baseline_block = {"spec": None, "bg_wav": None, "fg_wav": None, "m_wav": None}
+    for f in baseline_files:
+        parts = f.replace(baseline_inner + "/", "").split("/")
+        if parts and parts[0].startswith("generated_baseline_01_"):
+            full_prefix = baseline_inner + "/" + parts[0]
+            baseline_block = _collect_block(baseline_files, full_prefix)
+            break
+
     nn_list = []
-    for i, p in enumerate(prompts[:3]):
+    for i, p in enumerate(prompts[:top_k]):
         prompt = p.get("prompt", "")
         sim = p.get("similarity_score", p.get("retrieval_score"))
         gen_prefix = f"generated_{i+1:02d}_"
-        fg_path = bg_path = spec_path = None
+        fg_path = bg_path = m_path = spec_path = None
         for f in files:
             parts = f.replace(inner + "/", "").split("/")
             if len(parts) >= 2 and parts[0].startswith(gen_prefix):
@@ -86,19 +121,20 @@ def get_nn_demo_paths(bid: str) -> dict:
                     fg_path = f
                 elif name.endswith("_bg.wav"):
                     bg_path = f
+                elif name.endswith("_m.wav"):
+                    m_path = f
                 elif name.endswith(".png"):
                     spec_path = f
         nn_list.append({
             "fg_wav": _download_file(fg_path) if fg_path else None,
             "spec": _download_file(spec_path) if spec_path else None,
             "bg_wav": _download_file(bg_path) if bg_path else None,
+            "m_wav": _download_file(m_path) if m_path else None,
             "prompt": prompt,
             "similarity": sim,
         })
 
-    bg_wav = nn_list[0]["bg_wav"] if nn_list else None
-    bg_spec = nn_list[0]["spec"] if nn_list else None
-    return {"bg_wav": bg_wav, "bg_spec": bg_spec, "nn_list": nn_list}
+    return {"baseline": baseline_block, "nn_list": nn_list}
 
 
 def get_noise_demo_paths(bid: str) -> dict:
@@ -117,28 +153,6 @@ def get_noise_demo_paths(bid: str) -> dict:
     if not prompts:
         prompts = []
 
-    def collect_block(file_list: list, folder_prefix: str) -> dict:
-        """From files under folder_prefix, get spec + bg_wav, fg_wav, m_wav."""
-        spec = bg = fg = m = None
-        for f in file_list:
-            if folder_prefix not in f:
-                continue
-            name = f.split("/")[-1]
-            if name.endswith(".png"):
-                spec = f
-            elif name.endswith("_bg.wav"):
-                bg = f
-            elif name.endswith("_fg.wav"):
-                fg = f
-            elif name.endswith("_m.wav"):
-                m = f
-        return {
-            "spec": _download_file(spec) if spec else None,
-            "bg_wav": _download_file(bg) if bg else None,
-            "fg_wav": _download_file(fg) if fg else None,
-            "m_wav": _download_file(m) if m else None,
-        }
-
     # Find baseline folder names generated_baseline_01_*, 02_*, 03_*
     seen = set()
     baseline_folders = []
@@ -152,17 +166,15 @@ def get_noise_demo_paths(bid: str) -> dict:
     result = {}
     for i in range(1, 4):
         prompt_text = prompts[i - 1].get("prompt", "") if i <= len(prompts) else ""
-        # Baseline for this prompt: i-th baseline folder (01, 02, 03)
         bl_prefix = f"generated_baseline_{i:02d}_"
         baseline_block = {"spec": None, "bg_wav": None, "fg_wav": None, "m_wav": None}
         for folder_name, full_prefix in baseline_folders:
             if folder_name.startswith(bl_prefix):
-                baseline_block = collect_block(baseline_files, full_prefix)
+                baseline_block = _collect_block(baseline_files, full_prefix)
                 break
-        # Our method: generated_0{i}_*
         rel_prefix = f"generated_{i:02d}_"
         nn_files = [f for f in files if f.replace(inner + "/", "").startswith(rel_prefix)]
-        nn_block = collect_block(nn_files, rel_prefix)
+        nn_block = _collect_block(nn_files, rel_prefix)
         nn_block["prompt"] = prompt_text
         result[f"block{i}"] = {
             "prompt": prompt_text,
@@ -170,3 +182,81 @@ def get_noise_demo_paths(bid: str) -> dict:
             "nn": nn_block,
         }
     return result
+
+
+def get_results_demo_paths(bid: str) -> dict:
+    """
+    For Results view: 3 blocks (prompts 1-3), each with 4 columns:
+    Baseline (original), Gaussian, Youtube-noise, Ours.
+    """
+    inner = f"{ROOT_PREFIX}{bid}/{bid}"
+    files = _find_files(inner)
+    baseline_inner = f"{inner}/baseline"
+    gaussian_inner = f"{inner}/gaussian_baseline"
+    youtube_inner = f"{inner}/youtube_noise_baseline"
+    baseline_files = _find_files(baseline_inner) if any(f.startswith(baseline_inner) for f in files) else []
+    gaussian_files = _find_files(gaussian_inner) if any(f.startswith(gaussian_inner) for f in files) else []
+    youtube_files = _find_files(youtube_inner) if any(f.startswith(youtube_inner) for f in files) else []
+
+    prompts = _load_json_from_repo(f"{inner}/temp_retrieval.json")
+    if not prompts:
+        prompts = _load_json_from_repo(f"{inner}/natural_prompts.json")
+    if not prompts:
+        prompts = []
+
+    def get_baseline_folders(bl_inner, bl_files):
+        seen = set()
+        folders = []
+        for f in bl_files:
+            parts = f.replace(bl_inner + "/", "").split("/")
+            if parts and parts[0].startswith("generated_baseline_") and parts[0] not in seen:
+                seen.add(parts[0])
+                folders.append((parts[0], bl_inner + "/" + parts[0]))
+        folders.sort(key=lambda x: x[0])
+        return folders
+
+    def get_youtube_folders():
+        seen = set()
+        folders = []
+        for f in youtube_files:
+            parts = f.replace(youtube_inner + "/", "").split("/")
+            if parts and parts[0].startswith("generated_") and parts[0] not in seen:
+                seen.add(parts[0])
+                folders.append((parts[0], youtube_inner + "/" + parts[0]))
+        folders.sort(key=lambda x: x[0])
+        return folders
+
+    baseline_folders = get_baseline_folders(baseline_inner, baseline_files)
+    youtube_folders = get_youtube_folders()
+
+    result = {}
+    for i in range(1, 4):
+        prompt_text = prompts[i - 1].get("prompt", "") if i <= len(prompts) else ""
+        bl_prefix = f"generated_baseline_{i:02d}_"
+        rel_prefix = f"generated_{i:02d}_"
+
+        bl_orig = {"spec": None, "bg_wav": None, "fg_wav": None, "m_wav": None}
+        for fn, fp in baseline_folders:
+            if fn.startswith(bl_prefix):
+                bl_orig = _collect_block(baseline_files, fp)
+                break
+
+        gaussian_block = _collect_block(gaussian_files, gaussian_inner)
+
+        bl_youtube = {"spec": None, "bg_wav": None, "fg_wav": None, "m_wav": None}
+        for fn, fp in youtube_folders:
+            if fn.startswith(rel_prefix):
+                bl_youtube = _collect_block(youtube_files, fp)
+                break
+
+        nn_files = [f for f in files if f.replace(inner + "/", "").startswith(rel_prefix)]
+        ours_block = _collect_block(nn_files, rel_prefix)
+
+        result[f"block{i}"] = {
+            "prompt": prompt_text,
+            "baseline_original": bl_orig,
+            "baseline_gaussian": gaussian_block,
+            "baseline_youtube": bl_youtube,
+            "ours": ours_block,
+        }
+    return result