345 lines
15 KiB
Python
345 lines
15 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Batch-OCR (or batch-translate) for scanned pages of a Japanese light novel
|
|
(vertical text) 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.
|
|
|
|
This lets a multimodal LLM "read" a page image directly, including
|
|
vertical Japanese text, without any column-segmentation preprocessing.
|
|
|
|
Deliberately simple compared to manga_ocr_llm.py: light novel pages are
|
|
dense running prose read in one linear order, so there's no bubble
|
|
type/reading-order logic here — see manga_ocr_llm.py for that, which is a
|
|
genuinely different job (scattered speech bubbles across panels). The
|
|
--translate machinery (glossary, sliding-window context) is shared
|
|
infrastructure from ocr/common.py, used the same way by both scripts.
|
|
|
|
Usage (after filling in config.json):
|
|
python novel_ocr.py --input ./pages --output ./out
|
|
|
|
# translate instead of transcribing:
|
|
python novel_ocr.py --input ./pages --output ./out \\
|
|
--translate --target-lang Russian
|
|
|
|
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_novel.txt (transcription) or
|
|
prompt_novel_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, with --translate
|
|
--context-pages, a bit of sliding-window context) varies per page and goes
|
|
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).
|
|
"""
|
|
|
|
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 tqdm import tqdm
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent)) # so `import common` works when run from elsewhere
|
|
import common
|
|
|
|
IMAGE_EXTS = common.IMAGE_EXTS
|
|
SCRIPT_DIR = Path(__file__).resolve().parent
|
|
ROOT_DIR = SCRIPT_DIR.parent
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Batch-OCR/translate a light novel 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(
|
|
"--prompt-file", default=None,
|
|
help="Path to the prompt file (defaults to prompt_novel.txt, or "
|
|
"prompt_novel_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 just those pages."
|
|
)
|
|
parser.add_argument(
|
|
"--redo", action="store_true",
|
|
help="Re-OCR pages even if a transcription 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"
|
|
)
|
|
|
|
# --- Translation (mirrors manga_ocr_llm.py) ---
|
|
parser.add_argument(
|
|
"--translate", action="store_true",
|
|
help="Translate the page instead of transcribing the Japanese "
|
|
"(uses prompt_novel_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. Shared with manga_ocr_llm.py."
|
|
)
|
|
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 narration, "
|
|
"pronouns/referents). 0 disables it. Default: 2"
|
|
)
|
|
|
|
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_novel_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_novel.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)
|
|
|
|
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 now actually skips pages rather than only relabelling them.
|
|
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
|
|
|
|
client = OpenAI(base_url=base_url, api_key=api_key)
|
|
|
|
combined_path = output_dir / "combined.md"
|
|
failed = []
|
|
redone = 0
|
|
totals = {}
|
|
|
|
with open(combined_path, "w", encoding="utf-8") as combined_f:
|
|
for idx, img_path in enumerate(tqdm(images, desc="OCR"), start=args.start_page):
|
|
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:
|
|
user_text = None
|
|
context_used = []
|
|
if use_context and recent_context:
|
|
context_used = list(recent_context)
|
|
user_text = common.build_context_block(context_used, [])
|
|
|
|
log_entry = {
|
|
"timestamp": datetime.now().isoformat(),
|
|
"backend": "novel_ocr",
|
|
"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],
|
|
"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)
|
|
|
|
# Written once per run (not per page) so an interrupted run simply
|
|
# leaves the previous manifest in place rather than a partial one.
|
|
if not args.no_state:
|
|
common.save_state(output_dir, state, {
|
|
"backend": "novel_ocr",
|
|
"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 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()
|