404 lines
17 KiB
Python
404 lines
17 KiB
Python
"""
|
|
Shared infrastructure for the OCR scripts (novel_ocr.py, manga_ocr_llm.py).
|
|
|
|
Everything here is genuinely identical between the two "page layout"
|
|
scripts — config/state/logging plumbing, the OpenAI-compatible request
|
|
helper, and the translate-only helpers (glossary, sliding-window context).
|
|
What's deliberately NOT here is anything that depends on page layout
|
|
(reading order, output format, the prompt files themselves) — that stays
|
|
in each script, since novel and manga pages need genuinely different
|
|
logic there.
|
|
|
|
system/user prompt split (for provider-side prompt caching):
|
|
Everything that stays IDENTICAL across every page of a run — the
|
|
format/transcription instructions, and (for --translate) the glossary
|
|
— is sent as the `system` message. Everything that changes per page —
|
|
the image, and any sliding-window context — goes in `user`. Most
|
|
OpenAI-compatible providers cache a repeated system prefix, so with a
|
|
long book this avoids re-billing/re-processing the same instructions
|
|
on every single page. Per-page context (previous/upcoming pages) is
|
|
NOT cached-friendly by nature (it changes every page), so it belongs
|
|
in `user`, never appended to the cacheable `system` text.
|
|
"""
|
|
|
|
import base64
|
|
import io
|
|
import json
|
|
import re
|
|
import sys
|
|
import time
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
|
|
from openai import OpenAI
|
|
from PIL import Image
|
|
|
|
IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".tif", ".tiff", ".bmp"}
|
|
STATE_FILE = "state.json"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# config.json
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def load_config_section(config_path: Path, section: str) -> dict:
|
|
"""Loads the shared config.json and returns one top-level section
|
|
(e.g. "ocr"). Falls back to treating the whole file as that section if
|
|
the key is missing, so a bare {"api_key": ...} style file still works.
|
|
"""
|
|
if not config_path.exists():
|
|
return {}
|
|
try:
|
|
data = json.loads(config_path.read_text(encoding="utf-8"))
|
|
except json.JSONDecodeError as e:
|
|
print(f"Failed to parse {config_path}: {e}", file=sys.stderr)
|
|
sys.exit(1)
|
|
return data.get(section, data) if isinstance(data, dict) else {}
|
|
|
|
|
|
def load_prompt_file(prompt_path: Path) -> str:
|
|
if not prompt_path.exists():
|
|
print(f"Prompt file not found: {prompt_path}", file=sys.stderr)
|
|
sys.exit(1)
|
|
return prompt_path.read_text(encoding="utf-8").strip()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Page selection (--pages "3,7-12,20", --start-page)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def parse_page_selection(spec: str, images: list) -> set | None:
|
|
"""Parses a --pages selection like "3,7-12,20" into a set of page stems.
|
|
|
|
Numbers refer to the 1-based position in the naturally-sorted page list,
|
|
so "3" is the third page, not a file literally named "3". Returns None
|
|
when the spec is empty/absent, meaning "no explicit selection".
|
|
"""
|
|
if not spec:
|
|
return None
|
|
selected = set()
|
|
for part in spec.split(","):
|
|
part = part.strip()
|
|
if not part:
|
|
continue
|
|
m = re.fullmatch(r"(\d+)\s*(?:-\s*(\d+))?", part)
|
|
if not m:
|
|
raise ValueError(f"Invalid --pages entry: {part!r} (expected N or N-M)")
|
|
start = int(m.group(1))
|
|
end = int(m.group(2)) if m.group(2) else start
|
|
if start < 1 or end < start:
|
|
raise ValueError(f"Invalid --pages range: {part!r}")
|
|
for n in range(start, end + 1):
|
|
if n > len(images):
|
|
raise ValueError(f"--pages refers to page {n}, but only {len(images)} page(s) were found")
|
|
selected.add(images[n - 1].stem)
|
|
return selected
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# state.json manifest (resume/redo tracking, prompt-change detection)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def prompt_fingerprint(prompt: str) -> str:
|
|
"""Short stable hash of the prompt text, used to detect prompt edits."""
|
|
import hashlib
|
|
return hashlib.sha256(prompt.encode("utf-8")).hexdigest()[:12]
|
|
|
|
|
|
def load_state(output_dir: Path) -> dict:
|
|
"""Loads the pages manifest, mapping page stem -> what produced it.
|
|
|
|
Returns {} if the manifest is missing or unreadable; callers must treat
|
|
a page with no entry as "unverified" rather than "stale".
|
|
"""
|
|
path = output_dir / STATE_FILE
|
|
if not path.exists():
|
|
return {}
|
|
try:
|
|
data = json.loads(path.read_text(encoding="utf-8"))
|
|
except (json.JSONDecodeError, OSError) as e: # noqa: BLE001
|
|
print(f"Warning: could not read {path} ({e}); existing pages are reused unverified.", file=sys.stderr)
|
|
return {}
|
|
return data.get("pages", {}) if isinstance(data, dict) else {}
|
|
|
|
|
|
def save_state(output_dir: Path, pages: dict, meta: dict) -> None:
|
|
"""Atomically rewrites the pages manifest (temp file + replace)."""
|
|
path = output_dir / STATE_FILE
|
|
tmp = path.with_suffix(".json.tmp")
|
|
tmp.write_text(
|
|
json.dumps({"meta": meta, "pages": pages}, ensure_ascii=False, indent=2, default=str),
|
|
encoding="utf-8",
|
|
)
|
|
tmp.replace(path)
|
|
|
|
|
|
def warn_if_prompt_changed(state: dict, images: list, prompt_hash: str, prompt_file_name: str) -> None:
|
|
"""Warns (does not fail) if some already-transcribed pages were done
|
|
with a prompt whose hash no longer matches the current one."""
|
|
if not state:
|
|
return
|
|
stale = [
|
|
p.stem for p in images
|
|
if p.stem in state and state[p.stem].get("prompt_hash") not in (None, prompt_hash)
|
|
]
|
|
if stale:
|
|
print(
|
|
f"Note: {len(stale)} page(s) were transcribed with a different prompt "
|
|
f"({prompt_file_name} changed). Use --redo to re-OCR them.",
|
|
file=sys.stderr,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Request/response logging
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def write_log(log_dir: Path, page_name: str, entry: dict) -> None:
|
|
"""Writes one JSON log file per request/response.
|
|
|
|
The base64 image payload itself is never included (megabytes of no
|
|
debugging value, one per page) — only its path/size — but everything
|
|
else (full prompt text, model params, every retry attempt, the full raw
|
|
API response) is recorded as-is.
|
|
"""
|
|
log_dir.mkdir(parents=True, exist_ok=True)
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
|
|
safe_name = re.sub(r"[^A-Za-z0-9_.-]", "_", page_name)
|
|
log_path = log_dir / f"{timestamp}_{safe_name}.json"
|
|
log_path.write_text(json.dumps(entry, ensure_ascii=False, indent=2, default=str), encoding="utf-8")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Request params (CLI flag > config.json > default)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def build_request_kwargs(config: dict, args) -> tuple[dict, dict]:
|
|
"""Resolves OpenAI-API-style request params: CLI flag > config.json > a
|
|
sensible default. Returns (kwargs, extra_body) — standard params go
|
|
straight into the request; reasoning_effort goes through extra_body
|
|
since support for it varies by provider/model and extra_body is the
|
|
designed passthrough for exactly that.
|
|
"""
|
|
temperature = args.temperature if args.temperature is not None else config.get("temperature", 0)
|
|
max_tokens = args.max_tokens if args.max_tokens is not None else config.get("max_tokens")
|
|
top_p = args.top_p if args.top_p is not None else config.get("top_p")
|
|
reasoning_effort = args.reasoning_effort or config.get("reasoning_effort")
|
|
|
|
kwargs = {"temperature": temperature}
|
|
if max_tokens is not None:
|
|
kwargs["max_tokens"] = max_tokens
|
|
if top_p is not None:
|
|
kwargs["top_p"] = top_p
|
|
|
|
extra_body = {}
|
|
if reasoning_effort:
|
|
extra_body["reasoning_effort"] = reasoning_effort
|
|
|
|
return kwargs, extra_body
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Image encoding
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def image_to_data_url(path: Path, max_dim: int = 2200) -> str:
|
|
"""Downscale (if needed) and encode an image as a base64 data URL."""
|
|
img = Image.open(path)
|
|
if img.mode not in ("RGB", "L"):
|
|
img = img.convert("RGB")
|
|
if max(img.size) > max_dim:
|
|
ratio = max_dim / max(img.size)
|
|
img = img.resize((int(img.width * ratio), int(img.height * ratio)))
|
|
|
|
buf = io.BytesIO()
|
|
img.save(buf, format="JPEG", quality=90)
|
|
b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
|
|
return f"data:image/jpeg;base64,{b64}"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# LLM OCR call, with system/user split for prompt caching
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def usage_dict(response) -> dict | None:
|
|
"""Extracts token usage from a response, tolerating providers that omit it."""
|
|
usage = getattr(response, "usage", None)
|
|
if usage is None:
|
|
return None
|
|
try:
|
|
data = usage.model_dump() if hasattr(usage, "model_dump") else dict(usage)
|
|
except Exception: # noqa: BLE001
|
|
return None
|
|
return {k: v for k, v in data.items() if v is not None} or None
|
|
|
|
|
|
def ocr_image(
|
|
client: OpenAI, model: str, system_prompt: str, path: Path,
|
|
request_kwargs: dict, extra_body: dict,
|
|
user_text: str = None, retries: int = 3, log_entry: dict | None = None,
|
|
) -> tuple[str, dict]:
|
|
"""OCRs (or translates) one page image.
|
|
|
|
system_prompt: the part of the prompt that's IDENTICAL across every
|
|
page of this run (format instructions, glossary) — sent as the
|
|
`system` message so OpenAI-compatible providers can cache it
|
|
instead of reprocessing it on every page.
|
|
user_text: optional per-page text that varies from page to page
|
|
(sliding-window context) — sent alongside the image in `user`.
|
|
Never put per-page content in system_prompt, or every page
|
|
invalidates the cache instead of hitting it.
|
|
|
|
Returns (text, meta) where meta records what the provider reported
|
|
about the accepted response (finish_reason, usage) so the caller can
|
|
persist it (e.g. in the pages manifest).
|
|
"""
|
|
data_url = image_to_data_url(path)
|
|
user_content = [{"type": "image_url", "image_url": {"url": data_url}}]
|
|
if user_text:
|
|
user_content.insert(0, {"type": "text", "text": user_text})
|
|
|
|
messages = [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": user_content},
|
|
]
|
|
|
|
if log_entry is not None:
|
|
logged_user_content = [
|
|
{"type": "text", "text": user_text} if user_text else None,
|
|
{"type": "image_url", "image_url": {"url": f"<omitted: {path.name}, {path.stat().st_size} bytes>"}},
|
|
]
|
|
log_entry["request"] = {
|
|
"model": model,
|
|
**request_kwargs,
|
|
"extra_body": extra_body or None,
|
|
"messages": [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": [c for c in logged_user_content if c]},
|
|
],
|
|
}
|
|
log_entry["attempts"] = []
|
|
|
|
last_err = None
|
|
for attempt in range(retries):
|
|
try:
|
|
response = client.chat.completions.create(
|
|
model=model,
|
|
messages=messages,
|
|
extra_body=extra_body or None,
|
|
**request_kwargs,
|
|
)
|
|
choice = response.choices[0]
|
|
text = (choice.message.content or "").strip()
|
|
finish_reason = getattr(choice, "finish_reason", None)
|
|
|
|
# A truncated or empty response must not be accepted as a
|
|
# successful transcription — this file would otherwise be
|
|
# written to disk and silently skipped on every future retry.
|
|
if finish_reason == "length":
|
|
raise RuntimeError(f"response truncated (finish_reason=length, {len(text)} chars)")
|
|
if not text:
|
|
raise RuntimeError(f"empty response (finish_reason={finish_reason})")
|
|
|
|
if log_entry is not None:
|
|
try:
|
|
raw_response = response.model_dump()
|
|
except Exception: # noqa: BLE001
|
|
raw_response = {"content": text}
|
|
log_entry["attempts"].append({"attempt": attempt + 1, "success": True, "response": raw_response})
|
|
return text, {"finish_reason": finish_reason, "usage": usage_dict(response)}
|
|
except Exception as e: # noqa: BLE001
|
|
last_err = e
|
|
if log_entry is not None:
|
|
log_entry["attempts"].append({"attempt": attempt + 1, "success": False, "error": str(e)})
|
|
time.sleep(3 * (attempt + 1))
|
|
raise RuntimeError(f"Failed to OCR {path.name} after {retries} attempts: {last_err}")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Translation helpers: glossary + sliding-window context (--translate only)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def find_glossary(script_dir: Path, explicit_path: str) -> Path | None:
|
|
"""Resolves the glossary path: explicit --glossary if given, otherwise
|
|
an auto-detected glossary.md next to the ocr scripts themselves (a
|
|
stable, book-independent location — handy for a future UI to edit)."""
|
|
if explicit_path:
|
|
path = Path(explicit_path)
|
|
if not path.exists():
|
|
print(f"Glossary file not found: {path}", file=sys.stderr)
|
|
sys.exit(1)
|
|
return path
|
|
auto_path = script_dir / "glossary.md"
|
|
return auto_path if auto_path.exists() else None
|
|
|
|
|
|
def build_system_prompt(base_prompt: str, target_language: str = None, glossary: str = None) -> str:
|
|
"""Builds the cacheable `system` prompt: format instructions, with the
|
|
target language substituted in and the glossary appended.
|
|
|
|
Deliberately excludes per-page context (see build_context_block) —
|
|
that changes every page and must go in `user`, or every page would
|
|
invalidate the provider's cached system prefix instead of hitting it.
|
|
"""
|
|
prompt = base_prompt
|
|
if target_language:
|
|
prompt = prompt.replace("{target_language}", target_language)
|
|
if glossary:
|
|
prompt += (
|
|
"\n\n---\n\n"
|
|
"Glossary: use these established translations for the following "
|
|
"names/terms for consistency. Anything not listed here, translate "
|
|
"naturally.\n\n" + glossary
|
|
)
|
|
return prompt
|
|
|
|
|
|
def build_context_block(past_pages: list, ahead_pages: list) -> str:
|
|
"""Builds the per-page `user` text carrying sliding-window context:
|
|
- past_pages: the last N *translated* pages, for continuity (character
|
|
voice, an ongoing exchange, pronouns/referents that only make sense
|
|
given what was just said).
|
|
- ahead_pages: the next N pages' raw *Japanese* transcription (never a
|
|
draft translation — see the ocr/README.md design note on why), for
|
|
context that only becomes clear from what's about to happen (a
|
|
twist, who's actually being addressed, a joke that pays off a page
|
|
later).
|
|
|
|
Both are lists of (page_name, text) tuples, oldest/nearest first.
|
|
Returns "" if there's nothing to add.
|
|
|
|
This text is per-page by nature and must never be folded into the
|
|
cacheable system prompt (see build_system_prompt) — it changes on
|
|
every single page, so appending it to system would defeat prompt
|
|
caching entirely instead of the intended "cache the instructions,
|
|
vary only the image/context" split.
|
|
"""
|
|
if not past_pages and not ahead_pages:
|
|
return ""
|
|
parts = []
|
|
if past_pages:
|
|
parts.append(
|
|
"Context: translations of the immediately preceding page(s), for "
|
|
"continuity only (character voice, an ongoing conversation, "
|
|
"pronouns/referents). Do NOT re-translate or repeat any of this "
|
|
"in your output — translate only the NEW page shown in the image.\n"
|
|
)
|
|
for name, text in past_pages:
|
|
parts.append(f"\n[Previous page: {name}]\n{text}")
|
|
if ahead_pages:
|
|
if parts:
|
|
parts.append("\n\n---\n\n")
|
|
parts.append(
|
|
"Context: the ORIGINAL JAPANESE (not a translation) of the "
|
|
"upcoming page(s) that follow the one you're translating now. "
|
|
"Use this only to correctly resolve things that depend on what "
|
|
"happens next — pronoun/referent gender, who a line is actually "
|
|
"addressed to, a setup whose payoff lands later. Do NOT "
|
|
"translate or otherwise output any of this — translate only the "
|
|
"CURRENT page shown in the image.\n"
|
|
)
|
|
for name, text in ahead_pages:
|
|
parts.append(f"\n[Upcoming page ({name}), original Japanese]\n{text}")
|
|
return "".join(parts)
|