arabic-ocr/arabic_ocr_smart.py

#!/usr/bin/env python3
"""
Arabic OCR — universal prompting via Ollama vision model
=========================================================
Uses a single mixed prompt that handles all document types:
handwritten text, certificates, IDs, tables, forms, and printed Arabic.

Usage:
    python arabic_ocr_smart.py document.pdf [output.txt]
    python arabic_ocr_smart.py scan.jpg [output.txt]
    python arabic_ocr_smart.py scan.png [output.txt]
    python arabic_ocr_smart.py scan.pdf --model qwen2.5vl:7b    # default
    python arabic_ocr_smart.py scan.pdf --model llava:13b        # any Ollama model
    python arabic_ocr_smart.py scan.pdf --host http://localhost:11434
    python arabic_ocr_smart.py scan.pdf --dpi 300                # PDF only

Install:
    pip install pdf2image pillow
    sudo apt-get install -y poppler-utils
"""

import sys
import base64
import json
import io
import argparse
import urllib.request
import urllib.error
from pathlib import Path
from pdf2image import convert_from_path
from PIL import Image

# ── Config ────────────────────────────────────────────────────────────────────

DEFAULT_HOST  = "http://192.168.122.1:11434"
DEFAULT_MODEL = "qwen2.5vl:7b"
DPI           = 300
TIMEOUT       = 600   # seconds between streaming chunks (not total response time)

# ── Prompt ────────────────────────────────────────────────────────────────────

PROMPT = (
    # ── Persona ────────────────────────────────────────────────────────────
    "أنت عالِم متخصص في قراءة المخطوطات العربية وخبير بجميع أنماط الخط العربي: النسخ، الرقعة، الديواني، الإجازة، والكوفي. "
    "عملت عقوداً في فك رموز الوثائق الرسمية والمخطوطات اليدوية الصعبة. "
    "لا يوجد نص عربي يعجز عنك — حتى الخط الرديء أو المتلاشي تستطيع قراءته بالسياق والخبرة.\n\n"

    # ── Two-pass schema recognition ────────────────────────────────────────
    "قبل أن تبدأ النسخ، قم بهاتين الخطوتين الذهنيتين بصمت:\n"
    "الخطوة الأولى — تعرّف على نوع الوثيقة: ما هي؟ ما السياق العام؟ "
    "من المُرسِل والمُرسَل إليه إن وُجدا؟ ما الغرض من النص؟ "
    "لا تقيّد نفسك بأنواع محددة — الوثيقة تُعرّف نفسها بنفسها.\n"
    "الخطوة الثانية — توقّع المحتوى: بناءً على نوع الوثيقة الذي حددته، "
    "استحضر من خبرتك الصيغَ والمصطلحاتِ والتراكيبَ اللغوية الشائعة في هذا النوع تحديداً. "
    "استخدم هذه التوقعات كمرجع لحل الغموض في الخط — لا كقيد يمنعك من قراءة ما هو فعلاً مكتوب.\n"
    "بعد هاتين الخطوتين، ابدأ النسخ الفعلي.\n\n"

    # ── Task ───────────────────────────────────────────────────────────────
    "مهمتك: انسخ كل النص المرئي في هذه الصورة حرفياً.\n\n"

    # ── Strict rules ───────────────────────────────────────────────────────
    "قواعد صارمة:\n"
    "- اقرأ الصورة سطراً سطراً من الأعلى إلى الأسفل. كل سطر في الصورة = سطر جديد في الإخراج. لا تدمج أسطراً منفصلة في جملة واحدة.\n"
    "- اتجاه القراءة من اليمين إلى اليسار للعربية، ومن اليسار إلى اليمين للأرقام والنص اللاتيني.\n"
    "- الأرقام المضمّنة داخل النص العربي: اكتبها كما تظهر (عربية أو هندية أو لاتينية) دون تحويل.\n"
    "- للكلمات غير الواضحة: استخدم السياق والكلمات المحيطة لاستنتاج القراءة الأرجح.\n"
    "- الخطأ الأكبر هو كتابة كلمة خاطئة بدلاً من الاعتراف بالغموض. [؟] أفضل دائماً من كلمة مخترعة.\n"
    "- ممنوع منعاً باتاً اختراع كلمات أو أسماء لم تظهر في الصورة. إذا كانت الكلمة غير واضحة تماماً، ضع [؟].\n"
    "- الصيغ الرسمية الشائعة في أي نوع من الوثائق يجب التعرف عليها كوحدات متكاملة وليس كلمة كلمة.\n"
    "- احتفظ بالتشكيل (الحركات) إن ظهر في الصورة.\n"
    "- الأختام والطوابع الرسمية: إذا كان نصها مقروءاً، انسخه هكذا: [ختم: ...].\n"
    "- لا تعكس ترتيب الكلمات أو الأسطر — حافظ على الاتجاه الأصلي للنص.\n"
    "- لا تبدأ ردك بأي مقدمة أو جملة تمهيدية — ابدأ مباشرةً بأول كلمة في الصورة.\n"
    "- لا تُضِف أي تعليق أو وصف أو تفسير — فقط النص المستخرج.\n\n"

    # ── Output format ──────────────────────────────────────────────────────
    "تنسيق الإخراج حسب نوع المحتوى:\n"
    "- نص عادي (مطبوع أو يدوي): اكتبه سطراً سطراً كما هو.\n"
    "- عناوين وترويسات: ضعها على سطر منفصل مع نجمتين ** حولها.\n"
    "- جداول: أعد إنتاجها بتنسيق Markdown مع الحفاظ على جميع الأعمدة والصفوف.\n"
    "- حقول نماذج: اكتبها بصيغة 'اسم الحقل: القيمة'، وإذا كان فارغاً اكتب [فارغ].\n"
    "- مربعات اختيار: ✓ للمحدد، ☐ للفارغ.\n"
    "- بيانات هوية أو وثيقة رسمية: استخرج كل حقل مرئي بصيغة 'اسم الحقل: القيمة'.\n"
    "- منطقة MRZ (الأحرف الآلية في أسفل جواز السفر أو الهوية): انسخها كما هي في سطر منفصل.\n"
)

