Files
pf-ocr/ocr/manga_ocr_llm.py
T
2026-09-18 01:00:58 +03:00

470 lines
21 KiB
Python

#!/usr/bin/env python3
"""
Batch-OCR (or batch-translate) for scanned manga pages via any
OpenAI-compatible API (OpenRouter, a direct provider endpoint, a
self-hosted proxy, etc.) — including Gemini, GPT-4V-class models, or
anything else exposed through such an endpoint.
Manga pages split their text across scattered speech bubbles, thought
bubbles, narration boxes, and SFX rather than running prose, so this script
has genuinely different logic from novel_ocr.py: it transcribes each page
as a numbered, typed, reading-order list of bubbles. The --translate
machinery (glossary, sliding-window context) is shared infrastructure from
ocr/common.py — the same helpers novel_ocr.py uses for its own --translate.
Usage (after filling in config.json):
python manga_ocr_llm.py --input ./pages --output ./out
# translate instead of transcribing:
python manga_ocr_llm.py --input ./pages --output ./out \\
--translate --target-lang Russian
# ...with a glossary of established name/term translations:
python manga_ocr_llm.py --input ./pages --output ./out \\
--translate --target-lang Russian --glossary ../ocr/glossary.md
# (glossary.md sitting next to this script is picked up automatically
# even without --glossary — see ocr/glossary.example.md for the format)
# ...with more/less continuity context from previous pages (default: 2):
python manga_ocr_llm.py --input ./pages --output ./out \\
--translate --target-lang Russian --context-pages 4
# ...with lookahead context from upcoming (untranslated) pages, in the
# original Japanese — helps with twists, who's-talking-to-whom, jokes
# that pay off a page later:
python manga_ocr_llm.py --input ./pages --output ./out \\
--translate --target-lang Russian --context-pages-ahead 1
# re-OCR/re-translate just a few pages (e.g. after editing the prompt
# or the glossary), without redoing the whole book:
python manga_ocr_llm.py --input ./pages --output ./out \\
--translate --pages "12,40-42" --redo
Requirements:
pip install openai pillow natsort tqdm
One-time setup:
1. Copy config.example.json (repo root) -> config.json
2. Fill in the "ocr" section: api_key, base_url, model (temperature/
max_tokens/top_p/reasoning_effort are optional — see config.example.json)
3. Optionally edit prompt_manga.txt (transcription) or
prompt_manga_translate.txt (--translate) to fit your book / house style
System/user prompt split: the format instructions (and, with --translate,
the glossary) are sent as the `system` message, since they're identical on
every page of a run — most OpenAI-compatible providers cache a repeated
system prefix, so a long book isn't re-billed/re-processed for the same
instructions on every single page. Only the image and any sliding-window
context (which does vary page to page) go in `user`. See ocr/common.py's
module docstring for more.
Every request/response is logged as one JSON file under logs/ at the repo
root (--log-dir to change, --no-log to disable) — full prompt, model
params, every retry attempt, and the full raw API response (the base64
image itself is never included, only its path/size).
Progress/resume state: like novel_ocr.py, this writes a state.json
manifest to --output (model, prompt hash, usage, timestamps) so pages
transcribed under a now-stale prompt/glossary can be spotted, and so
--pages/--redo can target just the pages that actually need re-running.
"""
import argparse
import os
import sys
import time
from collections import deque
from datetime import datetime
from pathlib import Path
from natsort import natsorted
from openai import OpenAI
from PIL import Image
from tqdm import tqdm
SCRIPT_DIR = Path(__file__).resolve().parent
ROOT_DIR = SCRIPT_DIR.parent
sys.path.insert(0, str(SCRIPT_DIR)) # so `import common` / `from local_mangaocr_ocr import ...` work when run from elsewhere
import common
IMAGE_EXTS = common.IMAGE_EXTS
def get_japanese_context(
img_path: Path, context_src_dir: Path,
ahead_ocr_backend: str, ahead_ocr_state: dict,
) -> str:
"""Resolves the raw Japanese transcription of a page used as lookahead
context, in priority order:
1. A .txt already sitting in context_src_dir with the same stem —
hand-prepared, or cached from a previous run/page. Used as-is.
2. Otherwise, OCR it now with the chosen backend (--ahead-ocr-backend)
and cache the result into context_src_dir, so the next page that
needs this same page as context doesn't re-OCR it.
Returns "" (and prints a warning) if OCR isn't available/fails — a
missing bit of lookahead context isn't worth failing the whole page
over, it just means less context than requested.
"""
cached_path = context_src_dir / f"{img_path.stem}.txt"
if cached_path.exists():
return cached_path.read_text(encoding="utf-8").strip()
try:
if ahead_ocr_backend == "local":
if "mocr" not in ahead_ocr_state:
from local_mangaocr_ocr import ocr_page_manga
from manga_ocr import MangaOcr
print("Loading local manga-ocr model for lookahead context (first use only)...")
ahead_ocr_state["mocr"] = MangaOcr()
ahead_ocr_state["ocr_page_manga"] = ocr_page_manga
pil_img = Image.open(img_path)
text = ahead_ocr_state["ocr_page_manga"](ahead_ocr_state["mocr"], pil_img)
else: # "llm"
client = ahead_ocr_state["client"]
model = ahead_ocr_state["model"]
request_kwargs = ahead_ocr_state["request_kwargs"]
extra_body = ahead_ocr_state["extra_body"]
system_prompt = ahead_ocr_state["system_prompt"]
text, _meta = common.ocr_image(client, model, system_prompt, img_path, request_kwargs, extra_body)
except Exception as e: # noqa: BLE001
print(f"\nWarning: couldn't OCR {img_path.name} for lookahead context: {e}", file=sys.stderr)
return ""
text = text.strip()
context_src_dir.mkdir(parents=True, exist_ok=True)
cached_path.write_text(text, encoding="utf-8")
return text
def main():
parser = argparse.ArgumentParser(description="Batch-OCR/translate manga pages via an OpenAI-compatible API")
parser.add_argument("--input", required=True, help="Folder with scanned page images (jpg/png)")
parser.add_argument("--output", required=True, help="Folder for the OCR results")
parser.add_argument(
"--config", default=str(ROOT_DIR / "config.json"),
help="Path to config.json with an \"ocr\" section (api_key/base_url/model/...). "
"Defaults to config.json at the repo root."
)
parser.add_argument(
"--translate", action="store_true",
help="Translate the page instead of transcribing the Japanese "
"(uses prompt_manga_translate.txt by default)"
)
parser.add_argument(
"--target-lang", default="English",
help="Target language for --translate, written out as a plain name "
"(e.g. Russian, English, Spanish). Default: English"
)
parser.add_argument(
"--glossary", default=None,
help="[--translate only] Path to a glossary.md of established name/term "
"translations, sent along with the prompt for consistency (see "
"ocr/glossary.example.md). Auto-detected as glossary.md next to "
"this script if not given."
)
parser.add_argument(
"--context-pages", type=int, default=2,
help="[--translate only] Include the last N translated pages in the "
"prompt for continuity (character voice, ongoing dialogue, "
"pronouns/referents). 0 disables it. Default: 2"
)
parser.add_argument(
"--context-pages-ahead", type=int, default=0,
help="[--translate only] Also include the ORIGINAL JAPANESE (never a "
"draft translation — see ocr/README.md) of the next N pages, "
"for context that depends on what happens next (a twist, "
"who's actually being addressed, a joke that pays off later). "
"0 (default) disables it. Before translating a page, this "
"makes sure the next N pages' Japanese text is available "
"(OCR'ing it now if there's no cached/hand-prepared .txt yet), "
"so pages needing lookahead take a bit longer."
)
parser.add_argument(
"--ahead-ocr-backend", choices=["llm", "local"], default="llm",
help="[--context-pages-ahead only] How to OCR an upcoming page's "
"Japanese text when no cached/hand-prepared .txt is found for "
"it yet. 'llm': the same OpenAI-compatible API as the main "
"translation (transcription prompt, no translation). 'local': "
"manga-ocr running fully offline (requires the manga-ocr and "
"opencv-python packages). Default: llm"
)
parser.add_argument(
"--context-src-dir", default=None,
help="[--context-pages-ahead only] Folder holding the upcoming pages' "
"original-Japanese .txt files, named exactly like the page image "
"stems (e.g. page005.txt for page005.jpg) — hand-prepared ones "
"are used as-is; OCR'd-on-the-fly ones are cached here too. "
"Defaults to <o>/context_src/"
)
parser.add_argument(
"--prompt-file", default=None,
help="Path to the prompt file (defaults to prompt_manga.txt, or "
"prompt_manga_translate.txt with --translate)"
)
parser.add_argument("--base-url", default=None, help="Override base_url from config.json")
parser.add_argument("--model", default=None, help="Override model from config.json")
parser.add_argument("--api-key", default=None, help="Override api_key from config.json")
parser.add_argument(
"--temperature", type=float, default=None,
help="Override temperature from config.json (default if unset anywhere: 0)"
)
parser.add_argument("--max-tokens", type=int, default=None, help="Override max_tokens from config.json")
parser.add_argument("--top-p", type=float, default=None, help="Override top_p from config.json")
parser.add_argument(
"--reasoning-effort", default=None,
help="Override reasoning_effort from config.json (e.g. low/medium/high — support "
"depends on the model/provider; omitted from the request unless set)"
)
parser.add_argument(
"--start-page", type=int, default=1,
help="1-based index of the first page to process (pages before it are skipped entirely)"
)
parser.add_argument(
"--pages", default=None,
help="Only process these pages, e.g. \"3,7-12,20\" (1-based positions in the sorted page "
"list). Combine with --redo to re-OCR/re-translate just those pages."
)
parser.add_argument(
"--redo", action="store_true",
help="Re-OCR/re-translate pages even if output already exists (respects --pages if given)"
)
parser.add_argument(
"--no-state", action="store_true",
help="Ignore state.json and fall back to plain \"file exists = done\" behaviour"
)
parser.add_argument(
"--sleep", type=float, default=0.0,
help="Delay in seconds between requests (useful if you're hitting rate limits)"
)
parser.add_argument(
"--log-dir", default=str(ROOT_DIR / "logs"),
help="Folder for per-page request/response logs (one JSON file per page). "
"Defaults to logs/ at the repo root."
)
parser.add_argument(
"--no-log", action="store_true",
help="Disable request/response logging entirely"
)
args = parser.parse_args()
config = common.load_config_section(Path(args.config), "ocr")
api_key = args.api_key or config.get("api_key") or os.environ.get("API_KEY")
base_url = args.base_url or config.get("base_url")
model = args.model or config.get("model")
missing = [name for name, val in [("api_key", api_key), ("base_url", base_url), ("model", model)] if not val]
if missing:
print(
f"Missing settings: {', '.join(missing)}. "
f"Fill them in {args.config} (see config.example.json) or pass --api-key/--base-url/--model.",
file=sys.stderr,
)
sys.exit(1)
request_kwargs, extra_body = common.build_request_kwargs(config, args)
if args.translate:
default_prompt_name = "prompt_manga_translate.txt"
target_language = args.target_lang
glossary_path = common.find_glossary(SCRIPT_DIR, args.glossary)
glossary = glossary_path.read_text(encoding="utf-8").strip() if glossary_path else None
if glossary_path:
print(f"Using glossary: {glossary_path}")
else:
default_prompt_name = "prompt_manga.txt"
target_language = None
glossary = None
prompt_file = Path(args.prompt_file) if args.prompt_file else SCRIPT_DIR / default_prompt_name
base_prompt = common.load_prompt_file(prompt_file)
system_prompt = common.build_system_prompt(base_prompt, target_language, glossary)
prompt_hash = common.prompt_fingerprint(system_prompt)
# Plain transcription system prompt for OCR'ing lookahead pages — never
# the translate prompt (we want the raw Japanese, not a draft
# translation — see the design note in common.build_context_block).
ahead_system_prompt = (
system_prompt if not args.translate
else common.build_system_prompt(common.load_prompt_file(SCRIPT_DIR / "prompt_manga.txt"))
)
use_ahead_context = args.translate and args.context_pages_ahead > 0
ahead_ocr_state = {}
if use_ahead_context and args.ahead_ocr_backend == "llm":
ahead_ocr_state["system_prompt"] = ahead_system_prompt
ahead_ocr_state["request_kwargs"] = request_kwargs
ahead_ocr_state["extra_body"] = extra_body
# client/model are filled in below, once the real client exists.
input_dir = Path(args.input)
output_dir = Path(args.output)
output_dir.mkdir(parents=True, exist_ok=True)
pages_dir = output_dir / "pages_txt"
pages_dir.mkdir(exist_ok=True)
images = [p for p in input_dir.iterdir() if p.suffix.lower() in IMAGE_EXTS]
images = natsorted(images, key=lambda p: p.name)
if not images:
print(f"No images found in {input_dir}.", file=sys.stderr)
sys.exit(1)
if args.start_page < 1:
print("--start-page must be >= 1.", file=sys.stderr)
sys.exit(1)
total_found = len(images)
try:
selected = common.parse_page_selection(args.pages, images) if args.pages else None
except ValueError as e:
print(str(e), file=sys.stderr)
sys.exit(1)
# --start-page skips pages entirely (matches novel_ocr.py); numbering in
# combined.md picks back up from --start-page regardless.
if args.start_page > 1:
images = images[args.start_page - 1:]
if selected is not None:
images = [p for p in images if p.stem in selected]
if not images:
print("No pages left to process after applying --start-page/--pages.", file=sys.stderr)
sys.exit(1)
print(f"Pages found: {total_found}; processing {len(images)}")
state = {} if args.no_state else common.load_state(output_dir)
if state:
common.warn_if_prompt_changed(state, images, prompt_hash, prompt_file.name)
log_dir = None if args.no_log else Path(args.log_dir)
use_context = args.translate and args.context_pages > 0
recent_context = deque(maxlen=args.context_pages) if use_context else None
context_src_dir = Path(args.context_src_dir) if args.context_src_dir else output_dir / "context_src"
client = OpenAI(base_url=base_url, api_key=api_key)
if use_ahead_context and args.ahead_ocr_backend == "llm":
ahead_ocr_state["client"] = client
ahead_ocr_state["model"] = model
combined_path = output_dir / "combined.md"
failed = []
redone = 0
totals = {}
with open(combined_path, "w", encoding="utf-8") as combined_f:
for pos, img_path in enumerate(tqdm(images, desc="OCR")):
idx = args.start_page + pos
txt_out = pages_dir / f"{img_path.stem}.txt"
has_txt = txt_out.exists()
# A page is skipped if its file exists, unless --redo was passed.
# --no-state only disables the manifest, not this rule.
needs_run = args.redo or not has_txt
if not needs_run:
text = txt_out.read_text(encoding="utf-8")
else:
context_used_past = []
context_used_ahead = []
if use_context and recent_context:
context_used_past = list(recent_context)
if use_ahead_context:
for future_img in images[pos + 1: pos + 1 + args.context_pages_ahead]:
jp_text = get_japanese_context(future_img, context_src_dir, args.ahead_ocr_backend, ahead_ocr_state)
if jp_text and jp_text != "[NO_TEXT]":
context_used_ahead.append((future_img.stem, jp_text))
user_text = common.build_context_block(context_used_past, context_used_ahead) or None
log_entry = {
"timestamp": datetime.now().isoformat(),
"backend": "manga_ocr_llm",
"page": img_path.name,
"translate": args.translate,
"target_lang": args.target_lang if args.translate else None,
"glossary_used": bool(glossary),
"context_pages_used": [name for name, _ in context_used_past],
"context_pages_ahead_used": [name for name, _ in context_used_ahead],
"model": model,
"base_url": base_url,
} if log_dir is not None else None
try:
text, meta = common.ocr_image(
client, model, system_prompt, img_path, request_kwargs, extra_body,
user_text=user_text, log_entry=log_entry,
)
txt_out.write_text(text, encoding="utf-8")
state[img_path.stem] = {
"model": model,
"prompt_hash": prompt_hash,
"finish_reason": meta.get("finish_reason"),
"usage": meta.get("usage"),
"chars": len(text),
"timestamp": datetime.now().isoformat(),
}
if has_txt:
redone += 1
usage = meta.get("usage") or {}
for key, val in usage.items():
if isinstance(val, int):
totals[key] = totals.get(key, 0) + val
except Exception as e: # noqa: BLE001
tqdm.write(f"Error on {img_path.name}: {e}", file=sys.stderr)
failed.append(img_path.name)
text = ""
if log_entry is not None:
log_entry["error"] = str(e)
# Do NOT write a file to disk on failure — otherwise the next
# run would see the file exists and assume the page is already
# done, silently skipping a retry forever.
if log_entry is not None:
common.write_log(log_dir, img_path.stem, log_entry)
if args.sleep:
time.sleep(args.sleep)
# Feed this page's (already-completed-or-just-translated) text
# forward as context for the following pages. Blank/[NO_TEXT]
# pages carry no useful continuity, so skip adding those.
if use_context and text.strip() and text.strip() != "[NO_TEXT]":
recent_context.append((img_path.stem, text.strip()))
combined_f.write(f"\n\n<!-- page {idx}: {img_path.name} -->\n\n")
combined_f.write(text)
if not args.no_state:
common.save_state(output_dir, state, {
"backend": "manga_ocr_llm",
"model": model,
"translate": args.translate,
"target_lang": args.target_lang if args.translate else None,
"prompt_file": prompt_file.name,
"prompt_hash": prompt_hash,
"pages_total": total_found,
"updated": datetime.now().isoformat(),
})
print(f"\nDone. Combined file: {combined_path}")
print(f"Per-page files: {pages_dir}")
if use_ahead_context:
print(f"Cached lookahead Japanese context: {context_src_dir}")
if log_dir is not None:
print(f"Request/response logs: {log_dir}")
if redone:
print(f"Re-OCR'd: {redone} page(s)")
if totals:
print("Tokens: " + ", ".join(f"{k}={v}" for k, v in sorted(totals.items())))
if not args.no_state:
print(f"Manifest: {output_dir / common.STATE_FILE}")
if failed:
print(f"\nFailed to OCR {len(failed)} page(s):")
for name in failed:
print(f" - {name}")
print("Re-run the script with the same --output folder — already-done pages will not be redone.")
if __name__ == "__main__":
main()