Speed up beam, stream per-image results with progress, single upload

- pipeline: hoist per-ply digit marginals out of the per-hypothesis loop
(_ply_marginals), ~2.1x faster beam with identical output; thread a
progress_cb through run_pipeline and beam_decode (phase + per-ply)
- app: single multi-file uploader replaces the 5 fixed image slots; convert
is now a generator that streams each image's result as it finishes and
builds the unified zip only at the end; gr.Progress bar shows per-image
phase/ply; default beam lowered 8192 -> 2048 (a full game was ~2min at
8192, ~23s at 2048, and diagram evidence prunes well)

app.py

@@ -56,7 +56,9 @@ from togyz.pipeline import load_classifier, run_pipeline
 MAX_IMAGES = 5
 MODEL_DIR = Path(__file__).parent / "models"
 BEAM_CHOICES = [str(2**k) for k in range(10, 31)]  # 1024 ... 2^30
-DEFAULT_BEAM = "8192"
+# A full ~160-ply game costs roughly (width/1024) x 12s of beam search, so the
+# default stays modest; the board-diagram evidence prunes well at this width.
+DEFAULT_BEAM = "2048"
 
 # accepted spellings of a game result -> pipeline result code
 RESULT_MAP = {
@@ -170,10 +172,11 @@ def _pgn_tags(meta: dict, game: dict) -> str:
 
 
 def _process_one(image_path, game: dict, meta: dict, beam_width: int,
-                 base_name: str, out_dir: Path):
+                 base_name: str, out_dir: Path, progress_cb=None):
     """Run the pipeline on one image; return (row, gallery_item, files, warnings)."""
     out = run_pipeline(image_path, _MOVES, _DIAGRAM,
-                       result=game["result"], beam_width=beam_width)
+                       result=game["result"], beam_width=beam_width,
+                       progress_cb=progress_cb)
 
     stop = out["stopped"]
     stop_txt = stop.get("reason", "")
