NN view: NN1-NN10 from baseline (prompt order), no separate baseline block

app.py

@@ -14,18 +14,16 @@ TOP1_ID = SAMPLE_IDS[0] if SAMPLE_IDS else None
 
 
 def build_nn_view(sample_id: str | None):
-    """NN view: [Baseline] [NN1] [NN2] ... [NN10]. Each NN has spec + m_wav."""
+    """NN view: NN1-NN10 from baseline (in prompt order). Each has spec + m_wav."""
     if not sample_id:
-        return (None,) * (4 + 10 * 2)
+        return (None,) * (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:
+    while len(out) < 20:
         out.append(None)
-    return tuple(out[: 4 + 20])
+    return tuple(out[:20])
 
 
 def build_results_view(sample_id: str | None):
@@ -63,24 +61,16 @@ with gr.Blocks(title="NearestNeighbor Audio Demo", css=".gradio-container { max-
 **Audio labels**: **BG** = background noise | **FG** = generated foreground | **Mix** = BG + FG
 """)
 
-    # ---- NN View: Baseline + 10 NN ----
+    # ---- NN View: NN1-NN10 from baseline (in prompt order) ----
     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**")
+        gr.Markdown("### Nearest Neighbor: Baseline outputs (top 10 prompts)")
         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
+        nn_outputs = nn_items
 
     # ---- Results View: 3 prompts × 4 methods ----
     with gr.Column(visible=False) as res_col:

dataset_loader.py

@@ -84,57 +84,37 @@ def _collect_block(file_list: list, folder_prefix: str) -> dict:
 
 def get_nn_demo_paths(bid: str, top_k: int = 10) -> dict:
     """
-    For NN view: [Baseline] [NN1] [NN2] ... [NN10].
-    Returns {baseline: {spec, bg, fg, m}, nn_list: [{fg_wav, spec, bg_wav, prompt, similarity}, ...]}.
+    For NN view: NN1-NN10 from baseline (generated_baseline_01, 02, ..., 10) in prompt order.
+    Returns {nn_list: [{spec, bg_wav, fg_wav, m_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 {"baseline": {"spec": None, "bg_wav": None, "fg_wav": None, "m_wav": None}, "nn_list": []}
+        return {"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[: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 = m_path = spec_path = None
-        for f in files:
-            parts = f.replace(inner + "/", "").split("/")
-            if len(parts) >= 2 and parts[0].startswith(gen_prefix):
-                name = parts[-1]
-                if name.endswith("_fg.wav"):
-                    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,
-        })
-
-    return {"baseline": baseline_block, "nn_list": nn_list}
+        bl_prefix = f"generated_baseline_{i+1:02d}_"
+        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(bl_prefix):
+                full_prefix = baseline_inner + "/" + parts[0]
+                block = _collect_block(baseline_files, full_prefix)
+                break
+        block["prompt"] = prompt
+        block["similarity"] = sim
+        nn_list.append(block)
+
+    return {"nn_list": nn_list}
 
 
 def get_noise_demo_paths(bid: str) -> dict: