init

ocr/gemini_direct_ocr.py

@@ -0,0 +1,146 @@
+#!/usr/bin/env python3
+"""
+Batch-OCR for scanned pages of a Japanese light novel (vertical text) via the
+Gemini API directly (not through a proxy).
+
+Usage:
+    export GEMINI_API_KEY="your_key"
+    python gemini_direct_ocr.py --input ./pages --output ./out
+
+Requirements:
+    pip install google-genai pillow natsort tqdm
+
+If you access Gemini (or another model) through OpenRouter or a similar
+OpenAI-compatible proxy instead of a direct Google API key, use
+openrouter_ocr.py instead — it's the recommended entry point for this
+project and shares the same prompt.txt / config.json workflow.
+"""
+
+import argparse
+import os
+import sys
+import time
+from pathlib import Path
+
+from google import genai
+from google.genai import types
+from natsort import natsorted
+from PIL import Image
+from tqdm import tqdm
+
+IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".tif", ".tiff", ".bmp"}
+
+PROMPT = """\
+This is a scanned page from a Japanese light novel, printed in vertical text (縦書き).
+Transcribe ALL Japanese text on the page, preserving correct reading order
+(columns read top-to-bottom, then right-to-left).
+
+Rules:
+- Do not translate or summarize — output only the transcribed Japanese text.
+- No line numbers, no commentary, no explanations of your own.
+- Furigana may be omitted (only the base kanji/kana text is needed).
+- If the page has no text at all (illustration-only, blank, cover/technical
+  page), output exactly: [NO_TEXT]
+- Preserve paragraph breaks where the layout clearly shows them.
+"""
+
+
+def ocr_image(client: genai.Client, model: str, path: Path, retries: int = 3) -> str:
+    img = Image.open(path)
+    # Downscale very large scans — speeds up and cheapens the request without
+    # a noticeable loss of OCR quality.
+    max_dim = 2200
+    if max(img.size) > max_dim:
+        ratio = max_dim / max(img.size)
+        img = img.resize((int(img.width * ratio), int(img.height * ratio)))
+
+    last_err = None
+    for attempt in range(retries):
+        try:
+            response = client.models.generate_content(
+                model=model,
+                contents=[PROMPT, img],
+                config=types.GenerateContentConfig(temperature=0),
+            )
+            text = (response.text or "").strip()
+            return text
+        except Exception as e:  # noqa: BLE001
+            last_err = 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 the Gemini 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("--model", default="gemini-3.7-flash", help="Gemini model name")
+    parser.add_argument(
+        "--api-key", default=None,
+        help="API key (falls back to the GEMINI_API_KEY environment variable)"
+    )
+    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)"
+    )
+    args = parser.parse_args()
+
+    api_key = args.api_key or os.environ.get("GEMINI_API_KEY")
+    if not api_key:
+        print("No API key found. Pass --api-key or set GEMINI_API_KEY.", file=sys.stderr)
+        sys.exit(1)
+
+    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)}")
+
+    client = genai.Client(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:
+                try:
+                    text = ocr_image(client, args.model, img_path)
+                    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 = ""
+                    # Do NOT write a file to disk on failure — see openrouter_ocr.py
+                    # for the reasoning.
+                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 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()