pdf2zh/parser/utils/image.py

"""PDF page rendering utilities.

This module provides functions to render PDF pages to PIL Images
for Surya processing.
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

import fitz  # PyMuPDF
from PIL import Image

if TYPE_CHECKING:
    pass

logger = logging.getLogger(__name__)


def _extract_dominant_image(
    page: fitz.Page,
    coverage_threshold: float = 0.8,
) -> Image.Image | None:
    """Extract the dominant full-page embedded image, if one exists.

    For scanned PDFs the page content is typically a single high-resolution
    raster image embedded at 300+ DPI.  Re-rendering through fitz at 150 DPI
    halves the effective resolution and causes Surya to miss fine layout
    structure.  Extracting the raw embedded image preserves the original DPI.

    Args:
        page: fitz Page object
        coverage_threshold: Minimum fraction of page area the image must cover

    Returns:
        PIL Image if a dominant embedded image is found, else None
    """
    page_rect = page.rect
    page_area = page_rect.width * page_rect.height
    if page_area == 0:
        return None

    image_list = page.get_images(full=True)
    best_xref = None
    best_coverage = 0.0

    for img_info in image_list:
        xref = img_info[0]
        try:
            rects = list(page.get_image_rects(xref))
            if not rects:
                continue
            covered = sum(r.width * r.height for r in rects) / page_area
            if covered > best_coverage:
                best_coverage = covered
                best_xref = xref
        except Exception:
            continue

    if best_xref is None or best_coverage < coverage_threshold:
        return None

    try:
        pix = fitz.Pixmap(page.parent, best_xref)
        # Normalize to RGB. pix.n counts all components including alpha.
        # Grayscale (n=1), Gray+A (n=2), CMYK (n=4 no alpha), CMYK+A (n=5)
        # all need conversion — only pure RGB (n=3, alpha=0) is already correct.
        n_colors = pix.n - pix.alpha
        if n_colors != 3:
            pix = fitz.Pixmap(fitz.csRGB, pix)
        mode = "RGBA" if pix.alpha else "RGB"
        img = Image.frombytes(mode, (pix.width, pix.height), pix.samples)
        if pix.alpha:
            img = img.convert("RGB")
        logger.debug(
            f"Using embedded image (xref={best_xref}, "
            f"{pix.width}×{pix.height}px, n_colors={n_colors}, coverage={best_coverage:.0%})"
        )
        return img
    except Exception as e:
        logger.debug(f"Could not extract embedded image xref={best_xref}: {e}")
        return None


def render_page_to_image(
    page: fitz.Page,
    dpi: int = 150,
) -> Image.Image:
    """Render a PDF page for OCR — prefers native embedded image.

    Prefers extracting the dominant full-page embedded image at its native
    resolution.  Higher resolution means clearer characters for OCR.
    Falls back to fitz rendering when no dominant embedded image is found.

    Args:
        page: fitz Page object to render
        dpi: Resolution used for fitz fallback rendering (default 150)

    Returns:
        PIL Image in RGB mode
    """
    img = _extract_dominant_image(page)
    if img is not None:
        return img
    return _fitz_render(page, dpi)


def render_page_for_layout(
    page: fitz.Page,
    dpi: int = 96,
) -> Image.Image:
    """Render a PDF page for layout detection — always uses fitz at fixed DPI.

    Surya's layout model (and its sibling detection/reading-order models) is
    calibrated to receive images rendered at IMAGE_DPI = 96.  That is the DPI
    that Surya itself uses internally (see surya/input/processing.py:
    get_page_images).  Passing images at a much higher DPI (e.g. 150-300)
    causes the layout model to see the same document content at a larger pixel
    scale than it was trained on, which results in collapsed or missed layout
    bounding boxes.

    This function always renders through fitz so that the image scale is
    predictable regardless of whether embedded raster images exist.

    Args:
        page: fitz Page object to render
        dpi: Target DPI for layout rendering (default 96, matching Surya IMAGE_DPI)

    Returns:
        PIL Image in RGB mode
    """
    return _fitz_render(page, dpi)


def _fitz_render(page: fitz.Page, dpi: int) -> Image.Image:
    """Render a PDF page to a PIL Image using fitz at the specified DPI.

    Converts ``dpi`` to a scale factor relative to fitz's default 72 DPI and
    renders the page without an alpha channel (RGB mode).

    Args:
        page: fitz Page object to render.
        dpi: Output resolution in dots per inch.  Higher values produce larger
             images with finer detail but require more memory.

    Returns:
        PIL Image in ``RGB`` mode at the requested DPI.
    """
    zoom = dpi / 72.0
    matrix = fitz.Matrix(zoom, zoom)
    pixmap = page.get_pixmap(matrix=matrix, alpha=False)
    return Image.frombytes("RGB", (pixmap.width, pixmap.height), pixmap.samples)


def render_pages_batch(
    doc: fitz.Document,
    page_indices: list[int],
    dpi: int = 150,
) -> list[Image.Image]:
    """Render multiple PDF pages to PIL Images.

    Args:
        doc: fitz Document object
        page_indices: List of 0-based page indices to render
        dpi: Resolution in dots per inch

    Returns:
        List of PIL Images in same order as page_indices
    """
    images = []
    for page_idx in page_indices:
        page = doc[page_idx]
        img = render_page_to_image(page, dpi=dpi)
        images.append(img)
    return images


def get_page_dimensions(page: fitz.Page) -> tuple[float, float]:
    """Get page dimensions in PDF points.

    Args:
        page: fitz Page object

    Returns:
        (width, height) in PDF points
    """
    rect = page.rect
    return rect.width, rect.height


def crop_image_to_bbox(
    image: Image.Image,
    bbox: list[float],
    pdf_width: float,
    pdf_height: float,
) -> Image.Image:
    """Crop a rendered image to a bounding box.

    Args:
        image: PIL Image rendered from PDF page
        bbox: [x0, y0, x1, y1] in PDF points
        pdf_width: Original page width in PDF points
        pdf_height: Original page height in PDF points

    Returns:
        Cropped PIL Image
    """
    img_width, img_height = image.size

    # Calculate scale factors
    scale_x = img_width / pdf_width
    scale_y = img_height / pdf_height

    # Convert bbox to image pixels
    x0 = int(bbox[0] * scale_x)
    y0 = int(bbox[1] * scale_y)
    x1 = int(bbox[2] * scale_x)
    y1 = int(bbox[3] * scale_y)

    # Clamp to image bounds
    x0 = max(0, min(x0, img_width))
    y0 = max(0, min(y0, img_height))
    x1 = max(0, min(x1, img_width))
    y1 = max(0, min(y1, img_height))

    # Ensure valid crop region
    if x1 <= x0 or y1 <= y0:
        # Return a small placeholder image
        return Image.new("RGB", (1, 1), color=(255, 255, 255))

    return image.crop((x0, y0, x1, y1))