ocr_studio/export_pdf.py

from __future__ import annotations

from io import BytesIO
from pathlib import Path

from PIL import Image
from reportlab.lib.utils import ImageReader
from reportlab.pdfbase import pdfmetrics
from reportlab.pdfbase.ttfonts import TTFont
from reportlab.pdfgen.canvas import Canvas

from ocr_studio.config import ARABIC_FONT_PATH, LATIN_FONT_PATH, PDF_DPI
from ocr_studio.language import contains_arabic_script
from ocr_studio.spotting import TextSpan

_FONTS_REGISTERED = False
LATIN_FONT = "NotoSans"
ARABIC_FONT = "NotoNaskhArabic"


def _register_fonts() -> None:
    global _FONTS_REGISTERED
    if _FONTS_REGISTERED:
        return
    if LATIN_FONT_PATH.exists():
        pdfmetrics.registerFont(TTFont(LATIN_FONT, str(LATIN_FONT_PATH)))
    if ARABIC_FONT_PATH.exists():
        pdfmetrics.registerFont(TTFont(ARABIC_FONT, str(ARABIC_FONT_PATH)))
    _FONTS_REGISTERED = True


def _font_for(text: str) -> str:
    if contains_arabic_script(text) and ARABIC_FONT_PATH.exists():
        return ARABIC_FONT
    if LATIN_FONT_PATH.exists():
        return LATIN_FONT
    return "Helvetica"


def _fallback_spans(image: Image.Image, text: str, rtl: bool) -> list[TextSpan]:
    lines = [line.strip() for line in (text or "").splitlines() if line.strip()]
    if not lines:
        return []
    top = image.height * 0.06
    bottom = image.height * 0.94
    usable = max(bottom - top, float(len(lines) * 14))
    line_height = usable / max(len(lines), 1)
    left = image.width * 0.05
    right = image.width * 0.95
    spans: list[TextSpan] = []
    for index, line in enumerate(lines):
        y0 = top + index * line_height
        y1 = min(image.height - 2.0, y0 + max(12.0, line_height * 0.85))
        if rtl and contains_arabic_script(line):
            box = (left + (right - left) * 0.15, y0, right, y1)
        else:
            box = (left, y0, right, y1)
        spans.append(TextSpan(text=line, box=box))
    return spans


def _draw_page(
    canvas: Canvas,
    image: Image.Image,
    text: str,
    spans: list[TextSpan] | None,
    rtl: bool,
) -> None:
    page_w = image.width * 72.0 / PDF_DPI
    page_h = image.height * 72.0 / PDF_DPI
    scale = page_w / float(image.width)
    canvas.setPageSize((page_w, page_h))

    buffer = BytesIO()
    image.convert("RGB").save(buffer, format="JPEG", quality=92)
    buffer.seek(0)
    canvas.drawImage(
        ImageReader(buffer),
        0,
        0,
        width=page_w,
        height=page_h,
        preserveAspectRatio=True,
        mask="auto",
    )

    overlay = [span for span in (spans or []) if span.text]
    if not overlay:
        overlay = _fallback_spans(image, text, rtl)

    for span in overlay:
        box = span.box or (image.width * 0.05, image.height * 0.05, image.width * 0.95, image.height * 0.12)
        x0, y0, x1, y1 = box
        pdf_x = x0 * scale
        pdf_y = page_h - (y1 * scale)
        box_h = max(8.0, (y1 - y0) * scale)
        box_w = max(12.0, (x1 - x0) * scale)
        font_size = max(6.0, min(box_h * 0.82, 28.0))
        font_name = _font_for(span.text)
        canvas.setFont(font_name, font_size)
        canvas.setFillGray(0)
        text_obj = canvas.beginText()
        text_obj.setTextRenderMode(3)
        text_obj.setTextOrigin(pdf_x, pdf_y + max(1.0, (box_h - font_size) / 2.0))
        clipped = span.text
        try:
            text_width = canvas.stringWidth(clipped, font_name, font_size)
            if text_width > box_w and box_w > 0:
                font_size = max(5.0, font_size * (box_w / text_width))
                canvas.setFont(font_name, font_size)
                text_obj = canvas.beginText()
                text_obj.setTextRenderMode(3)
                text_obj.setTextOrigin(pdf_x, pdf_y + max(1.0, (box_h - font_size) / 2.0))
        except Exception:
            pass
        text_obj.textOut(clipped)
        canvas.drawText(text_obj)


def build_searchable_pdf(
    image: Image.Image | list[Image.Image],
    text: str | list[str],
    destination: Path,
    spans: list[TextSpan] | list[list[TextSpan]] | None = None,
    rtl: bool = False,
) -> Path:
    _register_fonts()
    images = image if isinstance(image, list) else [image]
    texts = text if isinstance(text, list) else [text]
    if len(texts) < len(images):
        texts = texts + [""] * (len(images) - len(texts))

    span_pages: list[list[TextSpan]]
    if spans is None:
        span_pages = [[] for _ in images]
    elif spans and isinstance(spans[0], TextSpan):
        span_pages = [spans]  # type: ignore[list-item]
        if len(images) > 1:
            span_pages = span_pages + [[] for _ in range(len(images) - 1)]
    else:
        span_pages = list(spans)  # type: ignore[arg-type]
        while len(span_pages) < len(images):
            span_pages.append([])

    first = images[0]
    page_w = first.width * 72.0 / PDF_DPI
    page_h = first.height * 72.0 / PDF_DPI
    canvas = Canvas(str(destination), pagesize=(page_w, page_h))
    canvas.setTitle("Searchable OCR scan")
    canvas.setAuthor("Alireza Aminzadeh")

    for index, page_image in enumerate(images):
        if index:
            canvas.showPage()
        _draw_page(canvas, page_image, texts[index], span_pages[index], rtl)

    canvas.save()
    return destination