Merge pull request #119 from weixlu/serilization

avoid decode/encode overhead in benchmark dataset loading

Author: liebharc

validation/ned_benchmark.py

@@ -41,7 +41,7 @@
 import statistics
 import sys
 import traceback
-from collections.abc import Callable, Iterable, Iterator
+from collections.abc import Callable, Iterator
 from pathlib import Path
 
 from homr.transformer.vocabulary import EncodedSymbol
@@ -240,7 +240,7 @@ def update_ned_scores(
 
 
 def run_benchmark(
-    samples: Iterable[tuple[str, str, Path | None]],
+    samples: list[tuple[str, str, Path]],
     tool: Callable[[str, Path | None], str],
     limit: int | None = None,
     verbose: bool = False,
@@ -253,7 +253,7 @@ def run_benchmark(
     """
     Wire data source, tool, and check together.
 
-    samples:      iterable of (sample_id, kern_text, image_path) from any data source
+    samples:      list of (sample_id, kern_text, image_path) from any data source
     tool:         (kern_text, image_path) -> output_text (MusicXML or **kern)
     kern_parser:  "native" (default) or "music21" - kern ground-truth parser
     xml_parser:   "native" (default), "music21", "musicdiff", or "musicdiff_detailed"
@@ -280,15 +280,9 @@ def run_benchmark(
     results: list[NedResult] = []
     failures: list[tuple[str, str]] = []
 
-    # Collect from iterator (applying limit and skips) so we can optionally batch.
-    all_from_iter: list[tuple[str, str, Path | None]] = []
-    for count, triple in enumerate(samples):
-        if limit is not None and count >= limit:
-            break
-        all_from_iter.append(triple)
-
-    skipped = sum(1 for s, _, _ in all_from_iter if s in skip_ids)
-    active = [(s, k, i) for s, k, i in all_from_iter if s not in skip_ids]
+    samples = samples[:limit]
+    active = [s for s in samples if s[0] not in skip_ids]
+    skipped = len(samples) - len(active)
 
     if xml_parser in ("musicdiff", "musicdiff_detailed"):
         _musicdiff_register_once()

validation/polish-scores.py

@@ -1,3 +1,5 @@
+# flake8: noqa: S101
+
 """
 btrkeks/polish-scores benchmark: measures OMR quality on the Polish Scores dataset.
 
@@ -8,7 +10,6 @@
 
 import argparse
 import tempfile
-from collections.abc import Iterator
 from pathlib import Path
 
 from validation.ned_benchmark import run_benchmark, update_ned_scores
@@ -28,23 +29,22 @@ def _add_kern_header(kern: str) -> str:
     return kern
 
 
-def iter_polish_scores(
-    image_dir: Path | None = None,
-) -> Iterator[tuple[str, str, Path | None]]:
+def get_polish_scores_samples(
+    image_dir: Path,
+) -> list[tuple[str, str, Path]]:
     """Yield (sample_id, kern_text, image_path) for each sample in btrkeks/polish-scores."""
-    from datasets import load_dataset  # type: ignore  # noqa: PLC0415
+    from datasets import Image, load_dataset  # type: ignore  # noqa: PLC0415
 
-    ds = load_dataset("btrkeks/polish-scores")
-    for i, sample in enumerate(ds["train"]):
+    ds = load_dataset("btrkeks/polish-scores")["train"]
+    ds = ds.cast_column("image", Image(decode=False))  # Don't decode on image colume.
+    result = []
+    for i, sample in enumerate(ds):
+        assert sample["transcription_kern"] and sample["image"]["bytes"]
         kern = _add_kern_header(sample["transcription_kern"])
-        image_path: Path | None = None
-        if image_dir is not None:
-            img = sample.get("image")
-            if img is not None and hasattr(img, "save"):
-                candidate = image_dir / f"{i}.png"
-                img.save(candidate)
-                image_path = candidate
-        yield str(i), kern, image_path
+        image_path = image_dir / f"{i}.png"
+        image_path.write_bytes(sample["image"]["bytes"])
+        result.append((str(i), kern, image_path))
+    return result
 
 
 def main() -> None:
@@ -114,7 +114,7 @@ def main() -> None:
     tool = TOOLS[args.tool]
     with tempfile.TemporaryDirectory() as image_dir:
         run_benchmark(
-            iter_polish_scores(image_dir=Path(image_dir)),
+            get_polish_scores_samples(Path(image_dir)),
             tool,
             limit=args.limit,
             verbose=args.verbose,

validation/smb.py

@@ -1,3 +1,5 @@
+# flake8: noqa: S101
+
 """
 PRAIG/SMB benchmark: measures OMR quality on the SMB dataset.
 
@@ -12,50 +14,41 @@
 
 import argparse
 import tempfile
-from collections.abc import Iterator
 from pathlib import Path
 
 from validation.ned_benchmark import run_benchmark, update_ned_scores
 from validation.tools import TOOLS
 
 
-def iter_smb(split: str, image_dir: Path | None = None) -> Iterator[tuple[str, str, Path | None]]:
+def get_smb_samples(image_dir: Path) -> list[tuple[str, str, Path]]:
     """Yield (sample_id, kern_text, image_path) for each page in PRAIG/SMB.
 
     Passes the full page image to allow proper end-to-end OMR evaluation.
     Ground truth priority: page.kern > top-level kern > concatenated region kern.
     image_path is only set when image_dir is provided and the page has an image.
     """
-    from datasets import load_dataset  # type: ignore  # noqa: PLC0415
-
-    ds = load_dataset("PRAIG/SMB")
-    for i, sample in enumerate(ds[split]):
-        kern = ""
-        page = sample.get("page", {})
-        if isinstance(page, dict):
-            kern = page.get("kern") or ""
-        if not kern:
-            kern = sample.get("kern") or ""
-        if not kern:
-            parts = [r.get("kern", "") for r in sample.get("regions", []) if r.get("kern")]
-            kern = "\n".join(parts)
-        if not kern:
-            continue
-
-        image_path: Path | None = None
-        if image_dir is not None:
-            img = sample.get("image")
-            if img is not None and hasattr(img, "save"):
-                candidate = image_dir / f"{i}.png"
-                img.save(candidate)
-                image_path = candidate
-
-        yield str(i), kern, image_path
+    from datasets import Image, load_dataset  # type: ignore  # noqa: PLC0415
+
+    ds = load_dataset("PRAIG/SMB")["test"]
+    ds = ds.cast_column("image", Image(decode=False))  # Don't decode on image colume.
+    result = []
+    for i, sample in enumerate(ds):
+        assert "kern" not in sample and isinstance(
+            sample["regions"], list
+        )  # not a global `kern` field, but a list of regions.
+        assert sample["image"]["bytes"]
+
+        parts = [r["kern"] for r in sample["regions"] if r["kern"]]
+        kern = "\n".join(parts)
+        image_path = image_dir / f"{i}.png"
+        image_path.write_bytes(sample["image"]["bytes"])
+
+        result.append((str(i), kern, image_path))
+    return result
 
 
 def main() -> None:
     parser = argparse.ArgumentParser(description="OMR-NED benchmark for PRAIG/SMB.")
-    parser.add_argument("--split", type=str, default="test", help="Dataset split (default: test).")
     parser.add_argument("--limit", type=int, default=None, help="Only process N samples.")
     parser.add_argument("--verbose", action="store_true", help="Print traceback on failure.")
     parser.add_argument("--output", type=str, default=None, help="Path to SQLite output file.")
@@ -121,7 +114,7 @@ def main() -> None:
     tool = TOOLS[args.tool]
     with tempfile.TemporaryDirectory() as image_dir:
         run_benchmark(
-            iter_smb(args.split, image_dir=Path(image_dir)),
+            get_smb_samples(image_dir=Path(image_dir)),
             tool,
             limit=args.limit,
             verbose=args.verbose,