Read full board diagrams, dynamic beam width, CSV metadata batch UI

- sheet.py: extract_diagram_cells reads the 2x9 pit grid + kazan boxes per
summary strip, table-guided strip location (splits merged components,
positional checkpoint numbering), median-area filter in _find_table_quads,
gridline export, richer render_overlay (kazans blue, pits orange, gridlines)
- classes/synth/dataset: unified 85-class diagram task (0-81, x, -, empty)
with overflow/neighbor-bleed synthesis; --task diagram everywhere
- pipeline: _diagram_logp scores kazans + pits per checkpoint, beam width
user-selectable with memory guard (safe_beam_width) + heapq.nlargest,
checkpoint_report/warnings in output
- app.py: batch of 5 images with shared Tournament,Location,Date,Round CSV
field, per-game White,Black,Result,WhiteTime,BlackTime lines, beam-width
dropdown 2^10..2^30, PGN tags in all outputs
- Colab notebook rewritten: clone repo, fetch glyphs at runtime, train both
tasks, export ONNX, download artifacts.zip

README.md

@@ -19,16 +19,20 @@ records, and a **Gradio demo** (`app.py`) packages it for a HuggingFace Space.
 
 ## Live demo (HuggingFace Space)
 
-`app.py` is a Gradio app: upload up to **5** scoresheet photos, optionally tag
-each game's known result, and download the reconstructed PGNs (beam / raw /
-legal) as a zip. Inference runs on the exported **ONNX** models (torch-free),
-so the Space needs only the light serving deps in `requirements.txt`.
+`app.py` is a Gradio app: upload up to **5** scoresheet photos of one round as
+a single batch, describe the tournament (`Tournament,Location,Date,Round`) and
+each game (`WhiteName,BlackName,Result,WhiteTime,BlackTime`, one CSV line per
+image) in two text fields, pick a beam width (1024 … 2^30; oversized requests
+are trimmed to fit memory), and download the reconstructed PGNs (beam / raw /
+legal, with PGN tags) as a zip. Inference runs on the exported **ONNX** models
+(torch-free), so the Space needs only the light serving deps in
+`requirements.txt`.
 
 ```bash
 python scripts/export_onnx.py     # checkpoints/*.pt -> single-file .onnx (+ .classes.json)
 cp checkpoints/best.onnx checkpoints/best.classes.json models/
-cp checkpoints/kazan/best.onnx models/kazan.onnx
-cp checkpoints/kazan/best.classes.json models/kazan.classes.json
+cp checkpoints/diagram/best.onnx models/diagram.onnx
+cp checkpoints/diagram/best.classes.json models/diagram.classes.json
 python app.py                     # serve locally at http://127.0.0.1:7860
 ```
 
@@ -85,12 +89,15 @@ A real training run needs a GPU — see Colab below. Defaults
 Open `notebooks/colab_train.ipynb` in Colab, or manually:
 
 ```bash
-git clone <this-repo> && cd 9OCR        # or upload a zip of the project
-pip install -r requirements.txt
+git clone https://github.com/ansarzeinulla/9OCR.git && cd 9OCR
+pip install -r requirements-train.txt
 python scripts/get_glyphs.py --emnist   # ingredients/ is gitignored, so also
                                         # rebuild the EMNIST pool from torchvision
-python train.py --epochs 30
-# download checkpoints/best.pt when done
+python train.py --task moves --epochs 30
+python train.py --task diagram --epochs 30
+python scripts/export_onnx.py
+# download checkpoints/best.onnx (+.classes.json) and
+# checkpoints/diagram/best.onnx (+.classes.json) when done
 ```
 
 `predict.py`/`eval.py` run anywhere (CPU is fine) with the downloaded
@@ -117,11 +124,15 @@ Pass `--result` (1-0 / 0-1 / draw, from the sheet footer) when known: the
 beam's final pool is re-ranked to prefer reconstructions whose end state is
 consistent with how the game actually ended.
 
-The summary strips between the tables record both kazan counts every 10
-moves (Black's box above the strip, White's below, always two digits 00-81).
-A second classifier reads them (`python train.py --task kazan`, saved to
-`checkpoints/kazan/best.pt`) and the beam gains likelihood when a
-hypothesis' computed kazans match these written checkpoints.
+The summary strips between the tables hold a full board diagram every 10
+moves: both kazan counts (Black's box above the strip, White's below, two
+digits 10-81) **and** the 2×9 pit grid (upper row = Black, pits 9…1
+left-to-right; lower row = White, pits 1…9; each cell is `x` = tuzdyk,
+`-` = 0, or a 1-2 digit count). A second, unified classifier reads all of it
+(`python train.py --task diagram`, saved to `checkpoints/diagram/best.pt`;
+classes 0-81 + `x` + `-` + `empty`, filtered per context at inference), and
+the beam gains likelihood when a hypothesis' computed kazans and pit counts
+match these written checkpoints.
 
 Finds the 8 printed move tables, classifies every cell, and writes:
 

app.py

@@ -1,8 +1,9 @@
 """Gradio demo for the Togyzkumalak scoresheet reader (HuggingFace Space).
 
-Upload up to 5 scoresheet photos, optionally tag each with its known result,
-and download the reconstructed PGNs. Inference runs on the exported ONNX
-models (torch-free) via `togyz.pipeline.run_pipeline`.
+Upload up to 5 scoresheet photos of one round as a single batch, describe the
+tournament and the games in two CSV text fields, and download the
+reconstructed PGNs (with proper PGN tags). Inference runs on the exported
+ONNX models (torch-free) via `togyz.pipeline.run_pipeline`.
 
 This is a demo, not production: state is per-session, work is capped at 5
 images per run, and requests are serialized through Gradio's queue so a shared
@@ -26,6 +27,8 @@ except Exception as e:
     print(f"[app] Failed to apply monkey-patch: {e}", flush=True)
 # --------------------------------------
 
+import csv
+import io
 import os
 import tempfile
 import zipfile
@@ -52,16 +55,29 @@ from togyz.pipeline import load_classifier, run_pipeline
 
 MAX_IMAGES = 5
 MODEL_DIR = Path(__file__).parent / "models"
-RESULT_CHOICES = ["unknown", "1-0", "0-1", "draw"]
+BEAM_CHOICES = [str(2**k) for k in range(10, 31)]  # 1024 ... 2^30
+DEFAULT_BEAM = "8192"
+
+# accepted spellings of a game result -> pipeline result code
+RESULT_MAP = {
+    "": None, "*": None,
+    "1": "1-0", "1-0": "1-0",
+    "0": "0-1", "0-1": "0-1",
+    "5": "draw", "0.5-0.5": "draw", "1/2-1/2": "draw", "draw": "draw",
+}
+RESULT_TAGS = {"1-0": "1-0", "0-1": "0-1", "draw": "1/2-1/2", None: "*"}
 
 # Load the ONNX sessions once at import - warm for the whole process lifetime.
 _log(f"loading move model from {MODEL_DIR / 'best.onnx'} ...")
 _MOVES = load_classifier(MODEL_DIR / "best.onnx")
-_KAZAN = None
-_kazan_path = MODEL_DIR / "kazan.onnx"
-if _kazan_path.exists():
-    _log("loading kazan model ...")
-    _KAZAN = load_classifier(_kazan_path)
+_DIAGRAM = None
+_diagram_path = MODEL_DIR / "diagram.onnx"
+if _diagram_path.exists():
+    _log("loading diagram model ...")
+    _DIAGRAM = load_classifier(_diagram_path)
+else:
+    # the old kazan.onnx has incompatible classes - do not fall back to it
+    _log("no models/diagram.onnx - checkpoint evidence disabled")
 _log("models loaded")
 
 
@@ -70,10 +86,94 @@ def _safe_slug(text: str) -> str:
     return keep.strip("_")
 
 
-def _process_one(image_path, result_choice, base_name, out_dir: Path):
-    """Run the pipeline on one image; return (row, gallery_item, pgn_files)."""
-    result = None if result_choice in (None, "unknown") else result_choice
-    out = run_pipeline(image_path, _MOVES, _KAZAN, result=result)
+def _csv_fields(line: str) -> list[str]:
+    """One CSV line -> stripped fields (handles quoted commas)."""
+    rows = list(csv.reader(io.StringIO(line)))
+    return [f.strip() for f in rows[0]] if rows else []
+
+
+def _parse_meta(text: str) -> dict:
+    """Shared metadata line: Tournament,Location,Date,Round (all optional)."""
+    fields = _csv_fields((text or "").strip())
+    if len(fields) > 4:
+        raise gr.Error(
+            "Metadata must be one CSV line: Tournament,Location,Date,Round "
+            f"(got {len(fields)} fields). Quote fields that contain commas."
+        )
+    fields += [""] * (4 - len(fields))
+    meta = {"event": fields[0], "site": fields[1], "date": fields[2],
+            "round": fields[3]}
+    if meta["round"]:
+        try:
+            rnd = int(meta["round"])
+        except ValueError:
+            raise gr.Error(f"Round must be a number 1-20, got {meta['round']!r}.")
+        if not 1 <= rnd <= 20:
+            raise gr.Error(f"Round must be between 1 and 20, got {rnd}.")
+        meta["round"] = str(rnd)
+    return meta
+
+
+def _parse_games(text: str, n_images: int) -> list[dict]:
+    """Per-game lines: WhiteName,BlackName,Result,WhiteTime,BlackTime.
+
+    One line per uploaded image, in order. Fewer lines than images is fine
+    (missing games get empty metadata); more lines is an error.
+    """
+    lines = [ln for ln in (text or "").splitlines() if ln.strip()]
+    if len(lines) > n_images:
+        raise gr.Error(
+            f"{len(lines)} game lines for {n_images} image(s). "
+            "Provide at most one line per uploaded image, in order."
+        )
+    games = []
+    for lineno, line in enumerate(lines, start=1):
+        fields = _csv_fields(line)
+        if len(fields) > 5:
+            raise gr.Error(
+                f"Game line {lineno}: expected at most 5 CSV fields "
+                "(White,Black,Result,WhiteTime,BlackTime), got "
+                f"{len(fields)}. Quote fields that contain commas."
+            )
+        fields += [""] * (5 - len(fields))
+        raw_result = fields[2]
+        if raw_result not in RESULT_MAP:
+            raise gr.Error(
+                f"Game line {lineno}: unknown result {raw_result!r}. Accepted: "
+                "1 or 1-0 (White won), 0 or 0-1 (Black won), "
+                "5 / 0.5-0.5 / 1/2-1/2 (draw), or empty."
+            )
+        games.append({"white": fields[0], "black": fields[1],
+                      "result": RESULT_MAP[raw_result],
+                      "white_time": fields[3], "black_time": fields[4]})
+    games += [{"white": "", "black": "", "result": None,
+               "white_time": "", "black_time": ""}] * (n_images - len(games))
+    return games
+
+
+def _pgn_tags(meta: dict, game: dict) -> str:
+    """Standard PGN tag section from the shared + per-game metadata."""
+    tags = [
+        ("Event", meta["event"] or "?"),
+        ("Site", meta["site"] or "?"),
+        ("Date", meta["date"] or "?"),
+        ("Round", meta["round"] or "?"),
+        ("White", game["white"] or "?"),
+        ("Black", game["black"] or "?"),
+        ("Result", RESULT_TAGS[game["result"]]),
+    ]
+    if game["white_time"]:
+        tags.append(("WhiteClock", game["white_time"]))
+    if game["black_time"]:
+        tags.append(("BlackClock", game["black_time"]))
+    return "".join(f'[{k} "{v}"]\n' for k, v in tags) + "\n"
+
+
+def _process_one(image_path, game: dict, meta: dict, beam_width: int,
+                 base_name: str, out_dir: Path):
+    """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)
 
     stop = out["stopped"]
     stop_txt = stop.get("reason", "")
