poison-flower/pf-ocr
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

update

Browse Source
This commit is contained in:
Poison Flower
2026-09-18 01:00:58 +03:00
parent 02da1023ee
commit 41c994277b
10 changed files with 843 additions and 456 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ocr/common.py
+403
View File
@@ -0,0 +1,403 @@
"""
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)