apis/baseimg2ocr.py

"""
baseimg2ocr: base64画像をOCRしてテキストを抽出。
HF版: VLM (Qwen2.5-VL) を使用。Google Vision API は使用しない。

NOTE: 元の実装は Google Vision API を使用。精度が異なる場合がある。
"""

import json
import os

from src.utils.tracer import customtracer


def _vlm_ocr(base64image: str, model: str = "Qwen/Qwen2.5-VL-7B-Instruct") -> str:
    """VLM でテキスト抽出（url2ocr 互換フォーマット）。"""
    from src.clients.llm_client import LLMClient
    from pydantic import BaseModel
    from typing import List

    class OcrEntry(BaseModel):
        text: str
        y: int
        size: int

    class OcrResult(BaseModel):
        items: List[OcrEntry]

    client = LLMClient()
    result = client.call(
        prompt=(
            "Extract all visible text from this image. "
            "For each text block, estimate its vertical position (y coordinate, 0=top) "
            "and approximate font size in pixels. "
            "Return results sorted by y position."
        ),
        schema=OcrResult,
        model=model,
        images=[base64image],
        temperature=0,
    )
    return json.dumps(
        [{"text": e.text, "y": e.y, "size": e.size, "rect": []} for e in result.items],
        ensure_ascii=False,
    )


@customtracer
def baseimg2ocr(base64image: str, margin: int = 120) -> str:
    """
    input1 (text): base64エンコードされた画像
    input2 (text): 120
    output1 (json): OCR結果

    NOTE: HF版は VLM ベースOCR。Google Vision API は使用しない。
    """
    return _vlm_ocr(base64image)