@@ -194,11 +197,32 @@ def _process_one(image_path, game: dict, meta: dict, beam_width: int,
     return row, (out["annotated_image"], caption), files, warnings
 
 
-def convert(meta_text, games_text, beam_choice, img1, img2, img3, img4, img5):
-    """One batch: up to 5 images sharing tournament metadata."""
-    images = [img for img in (img1, img2, img3, img4, img5) if img]
+def _as_paths(files) -> list[str]:
+    """Normalize the multi-file uploader value into a list of file paths."""
+    if not files:
+        return []
+    if isinstance(files, (str, os.PathLike)):
+        files = [files]
+    paths = []
+    for f in files:
+        # gr.File yields str paths (type="filepath") or objects with .name
+        paths.append(f if isinstance(f, str) else getattr(f, "name", str(f)))
+    return paths
+
+
+def convert(meta_text, games_text, beam_choice, files, progress=gr.Progress()):
+    """One batch: up to 5 images sharing tournament metadata.
+
+    A generator: it yields (table, gallery, zip, warnings) after each image so
+    results stream in one by one; the zip download is assembled only at the
+    end and contains every game's PGNs appended together.
+    """
+    images = _as_paths(files)
     if not images:
         raise gr.Error("Please upload at least one scoresheet image.")
+    if len(images) > MAX_IMAGES:
+        raise gr.Error(f"This demo handles at most {MAX_IMAGES} images per run "
+                       f"(got {len(images)}).")
 
     # parse everything up front - all input errors surface before any heavy work
     meta = _parse_meta(meta_text)
@@ -210,39 +234,54 @@ def convert(meta_text, games_text, beam_choice, img1, img2, img3, img4, img5):
 
     out_dir = Path(tempfile.mkdtemp(prefix="togyz_"))
     round_slug = _safe_slug(meta["round"])
+    n = len(images)
     rows, gallery, all_files, all_warnings = [], [], [], []
 
-    for i, (img, game) in enumerate(zip(images, games), start=1):
-        prefix = f"round{round_slug}_game{i}" if round_slug else f"game{i}"
+    def warn_md():
+        return "\n".join(f"⚠ {w}" for w in dict.fromkeys(all_warnings))
+
+    for i, (img, game) in enumerate(zip(images, games)):
+        name = Path(img).name
+        prefix = f"round{round_slug}_game{i + 1}" if round_slug else f"game{i + 1}"
+
+        def cb(frac, desc, i=i, name=name):
+            # blend per-image progress into an overall 0..1 bar
+            progress((i + frac) / n, desc=f"Image {i + 1}/{n} ({name}): {desc}")
+
+        cb(0.0, "starting")
         try:
             row, gal, files, warns = _process_one(
-                img, game, meta, beam_width, prefix, out_dir
+                img, game, meta, beam_width, prefix, out_dir, progress_cb=cb
             )
         except Exception as exc:  # one bad image must not kill the batch
-            rows.append([prefix, 0, f"error: {exc}", ""])
+            rows.append([f"{prefix} ({name})", 0, f"error: {exc}", ""])
+            yield rows[:], gallery[:], None, warn_md()
             continue
+        row[0] = f"{row[0]} ({name})"
         rows.append(row)
         gallery.append(gal)
         all_files.extend(files)
         all_warnings.extend(warns)
+        # stream this image's result immediately; zip is built only at the end
+        yield rows[:], gallery[:], None, warn_md()
 
-    warnings_md = "\n".join(f"⚠ {w}" for w in dict.fromkeys(all_warnings))
     if not all_files:
         # every image errored - still return the table so the user sees why
-        return rows, gallery, None, warnings_md
+        yield rows[:], gallery[:], None, warn_md()
+        return
 
     zip_path = out_dir / (f"round{round_slug}_pgns.zip" if round_slug else "pgns.zip")
     with zipfile.ZipFile(zip_path, "w") as zf:
         for f in all_files:
             zf.write(f, arcname=Path(f).name)
 
-    return rows, gallery, str(zip_path), warnings_md
+    yield rows[:], gallery[:], str(zip_path), warn_md()
 
 
-def _busy_wrapper(*args):
+def _busy_wrapper(meta_text, games_text, beam_choice, files, progress=gr.Progress()):
     """Turn infrastructure overload into a friendly message instead of a 500."""
     try:
-        return convert(*args)
+        yield from convert(meta_text, games_text, beam_choice, files, progress)
     except gr.Error:
         raise
     except Exception as exc:  # noqa: BLE001 - surface anything else gracefully
@@ -255,9 +294,10 @@ def _busy_wrapper(*args):
 with gr.Blocks(title="Togyzkumalak Scoresheet Reader") as demo:
     gr.Markdown(
         "# Togyzkumalak Scoresheet Reader\n"
-        "Upload up to **5** scoresheet photos of one round as a single batch, "
+        "Upload up to **5** scoresheet photos of one round in a single batch, "
         "describe the round and the games in the two text fields, then "
-        "**Convert** to download the PGNs.\n\n"
+        "**Convert**. Results stream in image by image with a live progress "
+        "bar; the combined PGN download appears once every image is done.\n\n"
         "Outputs per game: `beam` (best legal reconstruction), `raw` (pure OCR), "
         "`legal` (strict replay). Free demo — a first run may wake the Space, and "
         "images are processed one at a time."
@@ -280,16 +320,20 @@ with gr.Blocks(title="Togyzkumalak Scoresheet Reader") as demo:
               "thorough; very large values are trimmed to fit memory)",
     )
 
-    image_slots = [
-        gr.Image(label=f"Game {i + 1}", type="filepath", height=150)
-        for i in range(MAX_IMAGES)
-    ]
+    image_files = gr.File(
+        label=f"Scoresheet photos (up to {MAX_IMAGES}, in the same order as the "
+              "game lines above)",
+        file_count="multiple",
+        file_types=["image"],
+        type="filepath",
+    )
 
     convert_btn = gr.Button("Convert", variant="primary")
 
     results_table = gr.Dataframe(
         headers=["game", "legal plies", "stopped", "beam PGN"],
-        label="Results", wrap=True, interactive=False,
+        label="Results (stream in as each image finishes)",
+        wrap=True, interactive=False,
     )
     warnings_md = gr.Markdown()
     gallery = gr.Gallery(label="Annotated reconstruction", columns=2, height="auto")
@@ -297,7 +341,7 @@ with gr.Blocks(title="Togyzkumalak Scoresheet Reader") as demo:
 
     convert_btn.click(
         _busy_wrapper,
-        inputs=[meta_box, games_box, beam_dd, *image_slots],
+        inputs=[meta_box, games_box, beam_dd, image_files],
         outputs=[results_table, gallery, zip_out, warnings_md],
     )
 

togyz/pipeline.py

@@ -117,14 +117,12 @@ def classify_cells(images, clf: Classifier, temperature=1.0, tta=True, batch_siz
 # --------------------------------------------------------------------------- #
 # Legal-move scoring, kazan evidence, beam search  (numpy port)
 # --------------------------------------------------------------------------- #
-def _legal_scores(probs, legal_moves, classes, class_idx):
-    """Redistribute the 162-class distribution over the current legal moves.
+def _ply_marginals(probs, classes):
+    """Source/landing digit marginals and the x-mark mass for one ply.
 
-    A legal move earns credit from (a) its exact class string and (b) a
-    factorized term P(source digit) * P(landing digit), times agreement with
-    the written x-mark probability. Scores stay UNNORMALIZED for the beam so a
-    hypothesis whose legal set explains the observation poorly accumulates a
-    genuinely low joint likelihood (measured: 24/39 vs 29/39 when normalized).
+    Depends only on the ply's class distribution, so the beam computes it
+    once per ply and shares it across every hypothesis (this used to be
+    recomputed per hypothesis - the dominant cost at large beam widths).
     """
     source_marginal = [0.0] * 9
     landing_marginal = [0.0] * 9
@@ -136,7 +134,19 @@ def _legal_scores(probs, legal_moves, classes, class_idx):
         landing_marginal[int(c[1]) - 1] += p
         if c.endswith("x"):
             x_marginal += p
+    return source_marginal, landing_marginal, x_marginal
+
+
+def _legal_scores(probs, legal_moves, class_idx, marginals):
+    """Redistribute the 162-class distribution over the current legal moves.
 
+    A legal move earns credit from (a) its exact class string and (b) a
+    factorized term P(source digit) * P(landing digit), times agreement with
+    the written x-mark probability. Scores stay UNNORMALIZED for the beam so a
+    hypothesis whose legal set explains the observation poorly accumulates a
+    genuinely low joint likelihood (measured: 24/39 vs 29/39 when normalized).
+    """
+    source_marginal, landing_marginal, x_marginal = marginals
     scored = []
     for move in legal_moves:
         exact = float(probs[class_idx[move.notation]])
@@ -210,13 +220,14 @@ def safe_beam_width(beam_width: int, per_ply: int = 9) -> int:
 
 
 def beam_decode(ply_probs, classes, beam_width=1024, per_ply=9, result=None,
-                checkpoints=None):
+                checkpoints=None, progress_cb=None):
     """Longest fully-legal move sequence with the highest joint probability,
     exploring all `per_ply` legal continuations per ply; the final pool is
     re-ranked by diagram-checkpoint and known-result consistency.
 
     `beam_width` is capped by `safe_beam_width` so user-requested widths up
     to 2^30 degrade into a narrower beam instead of exhausting RAM.
+    `progress_cb(ply_index, total_plies)` is called once per ply if given.
 
     Returns (annotated_moves, log_prob, per_ply_share_list).
     """
@@ -225,11 +236,16 @@ def beam_decode(ply_probs, classes, beam_width=1024, per_ply=9, result=None,
     class_idx = {c: i for i, c in enumerate(classes)}
     active = [(Game(), [], [], 0.0)]  # (game, annotated_moves, chosen_shares, logp)
     finished = []
+    total_plies = len(ply_probs)
 
-    for probs in ply_probs:
+    for ply_index, probs in enumerate(ply_probs):
+        if progress_cb:
+            progress_cb(ply_index, total_plies)
+        # marginals are identical across all hypotheses this ply - compute once
+        marginals = _ply_marginals(probs, classes)
         expanded = []
         for game, moves, chosen, logp in active:
-            candidates = _legal_scores(probs, game.legal_moves(), classes, class_idx)
+            candidates = _legal_scores(probs, game.legal_moves(), class_idx, marginals)
             for raw, share, move in candidates[:per_ply]:
                 nxt = game.copy()
                 played = nxt.play(move.action)
@@ -263,17 +279,24 @@ def format_pgn(plies: list[str]) -> str:
 # --------------------------------------------------------------------------- #
 def run_pipeline(image, moves_clf: Classifier, diagram_clf: Classifier | None = None,
                  result=None, topk=5, beam_width=8192, per_ply=9,
-                 temperature=1.0, tta=True, save_cells_dir=None):
+                 temperature=1.0, tta=True, save_cells_dir=None, progress_cb=None):
     """Read one scoresheet image into game records.
 
     `image` is a path or a PIL.Image. Returns a dict with the reconstructed
     PGNs and diagnostics; no files are written unless `save_cells_dir` is set
     (the CLI uses it, the Gradio app does not).
 
+    `progress_cb(frac, desc)` - if given - is called with a 0..1 fraction and
+    a human-readable phase description as the read progresses.
+
     Keys: beam_pgn, raw_pgn, legal_pgn, game_json (dict), stopped,
     plies_scanned, legal_plies, beam_plies, checkpoint_report, warnings,
     annotated_image (PIL), median_cell_height, low_resolution (bool).
     """
+    def report(frac, desc):
+        if progress_cb:
+            progress_cb(frac, desc)
+
     classes = moves_clf.classes
     empty_idx = classes.index(EMPTY)
     warnings: list[str] = []
@@ -284,10 +307,12 @@ def run_pipeline(image, moves_clf: Classifier, diagram_clf: Classifier | None =
             "inside the memory budget"
         )
 
+    report(0.0, "detecting tables and cells")
     cells, sheet, gridlines = extract_cells(image, with_gridlines=True)
     cells.sort(key=lambda c: (c.move_no, c.side != "W"))  # game order: 1W 1B 2W ...
     median_h = sorted(c.bbox[3] for c in cells)[len(cells) // 2]
 
+    report(0.05, f"classifying {len(cells)} move cells")
     cleaned = [clean_cell(c.image) for c in cells]  # (image, ink_ratio) pairs
     all_probs = classify_cells(
         [img for img, _ in cleaned], moves_clf, temperature=temperature, tta=tta
@@ -355,6 +380,7 @@ def run_pipeline(image, moves_clf: Classifier, diagram_clf: Classifier | None =
     checkpoints, checkpoint_report = {}, []
     diagram_cells, diagram_labels = [], {}
     if diagram_clf is not None:
+        report(0.30, "reading board diagrams")
         dg_classes = diagram_clf.classes
         value_names = [str(v) for v in range(82)]
         kazan_allowed = set(str(v) for v in range(10, 82))
@@ -404,10 +430,15 @@ def run_pipeline(image, moves_clf: Classifier, diagram_clf: Classifier | None =
                     img.save(cells_dir /
                              f"{cell.kind}_{cell.move_no:02d}_{cell.side}{suffix}.png")
 
+    def beam_progress(ply_index, total):
+        report(0.45 + 0.53 * (ply_index / max(total, 1)),
+               f"reconstructing game (ply {ply_index + 1}/{total})")
+
     beam_moves, beam_logp, beam_shares = beam_decode(
         ply_prob_arrays, classes, effective_width, per_ply,
-        result=result, checkpoints=checkpoints,
+        result=result, checkpoints=checkpoints, progress_cb=beam_progress,
     )
+    report(1.0, "done")
     beam_detail = [
         {"move": p["move"], "side": p["side"], "chosen": m,
          "prob": round(q, 4), "agrees_with_raw": m.rstrip("+") == p["raw"]}