Fahad0saad/arabic-ocr-v2

1"""Qwen3-VL OCR integration via OpenRouter."""2from __future__ import annotations3 4import base645import os6import re7from pathlib import Path8from typing import Optional9 10import httpx11 12OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions"13 14DEFAULT_MODEL = "qwen/qwen3-vl-8b-instruct"15FALLBACK_MODEL = "qwen/qwen3-vl-235b-a22b-instruct"16 17OCR_PROMPT = (18    "أنت نظام متخصص في التعرف الضوئي على الحروف (OCR) للنصوص العربية. "19    "قم باستخراج كل النص العربي الظاهر في الصورة بدقة عالية، مع الحفاظ على:\n"20    "1. ترتيب الفقرات والأسطر كما تظهر في الصورة\n"21    "2. علامات الترقيم والمسافات\n"22    "3. الأرقام والرموز كما هي\n\n"23    "أعد فقط النص المستخرج بدون أي شرح أو مقدمة أو تعليق. "24    "إذا لم تجد نصاً عربياً، أعد النص بأي لغة موجود في الصورة."25)26 27TASHKEEL_PROMPT = (28    "أضف التشكيل الكامل (الحركات) إلى النص العربي التالي بشكل صحيح نحوياً وصرفياً، "29    "مع الحفاظ على المعنى الأصلي. أعد فقط النص المشكّل بدون أي تعليق:\n\n"30)31 32CORRECTION_PROMPT = (33    "صحّح الأخطاء الإملائية والنحوية في النص العربي التالي الناتج عن OCR، "34    "مع الحفاظ على المعنى الأصلي وترتيب الأسطر. أعد فقط النص المصحّح بدون تعليق:\n\n"35)36 37 38def _api_key() -> str:39    key = os.environ.get("OPENROUTER_API_KEY", "").strip()40    if not key:41        raise RuntimeError("OPENROUTER_API_KEY environment variable is not set")42    return key43 44 45def _image_to_data_url(image_path: str | Path) -> str:46    p = Path(image_path)47    ext = p.suffix.lower().lstrip(".") or "png"48    mime = "jpeg" if ext in ("jpg", "jpeg") else ext49    data = base64.b64encode(p.read_bytes()).decode("ascii")50    return f"data:image/{mime};base64,{data}"51 52 53async def _call_model(54    model: str, messages: list[dict], timeout: float = 120.055) -> str:56    headers = {57        "Authorization": f"Bearer {_api_key()}",58        "Content-Type": "application/json",59        # OpenRouter ranking headers60        "HTTP-Referer": os.environ.get("OPENROUTER_REFERER", "https://arabic-ocr.local"),61        "X-Title": "Arabic OCR v2",62    }63    payload = {64        "model": model,65        "messages": messages,66        "temperature": 0.1,67        "max_tokens": 4096,68    }69    async with httpx.AsyncClient(timeout=timeout) as client:70        resp = await client.post(OPENROUTER_URL, headers=headers, json=payload)71        resp.raise_for_status()72        data = resp.json()73 74    choices = data.get("choices") or []75    if not choices:76        return ""77    msg = choices[0].get("message", {})78    content = msg.get("content")79    if isinstance(content, list):80        # Some vision models return list of parts81        parts = []82        for part in content:83            if isinstance(part, dict) and part.get("type") == "text":84                parts.append(part.get("text", ""))85            elif isinstance(part, str):86                parts.append(part)87        return "".join(parts).strip()88    return (content or "").strip()89 90 91def _strip_fences(text: str) -> str:92    text = text.strip()93    if text.startswith("```"):94        text = re.sub(r"^```[a-zA-Z]*\n?", "", text)95        text = re.sub(r"\n?```$", "", text)96    return text.strip()97 98 99async def run_ocr(100    image_path: str | Path,101    model: str = DEFAULT_MODEL,102    fallback: str = FALLBACK_MODEL,103) -> str:104    """Run vision OCR on a single image, with fallback to a stronger model."""105    data_url = _image_to_data_url(image_path)106    messages = [107        {108            "role": "user",109            "content": [110                {"type": "text", "text": OCR_PROMPT},111                {"type": "image_url", "image_url": {"url": data_url}},112            ],113        }114    ]115 116    try:117        out = await _call_model(model, messages)118        text = _strip_fences(out)119        if text:120            return text121    except Exception:122        pass123 124    # Fallback model125    out = await _call_model(fallback, messages)126    return _strip_fences(out)127 128 129async def post_process(130    text: str,131    auto_tashkeel: bool = False,132    ai_correction: bool = False,133    model: str = DEFAULT_MODEL,134) -> str:135    """Apply optional AI post-processing: correction then tashkeel."""136    if not text.strip():137        return text138 139    if ai_correction:140        try:141            corrected = await _call_model(142                model,143                [{"role": "user", "content": CORRECTION_PROMPT + text}],144            )145            corrected = _strip_fences(corrected)146            if corrected:147                text = corrected148        except Exception:149            pass150 151    if auto_tashkeel:152        try:153            tashkeel = await _call_model(154                model,155                [{"role": "user", "content": TASHKEEL_PROMPT + text}],156            )157            tashkeel = _strip_fences(tashkeel)158            if tashkeel:159                text = tashkeel160        except Exception:161            pass162 163    return text164