@@ -81,53 +181,62 @@ def _process_one(image_path, result_choice, base_name, out_dir: Path):
         stop_txt += f" ({stop['winner']})"
     note = " ⚠ low-res" if out["low_resolution"] else ""
 
+    tags = _pgn_tags(meta, game)
     files = []
     for kind in ("beam", "raw", "legal"):
         f = out_dir / f"{base_name}_{kind}.pgn"
-        f.write_text(out[f"{kind}_pgn"])
+        f.write_text(tags + out[f"{kind}_pgn"])
         files.append(str(f))
 
     row = [base_name, out["beam_plies"], stop_txt + note, out["beam_pgn"].strip()]
     caption = f"{base_name}: {out['beam_plies']} plies"
-    return row, (out["annotated_image"], caption), files
-
+    warnings = [f"{base_name}: {w}" for w in out["warnings"]]
+    return row, (out["annotated_image"], caption), files, warnings
 
-def convert(round_label, *slot_values):
-    """slot_values = [img1, res1, img2, res2, ...] for the 5 fixed slots."""
-    images = slot_values[0::2]
-    results = slot_values[1::2]
-    provided = [(img, res) for img, res in zip(images, results) if img]
 
-    if not provided:
+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]
+    if not images:
         raise gr.Error("Please upload at least one scoresheet image.")
-    if len(provided) > MAX_IMAGES:  # defensive; the UI only exposes 5 slots
-        raise gr.Error(f"This demo handles at most {MAX_IMAGES} images per run.")
+
+    # parse everything up front - all input errors surface before any heavy work
+    meta = _parse_meta(meta_text)
+    games = _parse_games(games_text, len(images))
+    try:
+        beam_width = int(beam_choice)
+    except (TypeError, ValueError):
+        beam_width = int(DEFAULT_BEAM)
 
     out_dir = Path(tempfile.mkdtemp(prefix="togyz_"))
-    round_slug = _safe_slug(round_label)
-    rows, gallery, all_files = [], [], []
+    round_slug = _safe_slug(meta["round"])
+    rows, gallery, all_files, all_warnings = [], [], [], []
 
-    for i, (img, res) in enumerate(provided, start=1):
-        prefix = f"{round_slug}_table{i}" if round_slug else f"table{i}"
+    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}"
         try:
-            row, gal, files = _process_one(img, res, prefix, out_dir)
+            row, gal, files, warns = _process_one(
+                img, game, meta, beam_width, prefix, out_dir
+            )
         except Exception as exc:  # one bad image must not kill the batch
             rows.append([prefix, 0, f"error: {exc}", ""])
             continue
         rows.append(row)
         gallery.append(gal)
         all_files.extend(files)
+        all_warnings.extend(warns)
 
+    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
+        return rows, gallery, None, warnings_md
 
-    zip_path = out_dir / (f"{round_slug}_pgns.zip" if round_slug else "pgns.zip")
+    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)
+    return rows, gallery, str(zip_path), warnings_md
 
 
 def _busy_wrapper(*args):
@@ -146,22 +255,35 @@ 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, optionally tag each game's known "
-        "result (improves accuracy), then **Convert** to download the PGNs.\n\n"
+        "Upload up to **5** scoresheet photos of one round as a single batch, "
+        "describe the round and the games in the two text fields, then "
+        "**Convert** to download the PGNs.\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."
     )
-    round_label = gr.Textbox(label="Round (optional)", placeholder="e.g. 3",
-                             scale=1, max_lines=1)
+    meta_box = gr.Textbox(
+        label="Tournament metadata (CSV): Tournament,Location,Date,Round — all optional",
+        placeholder="World Championship among boys, Astana city, 2026 7 July, 7",
+        max_lines=1,
+    )
+    games_box = gr.Textbox(
+        label="Games (one CSV line per image, in order): "
+              "WhiteName,BlackName,Result,WhiteTime,BlackTime — all optional; "
+              "Result: 1 / 1-0, 0 / 0-1, 5 / 0.5-0.5 / 1/2-1/2",
+        placeholder="Zhanabay Korkem, Kubzhasar Marhabat, 0-1, 0:07:55, 0:24:05",
+        lines=MAX_IMAGES,
+    )
+    beam_dd = gr.Dropdown(
+        BEAM_CHOICES, value=DEFAULT_BEAM,
+        label="Beam width (game hypotheses kept; larger = slower but more "
+              "thorough; very large values are trimmed to fit memory)",
+    )
 
-    slots = []
-    for i in range(MAX_IMAGES):
-        with gr.Row():
-            img = gr.Image(label=f"Table {i + 1}", type="filepath", height=150)
-            res = gr.Dropdown(RESULT_CHOICES, value="unknown",
-                              label="Result (1-0 = White won)", scale=1)
-        slots.extend([img, res])
+    image_slots = [
+        gr.Image(label=f"Game {i + 1}", type="filepath", height=150)
+        for i in range(MAX_IMAGES)
+    ]
 
     convert_btn = gr.Button("Convert", variant="primary")
 
@@ -169,13 +291,14 @@ with gr.Blocks(title="Togyzkumalak Scoresheet Reader") as demo:
         headers=["game", "legal plies", "stopped", "beam PGN"],
         label="Results", wrap=True, interactive=False,
     )
