Poison Flower
35b7be680e
- Consolidate module configs into root config.example.json with ocr, vision, and epub sections - Split LLM OCR workflows into novel_ocr.py (prose) and manga_ocr_llm.py (manga) - Remove gemini_direct_ocr.py in favor of OpenAI-compatible API endpoints - Support direct manga translation via --translate, --target-lang, and glossary.md - Add bidirectional context support: past translations (--context-pages) and lookahead Japanese text (--context-pages-ahead) - Add per-page JSON audit logging under logs/ and expose OpenAI sampling parameters
321 lines
13 KiB
Python
321 lines
13 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Batch-OCR 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: light novel pages are dense running prose, and all
|
|
this script needs to do is transcribe them accurately. It has no
|
|
translate/glossary/context-page machinery — see manga_ocr_llm.py for that,
|
|
which is a genuinely different job (short, scattered dialogue that needs
|
|
continuity handling).
|
|
|
|
Usage (after filling in config.json):
|
|
python novel_ocr.py --input ./pages --output ./out
|
|
|
|
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 to fit your book / house style
|
|
|
|
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 base64
|
|
import io
|
|
import json
|
|
import os
|
|
import re
|
|
import sys
|
|
import time
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
|
|
from natsort import natsorted
|
|
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
|
|
|
|
|
|
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 load_prompt(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()
|
|
|
|
|
|
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 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 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=str(SCRIPT_DIR / "prompt_novel.txt"),
|
|
help="Path to the prompt file (defaults to prompt_novel.txt next to this script)"
|
|
)
|
|
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)
|
|
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 = load_config(Path(args.config))
|
|
|
|
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 = build_request_kwargs(config, args)
|
|
prompt = load_prompt(Path(args.prompt_file))
|
|
|
|
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)
|
|
|
|
print(f"Pages found: {len(images)}")
|
|
|
|
log_dir = None if args.no_log else Path(args.log_dir)
|
|
|
|
client = OpenAI(base_url=base_url, api_key=api_key)
|
|
|
|
combined_path = output_dir / "combined.md"
|
|
failed = []
|
|
|
|
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"
|
|
|
|
if txt_out.exists():
|
|
text = txt_out.read_text(encoding="utf-8")
|
|
else:
|
|
log_entry = {
|
|
"timestamp": datetime.now().isoformat(),
|
|
"backend": "novel_ocr",
|
|
"page": img_path.name,
|
|
"model": model,
|
|
"base_url": base_url,
|
|
} if log_dir is not None else None
|
|
try:
|
|
text = ocr_image(client, model, prompt, img_path, request_kwargs, extra_body, log_entry=log_entry)
|
|
txt_out.write_text(text, encoding="utf-8")
|
|
except Exception as e: # noqa: BLE001
|
|
print(f"\nError 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:
|
|
write_log(log_dir, img_path.stem, log_entry)
|
|
if args.sleep:
|
|
time.sleep(args.sleep)
|
|
|
|
combined_f.write(f"\n\n<!-- page {idx}: {img_path.name} -->\n\n")
|
|
combined_f.write(text)
|
|
|
|
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 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()
|