# ── Helpers ───────────────────────────────────────────────────────────────────

def image_to_base64(pil_image: Image.Image) -> str:
    buf = io.BytesIO()
    pil_image.save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode()


def call_ollama(host: str, model: str, pil_image: Image.Image, timeout: int = TIMEOUT, num_ctx: int = 8192) -> str:
    payload = {
        "model": model,
        "messages": [
            {
                "role": "user",
                "content": PROMPT,
                "images": [image_to_base64(pil_image)],
            }
        ],
        "stream": True,
        "options": {"num_ctx": num_ctx},
    }
    data = json.dumps(payload).encode()
    req  = urllib.request.Request(
        f"{host}/api/chat",
        data=data,
        headers={"Content-Type": "application/json"},
    )
    try:
        chunks = []
        eval_count = None
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            for line in resp:
                if not line.strip():
                    continue
                obj = json.loads(line)
                chunk = obj.get("message", {}).get("content", "")
                if chunk:
                    if not chunks:
                        print("generating", end="", flush=True)
                    chunks.append(chunk)
                    print(".", end="", flush=True)
                if obj.get("done"):
                    eval_count = obj.get("eval_count")
                    prompt_eval_count = obj.get("prompt_eval_count")
                    break
        return "".join(chunks).strip(), eval_count, prompt_eval_count
    except urllib.error.URLError as e:
        raise ConnectionError(
            f"Cannot reach Ollama at {host}. "
            f"Is it running? (OLLAMA_HOST=0.0.0.0:11434 ollama serve)\n{e}"
        ) from e


def ocr_page(host: str, model: str, pil_image: Image.Image, page_num: int, timeout: int = TIMEOUT, num_ctx: int = 8192) -> str:
    print(f"  Page {page_num}: prefill...", end=" ", flush=True)
    text, eval_count, prompt_eval_count = call_ollama(host, model, pil_image, timeout, num_ctx)
    tokens_info = f" (prompt: {prompt_eval_count}, output: {eval_count})" if eval_count is not None else ""
    print(f" done.{tokens_info}", flush=True)
    return text


# ── Main ──────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(
        description="Arabic OCR with universal mixed-content prompting.",
    )
    parser.add_argument("input",  help="Input file: PDF, JPEG, or PNG")
    parser.add_argument("output", nargs="?", help="Output .txt file (optional)")
    parser.add_argument(
        "--host", default=DEFAULT_HOST,
        help=f"Ollama host URL (default: {DEFAULT_HOST})"
    )
    parser.add_argument(
        "--model", default=DEFAULT_MODEL,
        help=f"Ollama model name (default: {DEFAULT_MODEL})"
    )
    parser.add_argument(
        "--dpi", type=int, default=DPI,
        help=f"PDF render resolution (default: {DPI})"
    )
    parser.add_argument(
        "--timeout", type=int, default=TIMEOUT,
        help=f"Seconds to wait between streaming chunks (default: {TIMEOUT})"
    )
    parser.add_argument(
        "--ctx", type=int, default=12288,
        help="Model context window in tokens (default: 12288)"
    )
    parser.add_argument(
        "--poppler", default=None,
        help="Path to poppler bin/ directory (Windows only, e.g. C:\\poppler\\bin)"
    )
    args = parser.parse_args()

    input_path  = Path(args.input)
    output_path = Path(args.output) if args.output \
                  else input_path.with_name(input_path.stem + "_ocr.txt")

    if not input_path.is_file():
        sys.exit(f"Error: file not found: {input_path}")

    suffix = input_path.suffix.lower()
    if suffix not in {".pdf", ".jpg", ".jpeg", ".png"}:
        sys.exit(f"Error: unsupported file type '{suffix}'. Supported: .pdf, .jpg, .jpeg, .png")

    print(f"\n[*] Input  : {input_path}")
    print(f"[*] Model  : {args.model}")
    print(f"[*] Ollama : {args.host}")
    print(f"[*] Output : {output_path}\n")

    if suffix == ".pdf":
        print("[*] Converting PDF to images...")
        pages = convert_from_path(str(input_path), dpi=args.dpi, poppler_path=args.poppler)
        print(f"    {len(pages)} page(s) found.\n")
    else:
        pages = [Image.open(input_path).convert("RGB")]
        print(f"[*] Loaded image ({suffix})\n")

    label  = "Page" if suffix == ".pdf" else "Image"
    sections = []
    for i, page_img in enumerate(pages, start=1):
        text = ocr_page(args.host, args.model, page_img, i, args.timeout, args.ctx)
        header = f"{'='*60}\n{label} {i}\n{'='*60}"
        sections.append(f"{header}\n\n{text}")
        print()

    output_path.write_text("\n\n".join(sections), encoding="utf-8")
    print(f"[✓] Done. Output saved to: {output_path}")


if __name__ == "__main__":
    main()