+    warnings_md = gr.Markdown()
     gallery = gr.Gallery(label="Annotated reconstruction", columns=2, height="auto")
     zip_out = gr.File(label="Download all PGNs (zip)")
 
     convert_btn.click(
         _busy_wrapper,
-        inputs=[round_label, *slots],
-        outputs=[results_table, gallery, zip_out],
+        inputs=[meta_box, games_box, beam_dd, *image_slots],
+        outputs=[results_table, gallery, zip_out, warnings_md],
     )
 
 # Serialize CPU-heavy runs: callers wait in a bounded queue instead of

notebooks/colab_train.ipynb

@@ -4,12 +4,15 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Togyzkumalak Move OCR — Colab training\n",
+    "# Togyzkumalak OCR — Colab training (moves + diagram)\n",
     "\n",
     "Runtime → Change runtime type → **GPU**, then run the cells in order.\n",
     "\n",
-    "Get the project into Colab either by cloning your repo or by uploading a zip\n",
-    "of the project directory (without `venv311/`, `train_data5/`, `legacy/`)."
+    "Trains both classifiers on on-the-fly synthetic data (nothing to upload):\n",
+    "1. **moves** — 163-class move-cell reader (`11`..`99x` + `empty`)\n",
+    "2. **diagram** — 85-class board-diagram reader (values `0`-`81`, `x`, `-`, `empty`) for the kazan boxes and 2×9 pit grids of the summary strips\n",
+    "\n",
+    "At the end you download one `artifacts.zip` with the ONNX models + class sidecars, ready to drop into `models/` of the HuggingFace Space."
    ]
   },
   {
@@ -18,13 +21,23 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Option A: clone (edit URL)\n",
-    "# !git clone https://github.com/<you>/9OCR.git project\n",
-    "\n",
-    "# Option B: upload a zip named project.zip via the Files sidebar, then:\n",
-    "# !unzip -q project.zip -d project\n",
+    "!nvidia-smi"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Fetch the code at runtime.\n",
+    "# Option A (default): clone from GitHub\n",
+    "!git clone https://github.com/ansarzeinulla/9OCR.git project\n",
+    "%cd project\n",
     "\n",
-    "%cd project"
+    "# Option B: upload a zip named project.zip instead, then:\n",
+    "# from google.colab import files; files.upload()\n",
+    "# !unzip -q project.zip -d project && cd project"
    ]
   },
   {
@@ -33,7 +46,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "!pip install -q -r requirements.txt"
+    "!pip install -q -r requirements-train.txt"
    ]
   },
   {
@@ -42,8 +55,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Build glyph pools: ARDIS (European handwriting) + EMNIST via torchvision\n",
-    "# (--emnist is needed because ingredients/ is gitignored and absent here)\n",
+    "# Build glyph pools at runtime: ARDIS (European handwriting) + EMNIST via\n",
+    "# torchvision (--emnist is needed on Colab because ingredients/ is gitignored;\n",
+    "# digit 0 comes from these pools).\n",
     "!python scripts/get_glyphs.py --emnist"
    ]
   },
