Results: add Results (CLAP) and Results (Dasheng) views; dataset_loader supports dasheng results

app.py

@@ -36,7 +36,7 @@ def build_nn_view(sample_id: str | None, use_dasheng: bool = False):
     return tuple(out[:50])
 
 
-def build_results_view(sample_id: str | None):
+def build_results_view(sample_id: str | None, use_dasheng: bool = False):
     """
     Results view: 3 blocks. Per block:
     - Row1: Gaussian | Youtube spec + their BG/FG/Mix
@@ -44,7 +44,8 @@ def build_results_view(sample_id: str | None):
     """
     if not sample_id:
         return (None,) * (3 * (1 + 4 * 4))
-    data = get_results_demo_paths(sample_id)
+    root_prefix = DASHENG_PREFIX if use_dasheng else None
+    data = get_results_demo_paths(sample_id, root_prefix=root_prefix)
     out = []
     for i in range(1, 4):
         block = data.get(f"block{i}", {})
@@ -66,7 +67,12 @@ with gr.Blocks(title="NearestNeighbor Audio Demo", css=".gradio-container { max-
     gr.Markdown("Data from [AE-W/batch_outputs](https://huggingface.co/datasets/AE-W/batch_outputs)")
 
     view_radio = gr.Radio(
-        choices=["Nearest Neighbor (CLAP)", "Nearest Neighbor (Dasheng)", "Results"],
+        choices=[
+            "Nearest Neighbor (CLAP)",
+            "Nearest Neighbor (Dasheng)",
+            "Results (CLAP)",
+            "Results (Dasheng)",
+        ],
         value="Nearest Neighbor (CLAP)",
         label="View",
     )
@@ -97,7 +103,7 @@ with gr.Blocks(title="NearestNeighbor Audio Demo", css=".gradio-container { max-
 
     # ---- Results View: 3 prompts, each with 2 rows (Gaussian|Youtube, Ours|NN baseline) ----
     with gr.Column(visible=False) as res_col:
-        gr.Markdown("### Results: 3 baselines + Ours (top 3 prompts)")
+        res_section_title = gr.Markdown("### Results (CLAP): 3 baselines + Ours (top 3 prompts)")
         res_outputs = []
         for i in range(1, 4):
             with gr.Group():
@@ -131,8 +137,8 @@ with gr.Blocks(title="NearestNeighbor Audio Demo", css=".gradio-container { max-
                         res_outputs.append(gr.Audio(label="Mix"))
 
     def on_change(sid, view):
-        use_dasheng = view == "Nearest Neighbor (Dasheng)"
-        # Dasheng view: only show IDs that exist in batch_outputs_dasheng (no fold*)
+        use_dasheng = view in ("Nearest Neighbor (Dasheng)", "Results (Dasheng)")
+        # Dasheng views: only show IDs that exist in batch_outputs_dasheng (no fold*)
         if use_dasheng:
             choices = DASHENG_SAMPLE_IDS
             if sid not in DASHENG_SAMPLE_IDS and DASHENG_SAMPLE_IDS:
@@ -141,17 +147,19 @@ with gr.Blocks(title="NearestNeighbor Audio Demo", css=".gradio-container { max-
             choices = SAMPLE_IDS
             if sid not in SAMPLE_IDS and SAMPLE_IDS:
                 sid = SAMPLE_IDS[0]
-        nn_vals = build_nn_view(sid, use_dasheng=use_dasheng)
-        res_vals = build_results_view(sid)
         is_nn = view in ("Nearest Neighbor (CLAP)", "Nearest Neighbor (Dasheng)")
-        nn_title = "### Nearest Neighbor (Dasheng): Baseline outputs (top 10 prompts)" if use_dasheng else "### Nearest Neighbor (CLAP): Baseline outputs (top 10 prompts)"
+        is_res = view in ("Results (CLAP)", "Results (Dasheng)")
+        nn_vals = build_nn_view(sid, use_dasheng=(view == "Nearest Neighbor (Dasheng)"))
+        res_vals = build_results_view(sid, use_dasheng=(view == "Results (Dasheng)"))
+        nn_title = "### Nearest Neighbor (Dasheng): Baseline outputs (top 10 prompts)" if view == "Nearest Neighbor (Dasheng)" else "### Nearest Neighbor (CLAP): Baseline outputs (top 10 prompts)"
+        res_title = "### Results (Dasheng): 3 baselines + Ours (top 3 prompts)" if view == "Results (Dasheng)" else "### Results (CLAP): 3 baselines + Ours (top 3 prompts)"
         dd_update = gr.update(choices=choices, value=sid)
         return (
-            [gr.update(value=nn_title)] + list(nn_vals) + list(res_vals) +
-            [gr.update(visible=is_nn), gr.update(visible=not is_nn), dd_update]
+            [gr.update(value=nn_title)] + list(nn_vals) + [gr.update(value=res_title)] + list(res_vals) +
+            [gr.update(visible=is_nn), gr.update(visible=is_res), dd_update]
         )
 
-    all_outputs = [nn_section_title] + nn_outputs + res_outputs + [nn_col, res_col, noise_dd]
+    all_outputs = [nn_section_title] + nn_outputs + [res_section_title] + res_outputs + [nn_col, res_col, noise_dd]
 
     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)

dataset_loader.py

@@ -191,12 +191,14 @@ def get_noise_demo_paths(bid: str) -> dict:
     return result
 
 
-def get_results_demo_paths(bid: str) -> dict:
+def get_results_demo_paths(bid: str, root_prefix: Optional[str] = None) -> dict:
     """
     For Results view: 3 blocks (prompts 1-3), each with 4 columns:
     Baseline (original), Gaussian, Youtube-noise, Ours.
+    root_prefix: None = batch_outputs (CLAP), DASHENG_PREFIX = batch_outputs_dasheng.
     """
-    inner = f"{ROOT_PREFIX}{bid}/{bid}"
+    prefix = root_prefix if root_prefix is not None else ROOT_PREFIX
+    inner = f"{prefix}{bid}/{bid}"
     files = _find_files(inner)
     baseline_inner = f"{inner}/baseline"
     gaussian_inner = f"{inner}/gaussian_baseline"
@@ -211,6 +213,8 @@ def get_results_demo_paths(bid: str) -> dict:
     if not prompts:
         prompts = []
 
+    use_dasheng = root_prefix == DASHENG_PREFIX
+
     def get_baseline_folders(bl_inner, bl_files):
         seen = set()
         folders = []
@@ -223,6 +227,17 @@ def get_results_demo_paths(bid: str) -> dict:
         return folders
 
     def get_youtube_folders():
+        if use_dasheng:
+            # Dasheng: subdirs are prompt names (underscores)
+            seen = set()
+            folders = []
+            for f in youtube_files:
+                parts = f.replace(youtube_inner + "/", "").split("/")
+                if parts 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
         seen = set()
         folders = []
         for f in youtube_files:
@@ -251,13 +266,28 @@ def get_results_demo_paths(bid: str) -> dict:
         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)
+        if use_dasheng:
+            # Dasheng: match by prompt -> folder name (spaces to underscores)
+            folder_name = prompt_text.replace(" ", "_") if prompt_text else ""
+            for fn, fp in youtube_folders:
+                if fn == folder_name:
+                    bl_youtube = _collect_block(youtube_files, fp)
+                    break
+        else:
+            for fn, fp in youtube_folders:
+                if fn.startswith(rel_prefix):
+                    bl_youtube = _collect_block(youtube_files, fp)
+                    break
+
+        if use_dasheng:
+            # Dasheng: "ours" = prompt-named folder under inner
+            folder_name = prompt_text.replace(" ", "_") if prompt_text else ""
+            ours_prefix = f"{inner}/{folder_name}"
+            nn_files = [f for f in files if f.startswith(ours_prefix + "/")]
+            ours_block = _collect_block(nn_files, ours_prefix)
+        else:
+            nn_files = [f for f in files if f.replace(inner + "/", "").startswith(rel_prefix)]
+            ours_block = _collect_block(nn_files, inner + "/" + rel_prefix)
 
         result[f"block{i}"] = {
             "prompt": prompt_text,