From 1d61fbe5359678744bda7bf25c83a32203f03ff1 Mon Sep 17 00:00:00 2001 From: 732642856 <732642856@qq.com> Date: Fri, 17 Jul 2026 00:27:57 +0800 Subject: [PATCH] Add OCR transcript extraction helper --- scripts/ocr_extract.py | 151 ++++++++++++++++++++++++++++++++++++++ tests/test_ocr_extract.py | 51 +++++++++++++ 2 files changed, 202 insertions(+) create mode 100644 scripts/ocr_extract.py create mode 100644 tests/test_ocr_extract.py diff --git a/scripts/ocr_extract.py b/scripts/ocr_extract.py new file mode 100644 index 00000000..b71b4334 --- /dev/null +++ b/scripts/ocr_extract.py @@ -0,0 +1,151 @@ +#!/usr/bin/env python3 +"""Extract text from screenshots without requiring Homebrew-installed Tesseract.""" + +from __future__ import annotations + +import argparse +import json +import shutil +import subprocess +import sys +from pathlib import Path +from typing import Any + + +DEFAULT_SHORTCUT_NAME = "Extract Text from Image" +VALID_BACKENDS = {"auto", "manual", "shortcuts", "tesseract"} + + +def choose_backend(requested: str = "auto") -> str: + if requested != "auto": + if requested not in VALID_BACKENDS: + raise ValueError(f"unsupported backend: {requested}") + return requested + if shutil.which("shortcuts"): + return "shortcuts" + if shutil.which("tesseract"): + return "tesseract" + return "manual" + + +def _manual_transcript_path(image: Path, transcript_dir: Path | None) -> Path: + base = transcript_dir or image.parent + return base / f"{image.stem}.txt" + + +def _extract_manual(image: Path, transcript_dir: Path | None) -> dict[str, Any]: + transcript = _manual_transcript_path(image, transcript_dir) + if not transcript.is_file(): + return { + "image_path": str(image), + "text": "", + "backend": "manual", + "status": "blocked", + "reason": "manual_transcript_missing", + "expected_transcript_path": str(transcript), + } + return { + "image_path": str(image), + "text": transcript.read_text(encoding="utf-8"), + "backend": "manual", + "status": "ok", + } + + +def _extract_shortcuts(image: Path, shortcut_name: str) -> dict[str, Any]: + if not shutil.which("shortcuts"): + return {"image_path": str(image), "text": "", "backend": "shortcuts", "status": "blocked", "reason": "shortcuts_cli_missing"} + completed = subprocess.run( + ["shortcuts", "run", shortcut_name, "-i", str(image)], + capture_output=True, + text=True, + check=False, + timeout=120, + ) + text = completed.stdout + if completed.returncode != 0: + return { + "image_path": str(image), + "text": text, + "backend": "shortcuts", + "status": "blocked", + "reason": "shortcuts_run_failed", + "stderr": completed.stderr.strip(), + "shortcut_name": shortcut_name, + } + return {"image_path": str(image), "text": text, "backend": "shortcuts", "status": "ok"} + + +def _extract_tesseract(image: Path) -> dict[str, Any]: + if not shutil.which("tesseract"): + return {"image_path": str(image), "text": "", "backend": "tesseract", "status": "blocked", "reason": "tesseract_missing"} + completed = subprocess.run( + ["tesseract", str(image), "stdout", "-l", "eng+chi_sim"], + capture_output=True, + text=True, + check=False, + timeout=120, + ) + if completed.returncode != 0: + return { + "image_path": str(image), + "text": completed.stdout, + "backend": "tesseract", + "status": "blocked", + "reason": "tesseract_run_failed", + "stderr": completed.stderr.strip(), + } + return {"image_path": str(image), "text": completed.stdout, "backend": "tesseract", "status": "ok"} + + +def extract_one(image: Path, *, backend: str = "auto", transcript_dir: Path | None = None, shortcut_name: str = DEFAULT_SHORTCUT_NAME) -> dict[str, Any]: + selected = choose_backend(backend) + if selected == "manual": + return _extract_manual(image, transcript_dir) + if selected == "shortcuts": + return _extract_shortcuts(image, shortcut_name) + if selected == "tesseract": + return _extract_tesseract(image) + raise ValueError(f"unsupported backend: {selected}") + + +def extract_many( + images: list[Path], + *, + output: Path | None = None, + backend: str = "auto", + transcript_dir: Path | None = None, + shortcut_name: str = DEFAULT_SHORTCUT_NAME, +) -> dict[str, Any]: + items = [extract_one(image, backend=backend, transcript_dir=transcript_dir, shortcut_name=shortcut_name) for image in images] + if output: + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text("\n".join(json.dumps(item, ensure_ascii=False, sort_keys=True) for item in items) + "\n", encoding="utf-8") + return { + "status": "ok" if items and all(item["status"] == "ok" for item in items) else "blocked", + "backend": choose_backend(backend), + "items": items, + } + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("images", nargs="+", type=Path) + parser.add_argument("--backend", choices=sorted(VALID_BACKENDS), default="auto") + parser.add_argument("--transcript-dir", type=Path) + parser.add_argument("--shortcut-name", default=DEFAULT_SHORTCUT_NAME) + parser.add_argument("--output", type=Path, default=Path("scratch/local/ocr_extract/ocr.jsonl")) + args = parser.parse_args(argv) + report = extract_many( + args.images, + output=args.output, + backend=args.backend, + transcript_dir=args.transcript_dir, + shortcut_name=args.shortcut_name, + ) + print(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True)) + return 0 if report["status"] == "ok" else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_ocr_extract.py b/tests/test_ocr_extract.py new file mode 100644 index 00000000..c48a5f6b --- /dev/null +++ b/tests/test_ocr_extract.py @@ -0,0 +1,51 @@ +from __future__ import annotations + +import json +import sys +from pathlib import Path + + +ROOT = Path(__file__).resolve().parents[1] +SCRIPTS = ROOT / "scripts" +if str(SCRIPTS) not in sys.path: + sys.path.insert(0, str(SCRIPTS)) + +import ocr_extract # noqa: E402 + + +def test_manual_transcript_backend_writes_jsonl(tmp_path: Path) -> None: + image = tmp_path / "IMG_3502.PNG" + image.write_bytes(b"not-a-real-image") + transcript = tmp_path / "IMG_3502.txt" + transcript.write_text("Arudha Lagna\nUpapada\n", encoding="utf-8") + output = tmp_path / "ocr.jsonl" + + report = ocr_extract.extract_many([image], output=output, transcript_dir=tmp_path, backend="manual") + + assert report["status"] == "ok" + assert report["items"][0]["backend"] == "manual" + rows = [json.loads(line) for line in output.read_text(encoding="utf-8").splitlines()] + assert rows == [ + { + "image_path": str(image), + "text": "Arudha Lagna\nUpapada\n", + "backend": "manual", + "status": "ok", + } + ] + + +def test_missing_manual_transcript_is_blocked(tmp_path: Path) -> None: + image = tmp_path / "IMG_3503.PNG" + image.write_bytes(b"not-a-real-image") + + report = ocr_extract.extract_many([image], transcript_dir=tmp_path, backend="manual") + + assert report["status"] == "blocked" + assert report["items"][0]["reason"] == "manual_transcript_missing" + + +def test_backend_auto_prefers_shortcuts_before_tesseract(monkeypatch) -> None: + monkeypatch.setattr(ocr_extract.shutil, "which", lambda name: f"/usr/bin/{name}" if name in {"shortcuts", "tesseract"} else None) + + assert ocr_extract.choose_backend("auto") == "shortcuts"