| |
| """Recognize cached Surya boxes from the same isolated handwriting layer.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import fcntl |
| import io |
| import json |
| import os |
| import re |
| import time |
| import traceback |
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
| import pyarrow.parquet as pq |
| from paddleocr import TextRecognition |
| from PIL import Image |
|
|
| from bina02_template_layers import TemplateCatalog, index_from_id |
|
|
|
|
| ROWS = 6_669 |
| LTR_RUN = re.compile(r"[a-zA-Z0-9 :*./%+-]") |
|
|
|
|
| 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("--detections-dir", type=Path, required=True) |
| parser.add_argument("--model-dir", type=Path, required=True) |
| parser.add_argument("--output-dir", type=Path, required=True) |
| parser.add_argument("--variant", required=True) |
| parser.add_argument("--device", default="gpu:0") |
| parser.add_argument("--crop-buffer", type=int, default=512) |
| parser.add_argument("--recognition-batch-size", type=int, default=64) |
| parser.add_argument("--pad-x-ratio", type=float, default=0.01) |
| parser.add_argument("--pad-y-ratio", type=float, default=0.08) |
| parser.add_argument("--row-factor", type=float, default=0.70) |
| 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 pred_reverse(text: str) -> str: |
| segments: list[str] = [] |
| current = "" |
| for character in text: |
| if LTR_RUN.search(character): |
| current += character |
| else: |
| if current: |
| segments.append(current) |
| current = "" |
| segments.append(character) |
| if current: |
| segments.append(current) |
| return "".join(reversed(segments)) |
|
|
|
|
| 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")) |
| 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("benchmark row load contract failed") |
| 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 center(line: dict[str, Any]) -> tuple[float, float]: |
| x0, y0, x1, y1 = (float(value) for value in line["bbox"]) |
| return (x0 + x1) / 2.0, (y0 + y1) / 2.0 |
|
|
|
|
| def height(line: dict[str, Any]) -> float: |
| return max(float(line["bbox"][3]) - float(line["bbox"][1]), 1.0) |
|
|
|
|
| def group_rows(lines: list[dict[str, Any]], factor: float) -> list[list[dict[str, Any]]]: |
| ordered = sorted(lines, key=lambda line: center(line)[1]) |
| rows: list[list[dict[str, Any]]] = [] |
| for line in ordered: |
| if not rows: |
| rows.append([line]) |
| continue |
| previous = rows[-1] |
| mean_y = sum(center(item)[1] for item in previous) / len(previous) |
| mean_height = sum(height(item) for item in previous) / len(previous) |
| if abs(center(line)[1] - mean_y) <= factor * max(height(line), mean_height): |
| previous.append(line) |
| else: |
| rows.append([line]) |
| for row in rows: |
| row.sort(key=lambda line: center(line)[0], reverse=True) |
| return rows |
|
|
|
|
| def grouped_lines( |
| lines: list[dict[str, Any]], |
| width: int, |
| height_value: int, |
| factor: float, |
| ) -> list[dict[str, Any]]: |
| aspect = width / max(height_value, 1) |
| if 1.03 <= aspect <= 1.42: |
| midpoint = width / 2.0 |
| groups = group_rows( |
| [line for line in lines if center(line)[0] >= midpoint], |
| factor, |
| ) + group_rows( |
| [line for line in lines if center(line)[0] < midpoint], |
| factor, |
| ) |
| else: |
| groups = group_rows(lines, factor) |
| return [ |
| { |
| "bbox": [ |
| min(float(line["bbox"][0]) for line in row), |
| min(float(line["bbox"][1]) for line in row), |
| max(float(line["bbox"][2]) for line in row), |
| max(float(line["bbox"][3]) for line in row), |
| ], |
| "confidence": sum(float(line.get("confidence", 0.0)) for line in row) |
| / len(row), |
| "member_count": len(row), |
| } |
| for row in groups |
| ] |
|
|
|
|
| def crop( |
| image: Image.Image, |
| bbox: list[float], |
| pad_x_ratio: float, |
| pad_y_ratio: float, |
| ) -> np.ndarray: |
| x0, y0, x1, y1 = (float(value) for value in bbox) |
| width = max(x1 - x0, 1.0) |
| height_value = max(y1 - y0, 1.0) |
| pad_x = max(2, int(width * pad_x_ratio)) |
| pad_y = max(2, int(height_value * pad_y_ratio)) |
| box = ( |
| max(0, int(x0) - pad_x), |
| max(0, int(y0) - pad_y), |
| min(image.width, int(x1) + pad_x), |
| min(image.height, int(y1) + pad_y), |
| ) |
| if box[2] <= box[0] or box[3] <= box[1]: |
| raise ValueError(f"invalid crop: {bbox}") |
| return np.asarray(image.crop(box).convert("RGB")) |
|
|
|
|
| def recognize( |
| model: TextRecognition, |
| crops: list[np.ndarray], |
| batch_size: int, |
| ) -> list[dict[str, Any]]: |
| if not crops: |
| return [] |
| last_error: Exception | None = None |
| for attempt in range(1, 4): |
| try: |
| results = list(model.predict(input=crops, batch_size=batch_size)) |
| if len(results) != len(crops): |
| raise RuntimeError("recognizer returned wrong result count") |
| output = [] |
| for result in results: |
| payload = result.json |
| if callable(payload): |
| payload = payload() |
| payload = payload.get("res", payload) |
| raw = str(payload.get("rec_text", "")) |
| output.append( |
| { |
| "text": pred_reverse(raw), |
| "raw_text": raw, |
| "rec_score": float(payload.get("rec_score", 0.0)), |
| } |
| ) |
| return output |
| except Exception as exc: |
| last_error = exc |
| traceback.print_exc() |
| time.sleep(attempt) |
| raise RuntimeError(f"recognition failed after retries: {last_error}") |
|
|
|
|
| def merge_if_complete(output_dir: Path, detections_dir: Path, world_size: int) -> None: |
| with (output_dir / "merge.lock").open("a+") as lock: |
| fcntl.flock(lock, fcntl.LOCK_EX) |
| for rank in range(world_size): |
| path = output_dir / f"status-rank{rank:02d}.json" |
| if not path.is_file() or json.loads(path.read_text()).get("state") != "complete": |
| return |
| expected = [ |
| json.loads(line)["id"] |
| for line in (detections_dir / "detections.jsonl").read_text( |
| encoding="utf-8" |
| ).splitlines() |
| if line.strip() |
| ] |
| rows = {} |
| for rank in range(world_size): |
| path = output_dir / f"recognized-rank{rank:02d}.jsonl" |
| for line in path.read_text(encoding="utf-8").splitlines(): |
| if line.strip(): |
| row = json.loads(line) |
| rows[row["id"]] = row |
| if set(rows) != set(expected): |
| raise RuntimeError("recognition merge contract failed") |
| temporary = output_dir / "recognized-lines.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 / "recognized-lines.jsonl") |
| atomic_json( |
| output_dir / "run-summary.json", |
| {"state": "complete", "rows": len(rows), "completed_unix": time.time()}, |
| ) |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| args.output_dir.mkdir(parents=True, exist_ok=True) |
| detections_path = args.detections_dir / f"detections-rank{args.rank:02d}.jsonl" |
| detections = [ |
| json.loads(line) |
| for line in detections_path.read_text(encoding="utf-8").splitlines() |
| if line.strip() |
| ] |
| indices = {index_from_id(row["id"]) for row in detections} |
| rows = load_rows(args.data, indices) |
| catalog = TemplateCatalog(args.background_root, args.ocr_root) |
| model = TextRecognition(model_dir=str(args.model_dir), device=args.device) |
| output_path = args.output_dir / f"recognized-rank{args.rank:02d}.jsonl" |
| status_path = args.output_dir / f"status-rank{args.rank:02d}.json" |
| started = time.monotonic() |
| failures = 0 |
| recognized_lines = 0 |
| page_buffer: list[dict[str, Any]] = [] |
| crop_buffer: list[np.ndarray] = [] |
|
|
| with output_path.open("w", encoding="utf-8", buffering=1) as output: |
| def flush() -> None: |
| nonlocal failures, recognized_lines |
| if not page_buffer: |
| return |
| try: |
| values = recognize(model, crop_buffer, args.recognition_batch_size) |
| recognition_error = None |
| except Exception as exc: |
| values = [ |
| {"text": "", "raw_text": "", "rec_score": 0.0} |
| for _ in crop_buffer |
| ] |
| recognition_error = f"{type(exc).__name__}: {exc}" |
| for page in page_buffer: |
| if recognition_error and not page.get("error"): |
| page["error"] = recognition_error |
| if page.get("error"): |
| failures += 1 |
| for line in page["lines"]: |
| result_index = line.pop("_result_index", None) |
| if result_index is not None: |
| line.update(values[result_index]) |
| recognized_lines += 1 |
| output.write(json.dumps(page, ensure_ascii=False) + "\n") |
| page_buffer.clear() |
| crop_buffer.clear() |
|
|
| for detection in detections: |
| index = index_from_id(detection["id"]) |
| page = { |
| key: detection.get(key) |
| for key in ( |
| "id", |
| "reference", |
| "width", |
| "height", |
| "variant", |
| "template_id", |
| "template_file", |
| "template_text", |
| "template_placeholder_count", |
| "match_median_abs_delta", |
| "match_p90_abs_delta", |
| "mask_threshold", |
| "mask_kept_fraction", |
| ) |
| } |
| page["lines"] = [] |
| if detection.get("error"): |
| page["error"] = detection["error"] |
| try: |
| image = decode_image(rows[index]) |
| template = catalog.by_id[str(detection["template_id"])] |
| isolated, _ = catalog.isolate_with_template( |
| image, |
| template, |
| args.variant, |
| ) |
| for source in grouped_lines( |
| detection.get("lines", []), |
| isolated.width, |
| isolated.height, |
| args.row_factor, |
| ): |
| line = dict(source) |
| line["_result_index"] = len(crop_buffer) |
| crop_buffer.append( |
| crop( |
| isolated, |
| source["bbox"], |
| args.pad_x_ratio, |
| args.pad_y_ratio, |
| ) |
| ) |
| page["lines"].append(line) |
| except Exception as exc: |
| traceback.print_exc() |
| page["error"] = f"{type(exc).__name__}: {exc}" |
| page_buffer.append(page) |
| if len(crop_buffer) >= args.crop_buffer: |
| flush() |
| flush() |
|
|
| elapsed = max(time.monotonic() - started, 1e-9) |
| atomic_json( |
| status_path, |
| { |
| "state": "complete", |
| "rank": args.rank, |
| "assigned": len(detections), |
| "failures": failures, |
| "recognized_lines": recognized_lines, |
| "rate_pages_per_second": len(detections) / elapsed, |
| "eta_seconds": 0, |
| }, |
| ) |
| merge_if_complete(args.output_dir, args.detections_dir, args.world_size) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
