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/manga_ocr_llm.py
+132 -230
View File
@@ -8,10 +8,9 @@ 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, and — since that's
what most people actually want out of a manga OCR pass — can translate
directly instead, with a glossary for consistent names/terms and a sliding
window of previous pages' translations for continuity.
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
@@ -36,6 +35,11 @@ Usage (after filling in config.json):
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
@@ -46,18 +50,27 @@ One-time setup:
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 base64
import io
import json
import os
import re
import sys
import time
from collections import deque
@@ -69,124 +82,12 @@ from openai import OpenAI
from PIL import Image
from tqdm import tqdm
IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".tif", ".tiff", ".bmp"}
SCRIPT_DIR = Path(__file__).resolve().parent
ROOT_DIR = SCRIPT_DIR.parent
sys.path.insert(0, str(SCRIPT_DIR)) # so `from local_mangaocr_ocr import ...` works when run from elsewhere
sys.path.insert(0, str(SCRIPT_DIR)) # so `import common` / `from local_mangaocr_ocr import ...` work when run from elsewhere
import common
def load_config(config_path: Path) -> dict:
"""Loads the shared config.json and returns its "ocr" section.
Falls back to treating the whole file as the ocr config if there's no
"ocr" key, 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("ocr", data) if isinstance(data, dict) else {}
def build_request_kwargs(config: dict, args: argparse.Namespace) -> 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
def load_prompt(prompt_path: Path, target_language: str = None, glossary: str = None) -> str:
if not prompt_path.exists():
print(f"Prompt file not found: {prompt_path}", file=sys.stderr)
sys.exit(1)
prompt = prompt_path.read_text(encoding="utf-8").strip()
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 find_glossary(explicit_path: str) -> Path:
"""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_context_block(past_pages: list, ahead_pages: list) -> str:
"""Builds a prompt suffix with:
- 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.
"""
if not past_pages and not ahead_pages:
return ""
parts = []
if past_pages:
parts.append(
"\n\n---\n\n"
"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:
parts.append(
"\n\n---\n\n"
"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)
IMAGE_EXTS = common.IMAGE_EXTS
def get_japanese_context(
@@ -224,8 +125,8 @@ def get_japanese_context(
model = ahead_ocr_state["model"]
request_kwargs = ahead_ocr_state["request_kwargs"]
extra_body = ahead_ocr_state["extra_body"]
prompt = ahead_ocr_state["prompt"]
text = ocr_image(client, model, prompt, img_path, request_kwargs, 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 ""
@@ -236,93 +137,6 @@ def get_japanese_context(
return text
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")
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}"
def ocr_image(
client: OpenAI, model: str, prompt: str, path: Path,
request_kwargs: dict, extra_body: dict,
retries: int = 3, log_entry: dict = None,
) -> str:
data_url = image_to_data_url(path)
if log_entry is not None:
log_entry["request"] = {
"model": model,
**request_kwargs,
"extra_body": extra_body or None,
"messages": [
{
"role": "user",
"content": [
{"type": "text", "text": prompt},
{"type": "image_url", "image_url": {"url": f"<omitted: {path.name}, {path.stat().st_size} bytes>"}},
],
}
],
}
log_entry["attempts"] = []
last_err = None
for attempt in range(retries):
try:
response = client.chat.completions.create(
model=model,
messages=[
{
"role": "user",
"content": [
{"type": "text", "text": prompt},
{"type": "image_url", "image_url": {"url": data_url}},
],
}
],
extra_body=extra_body or None,
**request_kwargs,
)
text = (response.choices[0].message.content or "").strip()
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
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}")
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)")
@@ -381,7 +195,7 @@ def main():
"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 <output>/context_src/"
"Defaults to <o>/context_src/"
)
parser.add_argument(
"--prompt-file", default=None,
@@ -402,7 +216,23 @@ def main():
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)
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)"
@@ -418,7 +248,7 @@ def main():
)
args = parser.parse_args()
config = load_config(Path(args.config))
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")
@@ -433,12 +263,12 @@ def main():
)
sys.exit(1)
request_kwargs, extra_body = build_request_kwargs(config, args)
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 = find_glossary(args.glossary)
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}")
@@ -447,15 +277,22 @@ def main():
target_language = None
glossary = None
prompt_file = Path(args.prompt_file) if args.prompt_file else SCRIPT_DIR / default_prompt_name
prompt = load_prompt(prompt_file, target_language, glossary)
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":
# Plain transcription prompt/params for OCR'ing lookahead pages —
# never the translate prompt (we want the raw Japanese, not a
# draft translation — see the design note in build_context_block).
ahead_ocr_state["prompt"] = load_prompt(SCRIPT_DIR / "prompt_manga.txt")
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.
@@ -473,7 +310,33 @@ def main():
print(f"No images found in {input_dir}.", file=sys.stderr)
sys.exit(1)
print(f"Pages found: {len(images)}")
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)
@@ -490,16 +353,21 @@ def main():
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 txt_out.exists():
if not needs_run:
text = txt_out.read_text(encoding="utf-8")
else:
page_prompt = prompt
context_used_past = []
context_used_ahead = []
if use_context and recent_context:
@@ -509,8 +377,7 @@ def main():
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))
if context_used_past or context_used_ahead:
page_prompt = prompt + build_context_block(context_used_past, context_used_ahead)
user_text = common.build_context_block(context_used_past, context_used_ahead) or None
log_entry = {
"timestamp": datetime.now().isoformat(),
@@ -525,10 +392,27 @@ def main():
"base_url": base_url,
} if log_dir is not None else None
try:
text = ocr_image(client, model, page_prompt, img_path, request_kwargs, extra_body, log_entry=log_entry)
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
print(f"\nError on {img_path.name}: {e}", file=sys.stderr)
tqdm.write(f"Error on {img_path.name}: {e}", file=sys.stderr)
failed.append(img_path.name)
text = ""
if log_entry is not None:
@@ -537,7 +421,7 @@ def main():
# run would see the file exists and assume the page is already
# done, silently skipping a retry forever.
if log_entry is not None:
write_log(log_dir, img_path.stem, log_entry)
common.write_log(log_dir, img_path.stem, log_entry)
if args.sleep:
time.sleep(args.sleep)
@@ -550,12 +434,30 @@ def main():
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: