| |
| """Resumable Surya detection on template-subtracted handwriting pages.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import fcntl |
| import io |
| import json |
| import os |
| import time |
| import traceback |
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
| import pyarrow.parquet as pq |
| from PIL import Image |
| from surya.detection import DetectionPredictor |
|
|
| from bina02_template_layers import TemplateCatalog |
|
|
|
|
| ROWS = 6_669 |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--rank", type=int, required=True) |
| parser.add_argument("--world-size", type=int, required=True) |
| parser.add_argument("--data", type=Path, required=True) |
| parser.add_argument("--background-root", type=Path, required=True) |
| parser.add_argument("--ocr-root", type=Path, required=True) |
| parser.add_argument("--output-dir", type=Path, required=True) |
| parser.add_argument("--variant", required=True) |
| parser.add_argument("--limit", type=int, default=ROWS) |
| parser.add_argument("--batch-size", type=int, default=8) |
| parser.add_argument("--progress-every", type=int, default=20) |
| return parser.parse_args() |
|
|
|
|
| def atomic_json(path: Path, payload: dict[str, Any]) -> None: |
| temporary = path.with_suffix(path.suffix + f".{os.getpid()}.tmp") |
| temporary.write_text( |
| json.dumps(payload, ensure_ascii=False, indent=2) + "\n", |
| encoding="utf-8", |
| ) |
| os.replace(temporary, path) |
|
|
|
|
| def selected_indices(limit: int) -> list[int]: |
| if not 1 <= limit <= ROWS: |
| raise ValueError(f"limit must be in [1, {ROWS}]") |
| if limit == ROWS: |
| return list(range(ROWS)) |
| return sorted( |
| { |
| int(round(value)) |
| for value in np.linspace(0, ROWS - 1, num=limit) |
| } |
| ) |
|
|
|
|
| def load_rows(data: Path, wanted: set[int]) -> dict[int, dict[str, Any]]: |
| rows: dict[int, dict[str, Any]] = {} |
| offset = 0 |
| paths = sorted((data / "data").glob("handwriting-*.parquet")) |
| if len(paths) != 14: |
| raise RuntimeError(f"expected 14 handwriting shards, found {len(paths)}") |
| for path in paths: |
| count = pq.ParquetFile(path).metadata.num_rows |
| local = sorted(index - offset for index in wanted if offset <= index < offset + count) |
| if local: |
| table = pq.read_table(path, columns=["image", "label"]) |
| for index in local: |
| rows[offset + index] = table.slice(index, 1).to_pylist()[0] |
| offset += count |
| if offset != ROWS or rows.keys() != wanted: |
| raise RuntimeError( |
| f"benchmark row contract failed: total={offset}, loaded={len(rows)}, wanted={len(wanted)}" |
| ) |
| return rows |
|
|
|
|
| def decode_image(row: dict[str, Any]) -> Image.Image: |
| value = row["image"] |
| encoded = value.get("bytes") if isinstance(value, dict) else value |
| if encoded is None and isinstance(value, dict) and value.get("path"): |
| encoded = Path(value["path"]).read_bytes() |
| if not encoded: |
| raise ValueError("dataset row has no image bytes") |
| return Image.open(io.BytesIO(encoded)).convert("RGB") |
|
|
|
|
| def line_payload(result: Any) -> list[dict[str, Any]]: |
| lines = [] |
| for item in getattr(result, "bboxes", []): |
| bbox = getattr(item, "bbox", None) |
| if bbox is None or len(bbox) != 4: |
| continue |
| polygon = getattr(item, "polygon", None) |
| lines.append( |
| { |
| "bbox": [float(value) for value in bbox], |
| "confidence": float(getattr(item, "confidence", 0.0)), |
| "polygon": ( |
| [[float(value) for value in point] for point in polygon] |
| if polygon is not None |
| else None |
| ), |
| } |
| ) |
| return lines |
|
|
|
|
| def load_done(path: Path) -> set[str]: |
| if not path.is_file(): |
| return set() |
| with path.open(encoding="utf-8") as handle: |
| return { |
| str(row["id"]) |
| for line in handle |
| if line.strip() |
| for row in [json.loads(line)] |
| } |
|
|
|
|
| def merge_if_complete( |
| output_dir: Path, |
| world_size: int, |
| selected: list[int], |
| variant: str, |
| ) -> None: |
| with (output_dir / "merge.lock").open("a+") as lock: |
| fcntl.flock(lock, fcntl.LOCK_EX) |
| statuses = [] |
| for rank in range(world_size): |
| path = output_dir / f"status-rank{rank:02d}.json" |
| if not path.is_file(): |
| return |
| value = json.loads(path.read_text(encoding="utf-8")) |
| if value.get("state") != "complete": |
| return |
| statuses.append(value) |
| rows: dict[str, dict[str, Any]] = {} |
| duplicates = [] |
| for rank in range(world_size): |
| path = output_dir / f"detections-rank{rank:02d}.jsonl" |
| with path.open(encoding="utf-8") as handle: |
| for line in handle: |
| row = json.loads(line) |
| row_id = str(row["id"]) |
| if row_id in rows: |
| duplicates.append(row_id) |
| rows[row_id] = row |
| expected = [f"handwriting:{index}" for index in selected] |
| missing = sorted(set(expected) - rows.keys()) |
| extras = sorted(rows.keys() - set(expected)) |
| if duplicates or missing or extras: |
| raise RuntimeError( |
| f"detection merge failed: duplicates={duplicates[:3]} missing={missing[:3]} extras={extras[:3]}" |
| ) |
| temporary = output_dir / "detections.jsonl.tmp" |
| with temporary.open("w", encoding="utf-8") as handle: |
| for row_id in expected: |
| handle.write(json.dumps(rows[row_id], ensure_ascii=False) + "\n") |
| os.replace(temporary, output_dir / "detections.jsonl") |
| atomic_json( |
| output_dir / "run-summary.json", |
| { |
| "state": "complete", |
| "variant": variant, |
| "selected_indices": selected, |
| "rows": len(rows), |
| "failed_pages": sum(row["failures"] for row in statuses), |
| "template_match_failures": sum( |
| row["template_match_failures"] for row in statuses |
| ), |
| "detected_lines": sum(row["detected_lines"] for row in statuses), |
| "completed_unix": time.time(), |
| }, |
| ) |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| selected = selected_indices(args.limit) |
| assigned = selected[args.rank :: args.world_size] |
| output_path = args.output_dir / f"detections-rank{args.rank:02d}.jsonl" |
| status_path = args.output_dir / f"status-rank{args.rank:02d}.json" |
| args.output_dir.mkdir(parents=True, exist_ok=True) |
| done = load_done(output_path) |
| wanted = { |
| index for index in assigned if f"handwriting:{index}" not in done |
| } |
| rows = load_rows(args.data, wanted) if wanted else {} |
| catalog = TemplateCatalog(args.background_root, args.ocr_root) |
| predictor = DetectionPredictor.local(device="cuda") |
| started = time.monotonic() |
| completed = 0 |
| failures = 0 |
| match_failures = 0 |
| detected_lines = 0 |
| pending: list[tuple[int, dict[str, Any], Image.Image, dict[str, Any]]] = [] |
|
|
| with output_path.open("a", encoding="utf-8", buffering=1) as output: |
| def flush() -> None: |
| nonlocal completed, failures, detected_lines |
| if not pending: |
| return |
| images = [item[2] for item in pending] |
| try: |
| results: list[Any | Exception] = list(predictor(images)) |
| if len(results) != len(images): |
| raise RuntimeError("detector returned wrong result count") |
| except Exception: |
| traceback.print_exc() |
| results = [] |
| for image in images: |
| try: |
| results.append(list(predictor([image]))[0]) |
| except Exception as exc: |
| results.append(exc) |
| for (index, row, image, layer), result in zip(pending, results): |
| lines = [] |
| error = None |
| if isinstance(result, Exception): |
| failures += 1 |
| error = f"{type(result).__name__}: {result}" |
| else: |
| lines = line_payload(result) |
| detected_lines += len(lines) |
| payload = { |
| "id": f"handwriting:{index}", |
| "split": "handwriting", |
| "reference": str(row.get("label") or ""), |
| "width": image.width, |
| "height": image.height, |
| "variant": args.variant, |
| "lines": lines, |
| **layer, |
| } |
| if error: |
| payload["error"] = error |
| output.write(json.dumps(payload, ensure_ascii=False) + "\n") |
| completed += 1 |
| pending.clear() |
|
|
| for index in assigned: |
| if index not in wanted: |
| continue |
| row = rows[index] |
| try: |
| image = decode_image(row) |
| isolated, match, metrics = catalog.isolate(image, args.variant) |
| layer = { |
| "template_id": match.template.id, |
| "template_file": match.template.file, |
| "template_text": match.template.template_text, |
| "template_placeholder_count": match.template.template_text.count( |
| "{{HANDWRITING}}" |
| ), |
| "match_median_abs_delta": match.median_abs_delta, |
| "match_p90_abs_delta": match.p90_abs_delta, |
| "mask_threshold": metrics["threshold"], |
| "mask_kept_fraction": metrics["kept_fraction"], |
| } |
| pending.append((index, row, isolated, layer)) |
| if len(pending) >= args.batch_size: |
| flush() |
| except Exception as exc: |
| traceback.print_exc() |
| failures += 1 |
| match_failures += 1 |
| completed += 1 |
| output.write( |
| json.dumps( |
| { |
| "id": f"handwriting:{index}", |
| "split": "handwriting", |
| "reference": str(row.get("label") or ""), |
| "width": 0, |
| "height": 0, |
| "variant": args.variant, |
| "lines": [], |
| "error": f"{type(exc).__name__}: {exc}", |
| }, |
| ensure_ascii=False, |
| ) |
| + "\n" |
| ) |
| flush() |
|
|
| elapsed = max(time.monotonic() - started, 1e-9) |
| atomic_json( |
| status_path, |
| { |
| "state": "complete", |
| "rank": args.rank, |
| "assigned": len(assigned), |
| "completed": completed, |
| "failures": failures, |
| "template_match_failures": match_failures, |
| "detected_lines": detected_lines, |
| "rate_pages_per_second": completed / elapsed, |
| "eta_seconds": 0, |
| "updated_unix": time.time(), |
| }, |
| ) |
| merge_if_complete(args.output_dir, args.world_size, selected, args.variant) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