@@ -53,10 +67,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Sanity check: synthetic cells should resemble data/real_crops/*.jpg\n",
-    "!python -m togyz.synth --preview preview.png\n",
-    "from IPython.display import Image as IPImage\n",
-    "IPImage('preview.png')"
+    "# Sanity check both synthesizers before burning GPU time.\n",
+    "!python -m togyz.synth --preview preview_moves.png --task moves\n",
+    "!python -m togyz.synth --preview preview_diagram.png --task diagram\n",
+    "from IPython.display import Image as IPImage, display\n",
+    "display(IPImage('preview_moves.png'))\n",
+    "display(IPImage('preview_diagram.png'))"
    ]
   },
   {
@@ -65,7 +81,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "!python train.py --epochs 30 --samples-per-epoch 50000 --batch-size 256 --workers 2"
+    "# Train the move classifier (~30 epochs x 50k synthetic cells).\n",
+    "!python train.py --task moves --epochs 30 --samples-per-epoch 50000 --batch-size 256 --workers 2"
    ]
   },
   {
@@ -74,11 +91,43 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "!python eval.py --ckpt checkpoints/best.pt\n",
+    "# Train the unified board-diagram classifier.\n",
+    "!python train.py --task diagram --epochs 30 --samples-per-epoch 50000 --batch-size 256 --workers 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Export both checkpoints to single-file ONNX (+ .classes.json sidecars).\n",
+    "!python scripts/export_onnx.py"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Bundle the serving artifacts and download.\n",
+    "# In the Space's models/ dir: best.onnx stays best.onnx, the diagram model\n",
+    "# is renamed to diagram.onnx (+ diagram.classes.json).\n",
+    "import shutil, zipfile\n",
+    "from pathlib import Path\n",
+    "\n",
+    "out = Path('artifacts'); out.mkdir(exist_ok=True)\n",
+    "shutil.copy('checkpoints/best.onnx', out / 'best.onnx')\n",
+    "shutil.copy('checkpoints/best.classes.json', out / 'best.classes.json')\n",
+    "shutil.copy('checkpoints/diagram/best.onnx', out / 'diagram.onnx')\n",
+    "shutil.copy('checkpoints/diagram/best.classes.json', out / 'diagram.classes.json')\n",
+    "with zipfile.ZipFile('artifacts.zip', 'w') as zf:\n",
+    "    for f in out.iterdir():\n",
+    "        zf.write(f, arcname=f.name)\n",
     "\n",
-    "# download the trained model\n",
     "from google.colab import files\n",
-    "files.download('checkpoints/best.pt')"
+    "files.download('artifacts.zip')"
    ]
   }
  ],
@@ -89,5 +138,5 @@
   "language_info": {"name": "python"}
  },
  "nbformat": 4,
- "nbformat_minor": 4
+ "nbformat_minor": 5
 }

read_game.py

@@ -30,9 +30,10 @@ def main() -> None:
     parser.add_argument("image", help="scoresheet photo")
     parser.add_argument("--onnx", default="checkpoints/best.onnx",
                         help="move classifier ONNX (with a .classes.json sidecar)")
-    parser.add_argument("--kazan-onnx", default="checkpoints/kazan/best.onnx",
-                        help="kazan-number ONNX; checkpoint matching is skipped "
-                             "if the file does not exist")
+    parser.add_argument("--diagram-onnx", default="checkpoints/diagram/best.onnx",
+                        help="board-diagram ONNX (kazan boxes + pit cells); "
+                             "checkpoint matching is skipped if the file "
+                             "does not exist")
     parser.add_argument("--out", default=None, help="output dir (default: out/<image stem>)")
     parser.add_argument("--topk", type=int, default=5)
     parser.add_argument("--beam-width", type=int, default=1024,
@@ -52,15 +53,15 @@ def main() -> None:
     out_dir.mkdir(parents=True, exist_ok=True)
 
     moves_clf = load_classifier(args.onnx)
-    kazan_clf = None
-    if Path(args.kazan_onnx).exists():
-        kazan_clf = load_classifier(args.kazan_onnx)
+    diagram_clf = None
+    if Path(args.diagram_onnx).exists():
+        diagram_clf = load_classifier(args.diagram_onnx)
     else:
-        print(f"No kazan classifier at {args.kazan_onnx} - checkpoint matching off.")
+        print(f"No diagram classifier at {args.diagram_onnx} - checkpoint matching off.")
 
     print(f"Reading {args.image} ...")
     out = run_pipeline(
-        args.image, moves_clf, kazan_clf,
+        args.image, moves_clf, diagram_clf,
         result=args.result, topk=args.topk,
         beam_width=args.beam_width, per_ply=args.beam_top,
         temperature=args.temperature, tta=not args.no_tta,
@@ -71,9 +72,13 @@ def main() -> None:
         print(f"WARNING: median cell height is only {out['median_cell_height']}px - "
               "accuracy suffers at this resolution; re-photograph at full camera "
               "resolution if possible.")
-    if out["kazan_report"]:
-        print("Kazan checkpoints read: "
-              f"{[(r['move'], r['side'], r['read']) for r in out['kazan_report']]}")
+    if out["checkpoint_report"]:
+        kaz = [(r["move"], r["side"], r["read"])
+               for r in out["checkpoint_report"] if r["kind"] == "kazan"]
+        pits = sum(1 for r in out["checkpoint_report"] if r["kind"] == "pit")
+        print(f"Diagram checkpoints read: kazans {kaz}, {pits} pit cells")
+    for warning in out["warnings"]:
+        print(f"WARNING: {warning}")
 
     result = {"image": args.image, "onnx": args.onnx, **out["game_json"]}
     (out_dir / "game.json").write_text(json.dumps(result, indent=1))

scripts/export_onnx.py

@@ -4,7 +4,7 @@
 
 Writes, next to each source checkpoint:
     checkpoints/best.onnx          + checkpoints/best.classes.json
-    checkpoints/kazan/best.onnx    + checkpoints/kazan/best.classes.json
+    checkpoints/diagram/best.onnx  + checkpoints/diagram/best.classes.json
 
 ONNX carries no Python metadata, so the class list (which maps output index
 -> move/kazan string) is saved as a JSON sidecar the serving code reads. The
@@ -66,8 +66,8 @@ def main() -> None:
     parser.add_argument(
         "checkpoints",
         nargs="*",
-        default=["checkpoints/best.pt", "checkpoints/kazan/best.pt"],
-        help="checkpoint(s) to export (default: moves + kazan best.pt)",
+        default=["checkpoints/best.pt", "checkpoints/diagram/best.pt"],
+        help="checkpoint(s) to export (default: moves + diagram best.pt)",
     )
     args = parser.parse_args()
     for path in args.checkpoints:

togyz/classes.py

@@ -18,6 +18,10 @@ CLASSES = MOVES + [EMPTY]
 CLASS_TO_IDX = {name: i for i, name in enumerate(CLASSES)}
 NUM_CLASSES = len(CLASSES)  # 163
 
-# Kazan checkpoint numbers on the scoresheet summary strips: always exactly
-# two digits (with leading zero), values 00-81, never an 'x'.
-KAZAN_CLASSES = [f"{v:02d}" for v in range(82)] + [EMPTY]
+# Unified board-diagram classes for the summary strips: pit/kazan values
+# 0-81 (rendered as 1 or 2 digits), 'x' = tuzdyk mark, '-' = zero mark.
+# At inference the pipeline filters per context: kazan boxes allow only
+# 10-81, pit cells allow '-', 'x' and 0-81.
+DASH = "-"
+TUZDYK = "x"
+DIAGRAM_CLASSES = [str(v) for v in range(82)] + [TUZDYK, DASH, EMPTY]

togyz/dataset.py

@@ -10,7 +10,7 @@ from torch.utils.data import Dataset
 from .classes import CLASSES, CLASS_TO_IDX
 from .glyphs import GlyphSampler
 from .preprocess import preprocess_file, preprocess_pil
-from .synth import synthesize_cell
+from .synth import synthesize_cell, synthesize_diagram_cell
 
 
 class SyntheticCellDataset(Dataset):
@@ -22,12 +22,13 @@ class SyntheticCellDataset(Dataset):
     """
 
     def __init__(self, num_samples: int, seed: int | None = None, project_root=".",
-                 classes: list[str] | None = None):
+                 classes: list[str] | None = None, task: str = "moves"):
         self.sampler = GlyphSampler(project_root)
         self.num_samples = num_samples
         self.seed = seed
         self.classes = classes or CLASSES
         self.class_to_idx = {c: i for i, c in enumerate(self.classes)}
+        self.synth_fn = synthesize_diagram_cell if task == "diagram" else synthesize_cell
 
     def __len__(self) -> int:
         return self.num_samples
@@ -38,7 +39,7 @@ class SyntheticCellDataset(Dataset):
         else:
             rng = random.Random(random.getrandbits(64))
         class_name = self.classes[rng.randrange(len(self.classes))]
-        img = synthesize_cell(class_name, self.sampler, rng)
+        img = self.synth_fn(class_name, self.sampler, rng)
         return torch.from_numpy(preprocess_pil(img)), self.class_to_idx[class_name]
 
 

togyz/pipeline.py

@@ -2,7 +2,7 @@
 
 This is the single inference core shared by the CLI (`read_game.py`) and the
 Gradio demo (`app.py`). It deliberately imports **no torch**: the move and
-kazan classifiers run as ONNX sessions, and every array op is numpy. The
+board-diagram classifiers run as ONNX sessions, and every array op is numpy. The
 heavy lifting of cell extraction (`togyz.sheet`) and the rules engine
 (`togyz.rules`) are already torch-free and imported as-is.
 
@@ -13,6 +13,7 @@ git history (sheet 1 = 36/39 exact plies). Keep it numerically equivalent.
 
 from __future__ import annotations
 
+import heapq
 import json
 import math
 from dataclasses import dataclass, field
@@ -22,10 +23,10 @@ import numpy as np
 import onnxruntime as ort
 from PIL import Image
 
-from .classes import EMPTY
+from .classes import DASH, EMPTY, TUZDYK
 from .preprocess import preprocess_pil
 from .rules import Game
-from .sheet import clean_cell, extract_cells, extract_kazan_cells, render_overlay
+from .sheet import clean_cell, extract_cells, extract_diagram_cells, render_overlay
 
 MIN_CELL_HEIGHT = 35  # px; below this, resolution is low (surfaced as a warning)
 MIN_INK_RATIO = 0.02  # cells with less handwriting ink than this are empty
@@ -36,8 +37,16 @@ TTA_SHIFTS = [(0, 0), (-2, 0), (2, 0), (0, -2), (0, 2), (-2, -2), (2, 2)]
 EXACT_WEIGHT = 0.5  # weight of the exact class-string probability
 FACTOR_WEIGHT = 0.5  # weight of the factorized digit-marginal probability
 KAZAN_FLOOR = 1e-4  # a misread checkpoint must not single-handedly kill truth
+PIT_FLOOR = 1e-3  # pit reads are noisier than kazans: weaker veto power
+PIT_WEIGHT = 0.5  # 18 pit terms per checkpoint must not swamp the rest
 RESULT_CODES = {"1-0": 0, "0-1": 1, "draw": -1}
 
+# Memory guard for user-selectable beam widths (the UI offers up to 2^30):
+# a late-game hypothesis (copied move lists etc.) costs roughly this many
+# bytes, and the expanded pool peaks at beam_width * per_ply hypotheses.
+HYPOTHESIS_BYTES = 25_000
+MAX_POOL_BYTES = 2 * 2**30  # ~2 GB hypothesis budget
+
 
 # --------------------------------------------------------------------------- #
 # ONNX session loading
@@ -152,27 +161,66 @@ def _result_consistency(game: Game, result: str | None) -> int:
     return 1 if leader == want else 0
 
 
-def _kazan_logp(game: Game, checkpoint_probs) -> float:
-    """Log-likelihood of a hypothesis' kazan counts under the checkpoint reads."""
+def _filter_probs(probs: np.ndarray, classes: list[str], allowed: set[str]):
+    """Zero out contextually impossible classes and renormalize.
+
+    Returns None when nothing observable remains (e.g. a blank crop whose
+    mass sat entirely on 'empty') - the cell is then treated as unobserved.
+    """
+    filtered = np.array(
+        [p if c in allowed else 0.0 for c, p in zip(classes, probs)], np.float64
+    )
+    total = float(filtered.sum())
+    if total < 1e-6:
+        return None
+    return filtered / total
+
+
+def _diagram_logp(game: Game, cp) -> float:
+    """Log-likelihood of a hypothesis' board state under one checkpoint's
+    diagram reads (kazan boxes + observed pit cells)."""
     logp = 0.0
     for side, player in (("W", 0), ("B", 1)):
-        probs = checkpoint_probs.get(side)
+        probs = cp["kazan"].get(side)
         if probs is None:
             continue
         value = game.kazans[player]
         p = float(probs[value]) if value < len(probs) else 0.0
         logp += math.log(max(p, KAZAN_FLOOR))
+
+    for (side, index), (values, p_x, p_dash) in cp["pits"].items():
+        player = 0 if side == "W" else 1
+        pit_pos = player * 9 + (index - 1)
+        # a pit owned as tuzdyk by the opponent is written as 'x'
+        if game.tuzdyks[1 - player] == index - 1:
+            p = p_x
+        else:
+            count = game.pits[pit_pos]
+            if count == 0:
+                p = p_dash + float(values[0])
+            else:
+                p = float(values[count]) if count < len(values) else 0.0
+        logp += PIT_WEIGHT * math.log(max(p, PIT_FLOOR))
     return logp
 
 
+def safe_beam_width(beam_width: int, per_ply: int = 9) -> int:
+    """Largest beam width that keeps the expanded pool inside the RAM budget."""
+    return max(1, min(beam_width, MAX_POOL_BYTES // (HYPOTHESIS_BYTES * max(per_ply, 1))))
+
+
 def beam_decode(ply_probs, classes, beam_width=1024, per_ply=9, result=None,
                 checkpoints=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 kazan-checkpoint and known-result consistency.
+    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.
 
     Returns (annotated_moves, log_prob, per_ply_share_list).
     """
+    beam_width = safe_beam_width(beam_width, per_ply)
     checkpoints = checkpoints or {}
     class_idx = {c: i for i, c in enumerate(classes)}
     active = [(Game(), [], [], 0.0)]  # (game, annotated_moves, chosen_shares, logp)
@@ -188,12 +236,13 @@ def beam_decode(ply_probs, classes, beam_width=1024, per_ply=9, result=None,
                 annotated = played.notation + ("+" if played.captured else "")
                 new_logp = logp + math.log(max(raw, 1e-12))
                 if len(moves) + 1 in checkpoints:
-                    new_logp += _kazan_logp(nxt, checkpoints[len(moves) + 1])
+                    new_logp += _diagram_logp(nxt, checkpoints[len(moves) + 1])
                 hyp = (nxt, moves + [annotated], chosen + [share], new_logp)
                 (finished if nxt.is_over else expanded).append(hyp)
         if not expanded:
             break
-        active = sorted(expanded, key=lambda h: -h[3])[:beam_width]
+        # nlargest instead of a full sort: wide beams stay tractable
+        active = heapq.nlargest(beam_width, expanded, key=lambda h: h[3])
 
     pool = active + finished
     best = max(pool, key=lambda h: (len(h[1]), _result_consistency(h[0], result), h[3]))
@@ -212,8 +261,8 @@ def format_pgn(plies: list[str]) -> str:
 # --------------------------------------------------------------------------- #
 # The single entry point
 # --------------------------------------------------------------------------- #
-def run_pipeline(image, moves_clf: Classifier, kazan_clf: Classifier | None = None,
-                 result=None, topk=5, beam_width=1024, per_ply=9,
+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):
     """Read one scoresheet image into game records.
 
@@ -222,13 +271,20 @@ def run_pipeline(image, moves_clf: Classifier, kazan_clf: Classifier | None = No
     (the CLI uses it, the Gradio app does not).
 
     Keys: beam_pgn, raw_pgn, legal_pgn, game_json (dict), stopped,
-    plies_scanned, legal_plies, beam_plies, kazan_report, annotated_image (PIL),
-    median_cell_height, low_resolution (bool).
+    plies_scanned, legal_plies, beam_plies, checkpoint_report, warnings,
+    annotated_image (PIL), median_cell_height, low_resolution (bool).
     """
     classes = moves_clf.classes
     empty_idx = classes.index(EMPTY)
+    warnings: list[str] = []
+    effective_width = safe_beam_width(beam_width, per_ply)
+    if effective_width < beam_width:
+        warnings.append(
+            f"beam width {beam_width} trimmed to {effective_width} to stay "
+            "inside the memory budget"
+        )
 
-    cells, sheet = extract_cells(image)
+    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]
 
@@ -294,31 +350,62 @@ def run_pipeline(image, moves_clf: Classifier, kazan_clf: Classifier | None = No
                 in_sync = False
             raw_pgn.append(raw)
 
-    # kazan checkpoint numbers from the summary strips (every 10 moves)
+    # board-diagram evidence from the summary strips (every 10 moves):
+    # kazan boxes (values 10-81) plus the 2x9 pit grid ('x'/'-'/0-81)
     checkpoints, checkpoint_report = {}, []
-    if kazan_clf is not None:
-        kz_cells = [(c, *clean_cell(c.image)) for c in extract_kazan_cells(sheet)]
-        kz_cells = [(c, img) for c, img, ink in kz_cells if ink >= MIN_KAZAN_INK]
-        if kz_cells:
-            kz_probs = classify_cells(
-                [img for _, img in kz_cells], kazan_clf,
+    diagram_cells, diagram_labels = [], {}
+    if diagram_clf is not None:
+        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))
+        pit_allowed = set(value_names) | {TUZDYK, DASH}
+        x_idx = dg_classes.index(TUZDYK)
+        dash_idx = dg_classes.index(DASH)
+
+        dg_cells = [(c, *clean_cell(c.image)) for c in extract_diagram_cells(sheet)]
+        dg_cells = [
+            (c, img) for c, img, ink in dg_cells
+            if ink >= (MIN_KAZAN_INK if c.kind == "kazan" else MIN_INK_RATIO)
+        ]
+        if dg_cells:
+            dg_probs = classify_cells(
+                [img for _, img in dg_cells], diagram_clf,
                 temperature=temperature, tta=tta,
             )
-            for (cell, img), probs in zip(kz_cells, kz_probs):
-                if kazan_clf.classes[int(np.argmax(probs))] == EMPTY:
+            for (cell, img), probs in zip(dg_cells, dg_probs):
+                if dg_classes[int(np.argmax(probs))] == EMPTY:
+                    continue
+                allowed = kazan_allowed if cell.kind == "kazan" else pit_allowed
+                filtered = _filter_probs(probs, dg_classes, allowed)
+                if filtered is None:
                     continue
-                values = probs[:82] / max(float(probs[:82].sum()), 1e-9)
-                checkpoints.setdefault(cell.move_no * 2, {})[cell.side] = values
-                top = int(np.argmax(values))
-                checkpoint_report.append(
-                    {"move": cell.move_no, "side": cell.side,
-                     "read": f"{top:02d}", "prob": round(float(values[top]), 4)}
+                cp = checkpoints.setdefault(
+                    cell.move_no * 2, {"kazan": {}, "pits": {}}
                 )
+                values = filtered[:82]
+                top_idx = int(np.argmax(filtered))
+                read = dg_classes[top_idx]
+                if cell.kind == "kazan":
+                    # renormalized over 10-81 only, value-indexed for scoring
+                    cp["kazan"][cell.side] = values / max(float(values.sum()), 1e-9)
+                else:
+                    cp["pits"][(cell.side, cell.pit_index)] = (
+                        values, float(filtered[x_idx]), float(filtered[dash_idx])
+                    )
+                diagram_cells.append(cell)
+                diagram_labels[(cell.move_no, cell.kind, cell.side, cell.pit_index)] = read
+                entry = {"move": cell.move_no, "kind": cell.kind, "side": cell.side,
+                         "read": read, "prob": round(float(filtered[top_idx]), 4)}
+                if cell.kind == "pit":
+                    entry["index"] = cell.pit_index
+                checkpoint_report.append(entry)
                 if cells_dir:
-                    img.save(cells_dir / f"kazan_{cell.move_no:02d}_{cell.side}.png")
+                    suffix = f"_{cell.pit_index}" if cell.kind == "pit" else ""
+                    img.save(cells_dir /
+                             f"{cell.kind}_{cell.move_no:02d}_{cell.side}{suffix}.png")
 
     beam_moves, beam_logp, beam_shares = beam_decode(
-        ply_prob_arrays, classes, beam_width, per_ply,
+        ply_prob_arrays, classes, effective_width, per_ply,
         result=result, checkpoints=checkpoints,
     )
     beam_detail = [
@@ -329,14 +416,19 @@ def run_pipeline(image, moves_clf: Classifier, kazan_clf: Classifier | None = No
     for p, m in zip(plies, beam_moves):
         labels[(p["move"], p["side"])] = m
 
-    annotated = render_overlay(sheet, cells[: len(plies)], labels)
+    annotated = render_overlay(
+        sheet, cells[: len(plies)], labels,
+        diagram_cells=diagram_cells, diagram_labels=diagram_labels,
+        gridlines=gridlines,
+    )
 
     game_json = {
         "plies_scanned": len(plies),
         "legal_plies": len(legal_pgn),
         "beam": {"plies": len(beam_moves), "log_prob": round(beam_logp, 3), "moves": beam_detail},
-        "kazan_checkpoints": checkpoint_report,
+        "checkpoints": checkpoint_report,
         "stopped": stopped or {"reason": "sheet exhausted"},
+        "warnings": warnings,
         "plies": plies,
     }
     return {
@@ -348,7 +440,8 @@ def run_pipeline(image, moves_clf: Classifier, kazan_clf: Classifier | None = No
         "plies_scanned": len(plies),
         "legal_plies": len(legal_pgn),
         "beam_plies": len(beam_moves),
-        "kazan_report": checkpoint_report,
+        "checkpoint_report": checkpoint_report,
+        "warnings": warnings,
         "annotated_image": annotated,
         "median_cell_height": int(median_h),
         "low_resolution": median_h < MIN_CELL_HEIGHT,

togyz/sheet.py

@@ -26,11 +26,13 @@ CELL_MARGIN = 0.15  # expand crops; handwriting overflows the printed cells
 
 @dataclass
 class Cell:
-    move_no: int  # 1-80
-    side: str  # "W" (Bast.) or "B" (Kost.)
+    move_no: int  # move cells: 1-80; diagram cells: checkpoint move (10..80)
+    side: str  # "W" (Bast.) or "B" (Kost.); for pits this is the row owner
     bbox: tuple[int, int, int, int]  # x, y, w, h in original image coords
     image: Image.Image
     quad: np.ndarray | None = None  # 4x2 original-image corners (tilt-aware)
+    kind: str = "move"  # "move" | "kazan" | "pit"
+    pit_index: int = 0  # 1-9 for pit cells (scoresheet numbering), else 0
 
 
 def _grid_mask(gray: np.ndarray) -> np.ndarray:
@@ -82,6 +84,12 @@ def _find_table_quads(gray: np.ndarray) -> list[np.ndarray]:
             f"Found only {len(quads)} move tables (need {TABLES}). "
             "Check photo quality/framing."
         )
+    # the 8 move tables all have near-identical area; drop outliers like the
+    # footer/signature block, which can out-size a genuine table
+    median_area = float(np.median([a for a, _ in quads]))
+    consistent = [(a, q) for a, q in quads if 0.5 * median_area < a < 2.0 * median_area]
+    if len(consistent) >= TABLES:
+        quads = consistent
     quads = [q for _, q in sorted(quads, key=lambda t: -t[0])[:TABLES]]
     # split into top/bottom bands by y, order each band left to right
     quads.sort(key=lambda q: q[:, 1].mean())
@@ -139,12 +147,26 @@ def _row_lines(warped: np.ndarray) -> list[int]:
     return rows
 
 
-def extract_cells(image) -> tuple[list[Cell], Image.Image]:
+def _map_segment(inverse: np.ndarray, p0, p1) -> tuple[tuple[int, int], tuple[int, int]]:
+    """Map a rectified-table segment back onto the original sheet."""
+    pts = np.array([p0, p1], np.float32).reshape(-1, 1, 2)
+    mapped = cv2.perspectiveTransform(pts, inverse).reshape(2, 2)
+    return (
+        (int(mapped[0][0]), int(mapped[0][1])),
+        (int(mapped[1][0]), int(mapped[1][1])),
+    )
+
+
+def extract_cells(image, with_gridlines: bool = False):
     """All 160 move cells (80 moves x W/B) plus the upright sheet image.
 
     `image` is a path (str/Path) or a PIL.Image. Each table is
     perspective-rectified before being split, so cell crops stay aligned even
     on tilted phone photos.
+
+    Returns (cells, sheet) or, with `with_gridlines`, (cells, sheet, gridlines)
+    where gridlines is a list of ((x0, y0), (x1, y1)) sheet-coordinate segments
+    for every row/column separator used during segmentation.
     """
     if isinstance(image, Image.Image):
         pil = ImageOps.exif_transpose(image).convert("L")
@@ -154,11 +176,17 @@ def extract_cells(image) -> tuple[list[Cell], Image.Image]:
     gray = np.asarray(pil)
 
     cells: list[Cell] = []
+    gridlines: list[tuple[tuple[int, int], tuple[int, int]]] = []
     for t, quad in enumerate(_find_table_quads(gray)):
         warped, inverse = _rectify_table(gray, quad)
         th, tw = warped.shape
         rows = _row_lines(warped)
         cols = [round(f * tw) for f in COLUMN_FRACTIONS]
+        if with_gridlines:
+            for y in rows:
+                gridlines.append(_map_segment(inverse, (0, y), (tw, y)))
+            for x in cols:
+                gridlines.append(_map_segment(inverse, (x, rows[0]), (x, rows[-1])))
         for r in range(ROWS_PER_TABLE):
             y0, y1 = rows[r + 1], rows[r + 2]  # rows[0..1] is the header
             for col_index, side in ((1, "W"), (2, "B")):
@@ -180,37 +208,132 @@ def extract_cells(image) -> tuple[list[Cell], Image.Image]:
                 cells.append(
                     Cell(t * ROWS_PER_TABLE + r + 1, side, bbox, crop, sheet_quad)
                 )
+    if with_gridlines:
+        return cells, pil, gridlines
     return cells, pil
 
 
-def extract_kazan_cells(sheet: Image.Image) -> list[Cell]:
-    """Kazan checkpoint boxes from the 8 summary strips between table bands.
+PIT_MARGIN_Y = 0.25  # pit digits regularly overflow the printed row height
+PIT_MARGIN_X = 0.15  # and bleed a little into neighboring columns
 
-    After every 10 moves the scorer records both kazan counts next to a small
-    board diagram: the box protruding ABOVE the strip is the Kost. (Black)
-    kazan, the box BELOW is the Bast. (White) kazan. Returned as Cell objects
-    with move_no = the checkpoint move (10, 20, ... 80) and side "B"/"W".
-    The 2x9 board-pit grid itself is ignored (too messy to read reliably).
+
+def _strip_pit_grid(mask: np.ndarray, gray: np.ndarray,
+                    strip_bbox: tuple[int, int, int, int], checkpoint: int) -> list[Cell]:
+    """The 18 pit cells of the 2x9 board diagram inside one summary strip.
+
+    Upper row = Kost. (Black), pit indexes 9..1 left-to-right; lower row =
+    Bast. (White), indexes 1..9. Grid lines are detected inside the strip;
+    when detection is incomplete the known uniform 2x9 structure is used.
+    """
+    x, y, bw, bh = strip_bbox
+    h, w = gray.shape
+    sub = mask[y : y + bh, x : x + bw]
+
+    # horizontal separators: top / middle / bottom of the 2x9 grid
+    hlines = _detect_lines(sub, axis=0, min_frac=0.55)
+    if len(hlines) < 2:
+        return []
+    top, bottom = hlines[0], hlines[-1]
+    if bottom - top < 6:
+        return []
+    mid_target = (top + bottom) / 2
+    inner = [v for v in hlines if top + 2 < v < bottom - 2]
+    middle = min(inner, key=lambda v: abs(v - mid_target)) if inner else int(round(mid_target))
+
+    # vertical separators: 10 column lines across the grid band
+    band = sub[top : bottom + 1]
+    vlines = _detect_lines(band, axis=1, min_frac=0.5)
+    if len(vlines) != 10:
+        if len(vlines) >= 2:
+            x0, x1 = vlines[0], vlines[-1]
+        else:
+            xs = np.where(band.sum(axis=0) > 0)[0]
+            x0, x1 = (int(xs.min()), int(xs.max())) if len(xs) else (0, bw - 1)
+        vlines = [round(x0 + j * (x1 - x0) / 9) for j in range(10)]
+
+    cells: list[Cell] = []
+    rows = ((top, middle, "B"), (middle, bottom, "W"))
+    for ry0, ry1, owner in rows:
+        my = round((ry1 - ry0) * PIT_MARGIN_Y)
+        for j in range(9):
+            cx0, cx1 = vlines[j], vlines[j + 1]
+            mx = round((cx1 - cx0) * PIT_MARGIN_X)
+            gx0 = max(0, x + cx0 - mx)
+            gx1 = min(w, x + cx1 + mx)
+            gy0 = max(0, y + ry0 - my)
+            gy1 = min(h, y + ry1 + my)
+            if gx1 - gx0 < 4 or gy1 - gy0 < 4:
+                continue
+            pit_index = 9 - j if owner == "B" else j + 1
+            cells.append(Cell(
+                checkpoint, owner, (gx0, gy0, gx1 - gx0, gy1 - gy0),
+                Image.fromarray(gray[gy0:gy1, gx0:gx1]),
+                kind="pit", pit_index=pit_index,
+            ))
+    return cells
+
+
+def extract_diagram_cells(sheet: Image.Image) -> list[Cell]:
+    """All board-diagram cells from the 8 summary strips between table bands.
+
+    After every 10 moves the scorer fills a small board diagram: the box
+    protruding ABOVE the strip is the Kost. (Black) kazan, the box BELOW is
+    the Bast. (White) kazan, and the 2x9 grid between them holds the pit
+    counts ('x' = tuzdyk, '-' = 0, or a 1-2 digit number). Returned as Cell
+    objects with move_no = the checkpoint move (10, 20, ... 80), kind
+    "kazan"/"pit", side "B"/"W" (row owner for pits) and pit_index 1-9.
     """
     gray = np.asarray(sheet)
     h, w = gray.shape
     mask = _grid_mask(gray)
-    table_area = float(np.median([cv2.contourArea(q) for q in _find_table_quads(gray)]))
-
+    quads = _find_table_quads(gray)
+    table_area = float(np.median([cv2.contourArea(q) for q in quads]))
+
+    # Diagram strips are located by table position, not by enumeration order:
+    # the strip for checkpoint 10*(4*band + col + 1) sits below table `col` of
+    # band `band`, in the gap under that band. Handwriting sometimes bridges
+    # neighboring strips into one connected component, so wide components are
+    # kept and carved by each table's x-range.
     n, _, stats, _ = cv2.connectedComponentsWithStats(mask)
-    strips = [
+    components = [
         tuple(int(v) for v in stats[i][:4])
         for i in range(1, n)
-        if stats[i][2] > w * 0.1
+        if stats[i][2] > w * 0.05
         and h * 0.02 < stats[i][3] < h * 0.12
         and stats[i][2] > 1.5 * stats[i][3]
-        and stats[i][2] * stats[i][3] < 0.9 * table_area
+        and stats[i][2] * stats[i][3] < 2.5 * table_area
     ]
-    strips.sort(key=lambda b: (b[1] > h / 2, b[0]))  # band, then left-to-right
+
+    top_quads, bottom_quads = quads[:4], quads[4:]
+    zones = (
+        (top_quads, (max(q[:, 1].max() for q in top_quads),
+                     min(q[:, 1].min() for q in bottom_quads))),
+        (bottom_quads, (max(q[:, 1].max() for q in bottom_quads), h)),
+    )
+
+    strips: list[tuple[int, tuple[int, int, int, int]]] = []
+    for band, (band_quads, (zy0, zy1)) in enumerate(zones):
+        for col, quad in enumerate(band_quads):
+            tx0, tx1 = float(quad[:, 0].min()), float(quad[:, 0].max())
+            parts = []
+            for cx, cy, cw, ch in components:
+                yc = cy + ch / 2
+                overlap = min(cx + cw, tx1) - max(cx, tx0)
+                if zy0 < yc < zy1 and overlap > 0.3 * (tx1 - tx0):
+                    parts.append((cx, cy, cw, ch))
+            if not parts:
+                continue
+            pad = int(0.02 * w)
+            sx0 = max(int(tx0) - pad, min(cx for cx, *_ in parts))
+            sx1 = min(int(tx1) + pad, max(cx + cw for cx, _, cw, _ in parts))
+            sy0 = min(cy for _, cy, *_ in parts)
+            sy1 = max(cy + ch for _, cy, _, ch in parts)
+            checkpoint = (band * 4 + col + 1) * 10
+            strips.append((checkpoint, (sx0, sy0, sx1 - sx0, sy1 - sy0)))
 
     cells: list[Cell] = []
-    for k, (x, y, bw, bh) in enumerate(strips[:8]):
-        checkpoint = (k + 1) * 10
+    for checkpoint, (x, y, bw, bh) in strips:
+        cells.extend(_strip_pit_grid(mask, gray, (x, y, bw, bh), checkpoint))
         sub = mask[y : y + bh, x : x + bw]
         long_rows = np.where((sub > 0).sum(axis=1) > bw * 0.55)[0]
         if len(long_rows) == 0:
@@ -245,6 +368,7 @@ def extract_kazan_cells(sheet: Image.Image) -> list[Cell]:
             cells.append(Cell(
                 checkpoint, side, (cx0, cy0, cx1 - cx0, cy1 - cy0),
                 Image.fromarray(gray[cy0:cy1, cx0:cx1]),
+                kind="kazan",
             ))
     return cells
 
@@ -281,9 +405,19 @@ def clean_cell(image: Image.Image) -> tuple[Image.Image, float]:
     return Image.fromarray(cleaned), ink_ratio
 
 
-def render_overlay(sheet: Image.Image, cells: list[Cell], labels: dict | None = None) -> Image.Image:
-    """Debug image: cell boxes (green) with optional predicted labels (red)."""
+def render_overlay(sheet: Image.Image, cells: list[Cell], labels: dict | None = None,
+                   diagram_cells: list[Cell] | None = None,
+                   diagram_labels: dict | None = None,
+                   gridlines: list | None = None) -> Image.Image:
+    """Debug image: move-cell boxes (green) with predicted labels (red),
+    kazan boxes (blue), pit cells (orange) with their reads, and the
+    segmentation gridlines (gray).
+
+    `diagram_labels` is keyed by (move_no, kind, side, pit_index).
+    """
     vis = cv2.cvtColor(np.asarray(sheet), cv2.COLOR_GRAY2BGR)
+    for p0, p1 in gridlines or []:
+        cv2.line(vis, p0, p1, (160, 160, 160), 1, cv2.LINE_AA)
     for cell in cells:
         x, y, w, h = cell.bbox
         if cell.quad is not None:
@@ -295,4 +429,13 @@ def render_overlay(sheet: Image.Image, cells: list[Cell], labels: dict | None =
             if text:
                 cv2.putText(vis, text, (x + 2, y + h - 3),
                             cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 0, 255), 1, cv2.LINE_AA)
+    for cell in diagram_cells or []:
+        x, y, w, h = cell.bbox
+        color = (200, 120, 0) if cell.kind == "kazan" else (0, 140, 255)  # BGR
+        cv2.rectangle(vis, (x, y), (x + w, y + h), color, 1)
+        if diagram_labels:
+            text = diagram_labels.get((cell.move_no, cell.kind, cell.side, cell.pit_index))
+            if text:
+                cv2.putText(vis, text, (x + 1, y + h - 2),
+                            cv2.FONT_HERSHEY_SIMPLEX, 0.35, (0, 0, 255), 1, cv2.LINE_AA)
     return Image.fromarray(cv2.cvtColor(vis, cv2.COLOR_BGR2RGB))

togyz/synth.py

@@ -20,7 +20,7 @@ import cv2
 import numpy as np
 from PIL import Image
 
-from .classes import CLASSES, EMPTY
+from .classes import CLASSES, DASH, DIAGRAM_CLASSES, EMPTY, TUZDYK
 from .glyphs import GlyphSampler
 
 # working resolution: glyphs are rendered at this digit height, the finished
@@ -206,12 +206,116 @@ def synthesize_cell(
     return Image.fromarray(_camera_effects(img, rng))
 
 
-def _preview(path: str, seed: int, count: int) -> None:
+def _paste_clipped(alpha: np.ndarray, mask: np.ndarray, x: int, y: int) -> None:
+    """max-blend `mask` onto `alpha` at (x, y), clipping at the canvas edges
+    (diagram digits regularly overflow the printed cell borders)."""
+    h, w = alpha.shape
+    mh, mw = mask.shape
+    x0, y0 = max(0, x), max(0, y)
+    x1, y1 = min(w, x + mw), min(h, y + mh)
+    if x1 <= x0 or y1 <= y0:
+        return
+    region = alpha[y0:y1, x0:x1]
+    np.maximum(region, mask[y0 - y : y1 - y, x0 - x : x1 - x], out=region)
+
+
+def _dash_mask(rng: random.Random) -> np.ndarray:
+    """A handwritten '-' (zero mark): one short, slightly curved stroke."""
+    s = 64
+    canvas = np.zeros((s // 2, s), np.float32)
+    y = s // 4 + rng.randint(-4, 4)
+    x0, x1 = rng.randint(2, 12), rng.randint(52, 62)
+    mid = ((x0 + x1) // 2, y + rng.randint(-6, 6))
+    pts = np.array([(x0, y), mid, (x1, y + rng.randint(-4, 4))], np.int32)
+    cv2.polylines(canvas, [pts], False, 1.0, rng.randint(3, 6), cv2.LINE_AA)
+    ys, xs = np.where(canvas > 0.1)
+    return canvas[ys.min() : ys.max() + 1, xs.min() : xs.max() + 1]
+
+
+def synthesize_diagram_cell(
+    class_name: str, sampler: GlyphSampler, rng: random.Random
+) -> Image.Image:
+    """Render one board-diagram cell photo: a pit/kazan value 0-81, a tuzdyk
+    'x', a zero '-', or 'empty'.
+
+    Diagram crops differ from move cells: they are squarer, always framed by
+    printed grid lines (drawn on several edges), digits often overflow the
+    row height and neighboring digits bleed into the crop sides.
+    """
+    # --- ink masks for the written symbol ---
+    if class_name == EMPTY:
+        masks: list[np.ndarray] = []
+    elif class_name == DASH:
+        m = _dash_mask(rng)
+        target_w = max(10, int(RENDER_DIGIT_H * rng.uniform(0.35, 0.8)))
+        target_h = max(4, int(m.shape[0] * target_w / m.shape[1]))
+        masks = [cv2.resize(m, (target_w, target_h), interpolation=cv2.INTER_LINEAR)]
+    else:
+        chars = class_name if class_name == TUZDYK else str(int(class_name))
+        slant = rng.uniform(-0.10, 0.40)
+        masks = []
+        for char in chars:
+            mask = sampler.sample(char, rng)
+            rel_h = rng.uniform(0.55, 0.95) if char == "x" else rng.uniform(0.85, 1.15)
+            target_h = max(8, int(RENDER_DIGIT_H * rel_h))
+            target_w = max(4, int(mask.shape[1] * target_h / mask.shape[0]))
+            mask = cv2.resize(mask, (target_w, target_h), interpolation=cv2.INTER_LINEAR)
+            mask = _transform_mask(mask, rng.uniform(-7, 7), slant + rng.uniform(-0.05, 0.05))
+            masks.append(_vary_thickness(mask, rng))
+
+    gaps = [int(rng.uniform(-0.22, 0.10) * masks[i].shape[1]) for i in range(len(masks) - 1)]
+    total_w = sum(m.shape[1] for m in masks) + sum(gaps)
+    max_h = max((m.shape[0] for m in masks), default=RENDER_DIGIT_H)
+
+    # --- cell geometry: squarish crop, digits may overflow it vertically ---
+    fill = rng.uniform(0.55, 1.25)  # >1: taller than the printed row
+    cell_h = max(24, int(max_h / fill))
+    aspect = rng.uniform(0.9, 1.9)
+    cell_w = max(int(cell_h * aspect), int(total_w * rng.uniform(0.85, 1.15)))
+
+    alpha = np.zeros((cell_h, cell_w), np.float32)
+    x = int((cell_w - total_w) * rng.uniform(0.1, 0.9)) if masks else 0
+    for i, mask in enumerate(masks):
+        mh = mask.shape[0]
+        y = int((cell_h - mh) * rng.uniform(0.25, 0.75)) if cell_h > mh else int(
+            (cell_h - mh) * rng.uniform(0.3, 0.7)
+        )
+        _paste_clipped(alpha, mask, x, y)
+        if i < len(gaps):
+            x += mask.shape[1] + gaps[i]
+
+    # --- neighbor bleed: fragments of adjacent cells' digits at the sides ---
+    if rng.random() < 0.3:
+        frag = sampler.sample(rng.choice("0123456789"), rng)
+        fh = max(8, int(RENDER_DIGIT_H * rng.uniform(0.7, 1.1)))
+        fw = max(4, int(frag.shape[1] * fh / frag.shape[0]))
+        frag = cv2.resize(frag, (fw, fh), interpolation=cv2.INTER_LINEAR)
+        outside = rng.uniform(0.6, 0.9)
+        fx = -int(fw * outside) if rng.random() < 0.5 else cell_w - int(fw * (1 - outside))
+        _paste_clipped(alpha, frag, fx, int((cell_h - fh) * rng.uniform(0.2, 0.8)))
+
+    img = _compose_ink(_paper(rng, cell_h, cell_w), alpha, rng)
+    # printed grid borders: diagram crops nearly always contain them
+    for _ in range(rng.randint(1, 4) if rng.random() < 0.85 else 0):
+        img = _add_border_lines(img, rng)
+
+    out_h = rng.randint(20, 90)  # pit rows are small on the sheet
+    out_w = max(12, int(cell_w * out_h / cell_h))
+    img = cv2.resize(img, (out_w, out_h), interpolation=cv2.INTER_AREA)
+    return Image.fromarray(_camera_effects(img, rng))
+
+
+def _preview(path: str, seed: int, count: int, task: str = "moves") -> None:
     """Render a comparison grid: the real-crop classes first, then random ones."""
     sampler = GlyphSampler()
     rng = random.Random(seed)
-    fixed = ["96", "77", "13x", "37", "85", "42"] * 2
-    names = fixed + [rng.choice(CLASSES) for _ in range(max(0, count - len(fixed)))]
+    if task == "diagram":
+        fixed = ["x", "-", "0", "7", "12", "45", "81", "empty"]
+        pool, synth_fn = DIAGRAM_CLASSES, synthesize_diagram_cell
+    else:
+        fixed = ["96", "77", "13x", "37", "85", "42"] * 2
+        pool, synth_fn = CLASSES, synthesize_cell
+    names = fixed + [rng.choice(pool) for _ in range(max(0, count - len(fixed)))]
 
     tile_w, tile_h, caption = 200, 100, 14
     cols = 6
@@ -221,7 +325,7 @@ def _preview(path: str, seed: int, count: int) -> None:
 
     draw = ImageDraw.Draw(sheet)
     for i, name in enumerate(names):
-        cell = synthesize_cell(name, sampler, rng).resize((tile_w, tile_h))
+        cell = synth_fn(name, sampler, rng).resize((tile_w, tile_h))
         cx, cy = (i % cols) * tile_w, (i // cols) * (tile_h + caption)
         sheet.paste(cell, (cx, cy + caption))
         draw.text((cx + 4, cy + 1), name, fill=0)
@@ -234,5 +338,6 @@ if __name__ == "__main__":
     parser.add_argument("--preview", default="preview.png", help="output image path")
     parser.add_argument("--seed", type=int, default=0)
     parser.add_argument("--n", type=int, default=48, help="number of cells")
+    parser.add_argument("--task", choices=["moves", "diagram"], default="moves")
     args = parser.parse_args()
-    _preview(args.preview, args.seed, args.n)
+    _preview(args.preview, args.seed, args.n, args.task)

train.py

@@ -19,7 +19,7 @@ import torch
 import torch.nn as nn
 from torch.utils.data import DataLoader
 
-from togyz.classes import CLASSES, KAZAN_CLASSES
+from togyz.classes import CLASSES, DIAGRAM_CLASSES
 from togyz.dataset import RealCropDataset, SyntheticCellDataset
 from togyz.model import auto_device, build_model, save_checkpoint
 
@@ -52,9 +52,11 @@ def evaluate_real(model, real: RealCropDataset, device) -> tuple[float, float]:
 
 def main() -> None:
     parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--task", choices=["moves", "kazan"], default="moves",
-                        help="moves: 163-class cell classifier; kazan: 00-81 "
-                             "checkpoint numbers from the summary strips")
+    parser.add_argument("--task", choices=["moves", "diagram"], default="moves",
+                        help="moves: 163-class cell classifier; diagram: unified "
+                             "board-diagram reader (0-81, 'x', '-') for the kazan "
+                             "boxes and pit cells of the summary strips "
+                             "(replaces the old 'kazan' task)")
     parser.add_argument("--epochs", type=int, default=20)
     parser.add_argument("--samples-per-epoch", type=int, default=50_000)
     parser.add_argument("--val-size", type=int, default=4_000)
@@ -63,7 +65,7 @@ def main() -> None:
     parser.add_argument("--weight-decay", type=float, default=1e-4)
     parser.add_argument("--workers", type=int, default=4)
     parser.add_argument("--out", default=None,
-                        help="default: checkpoints (moves) / checkpoints/kazan")
+                        help="default: checkpoints (moves) / checkpoints/diagram")
     parser.add_argument("--resume", default=None, help="path to last.pt to continue")
     parser.add_argument("--device", default=None, help="cuda / mps / cpu (default: auto)")
     args = parser.parse_args()
@@ -71,13 +73,16 @@ def main() -> None:
     device = torch.device(args.device) if args.device else auto_device()
     print(f"Device: {device}, task: {args.task}")
 
-    classes = KAZAN_CLASSES if args.task == "kazan" else CLASSES
-    out_dir = Path(args.out or ("checkpoints/kazan" if args.task == "kazan" else "checkpoints"))
+    classes = DIAGRAM_CLASSES if args.task == "diagram" else CLASSES
+    out_dir = Path(args.out or ("checkpoints/diagram" if args.task == "diagram" else "checkpoints"))
     out_dir.mkdir(parents=True, exist_ok=True)
     (out_dir / "classes.json").write_text(json.dumps(classes))
 
-    train_set = SyntheticCellDataset(args.samples_per_epoch, seed=None, classes=classes)
-    val_set = SyntheticCellDataset(args.val_size, seed=1234, classes=classes)
+    train_set = SyntheticCellDataset(args.samples_per_epoch, seed=None,
+                                     classes=classes, task=args.task)
+    val_set = SyntheticCellDataset(args.val_size, seed=1234,
+                                   classes=classes, task=args.task)
+    # the labeled real crops are move cells; other tasks have no real eval set
     real_set = RealCropDataset() if args.task == "moves" else RealCropDataset("/nonexistent")
     print(f"Glyph pools:\n{train_set.sampler.describe()}")
     print(f"Real eval crops: {len(real_set)}")