ocr_studio/image_io.py

from __future__ import annotations

from io import BytesIO
from pathlib import Path

from PIL import Image, ImageOps, UnidentifiedImageError

from ocr_studio.config import (
    ALLOWED_DOCUMENT_SUFFIXES,
    ALLOWED_IMAGE_FORMATS,
    ALLOWED_IMAGE_SUFFIXES,
    MAX_BATCH_ITEMS,
    MAX_IMAGE_SIDE,
    MAX_IMAGE_SIDE_ACCURATE,
    MAX_PDF_PAGES,
    MAX_UPLOAD_BYTES,
    MAX_UPLOAD_MB,
    PDF_RENDER_SCALE,
)
from ocr_studio.errors import OcrError

_HEIF_READY = False


def register_heif() -> None:
    global _HEIF_READY
    if _HEIF_READY:
        return
    try:
        from pillow_heif import register_heif_opener

        register_heif_opener()
        _HEIF_READY = True
    except Exception:
        _HEIF_READY = False


register_heif()


def _as_file_list(files: object) -> list[Path]:
    if files is None:
        return []
    if isinstance(files, (str, Path)):
        raw = str(files).strip()
        if not raw or raw in {".", "./"}:
            return []
        return [Path(raw)]
    if isinstance(files, dict) and files.get("path"):
        return [Path(str(files["path"]))]
    if isinstance(files, (list, tuple)):
        paths: list[Path] = []
        for item in files:
            paths.extend(_as_file_list(item))
        return paths
    return []


def _open_image(source: Image.Image | str | Path | bytes) -> Image.Image:
    if isinstance(source, Image.Image):
        image = source
    elif hasattr(source, "shape") and hasattr(source, "dtype"):
        image = Image.fromarray(source)
    elif isinstance(source, (str, Path)):
        path = Path(source)
        if not path.exists():
            raise OcrError("The uploaded image could not be read.")
        if path.stat().st_size > MAX_UPLOAD_BYTES:
            raise OcrError(f"Image is larger than {MAX_UPLOAD_MB} MB.")
        suffix = path.suffix.lower()
        if suffix not in ALLOWED_IMAGE_SUFFIXES:
            raise OcrError("Use JPG, PNG, WEBP, BMP, TIFF, or HEIC.")
        try:
            image = Image.open(path)
        except UnidentifiedImageError as exc:
            raise OcrError("The file is not a valid image.") from exc
    elif isinstance(source, (bytes, bytearray)):
        if len(source) > MAX_UPLOAD_BYTES:
            raise OcrError(f"Image is larger than {MAX_UPLOAD_MB} MB.")
        try:
            image = Image.open(BytesIO(source))
        except UnidentifiedImageError as exc:
            raise OcrError("The file is not a valid image.") from exc
    else:
        raise OcrError("Please upload one image.", key="err.no_input")

    try:
        image.load()
    except OSError as exc:
        raise OcrError("The image file is damaged.") from exc
    return image


def prepare_image(
    source: Image.Image | str | Path | bytes | None,
    accurate: bool = False,
) -> Image.Image:
    if source is None:
        raise OcrError("Please upload one image.", key="err.no_input")
    image = _open_image(source)
    fmt = (image.format or "").upper()
    if fmt and fmt not in ALLOWED_IMAGE_FORMATS:
        raise OcrError("Use JPG, PNG, WEBP, BMP, TIFF, or HEIC.")

    image = ImageOps.exif_transpose(image).convert("RGB")
    longest = max(image.size)
    limit = MAX_IMAGE_SIDE_ACCURATE if accurate else MAX_IMAGE_SIDE
    if longest > limit:
        scale = limit / float(longest)
        image = image.resize(
            (max(1, int(image.width * scale)), max(1, int(image.height * scale))),
            Image.Resampling.LANCZOS,
        )
    return image


def render_pdf_pages(path: Path, accurate: bool = False) -> tuple[list[Image.Image], int]:
    if path.stat().st_size > MAX_UPLOAD_BYTES:
        raise OcrError(f"PDF is larger than {MAX_UPLOAD_MB} MB.")
    try:
        import pypdfium2 as pdfium
    except Exception as exc:
        raise OcrError("PDF support is not available in this environment.") from exc

    try:
        document = pdfium.PdfDocument(str(path))
    except Exception as exc:
        raise OcrError("The PDF file could not be opened.") from exc

    total = len(document)
    count = min(total, MAX_PDF_PAGES)
    pages: list[Image.Image] = []
    try:
        for index in range(count):
            page = document[index]
            bitmap = page.render(scale=PDF_RENDER_SCALE)
            rendered = bitmap.to_pil()
            pages.append(prepare_image(rendered, accurate=accurate))
    finally:
        document.close()
    return pages, total


def collect_pages(
    image: Image.Image | str | Path | bytes | None,
    files: object,
    accurate: bool = False,
) -> tuple[list[Image.Image], int, bool]:
    pages: list[Image.Image] = []
    source_files = _as_file_list(files)
    if len(source_files) > MAX_BATCH_ITEMS:
            raise OcrError(
                f"Upload at most {MAX_BATCH_ITEMS} files per job.",
                key="err.too_many_files",
                max_items=MAX_BATCH_ITEMS,
            )

    if image is not None:
        pages.append(prepare_image(image, accurate=accurate))

    total_pdf_pages = 0
    truncated = False
    for path in source_files:
        if not path.exists():
            raise OcrError("The uploaded file could not be read.")
        if path.stat().st_size > MAX_UPLOAD_BYTES:
            raise OcrError(f"A file is larger than {MAX_UPLOAD_MB} MB.")
        suffix = path.suffix.lower()
        if suffix not in ALLOWED_DOCUMENT_SUFFIXES:
            raise OcrError("Use JPG, PNG, WEBP, BMP, TIFF, HEIC, or PDF.")
        if suffix == ".pdf":
            rendered, total = render_pdf_pages(path, accurate=accurate)
            total_pdf_pages += total
            if total > MAX_PDF_PAGES:
                truncated = True
            remaining = MAX_PDF_PAGES - len(pages)
            if remaining <= 0:
                truncated = True
                break
            pages.extend(rendered[:remaining])
            if total > remaining:
                truncated = True
        else:
            if len(pages) >= MAX_PDF_PAGES:
                truncated = True
                break
            pages.append(prepare_image(path, accurate=accurate))

    if not pages:
        raise OcrError("Please upload one image.", key="err.no_input")
    if len(pages) > MAX_PDF_PAGES:
        truncated = True
        pages = pages[:MAX_PDF_PAGES]
    return pages, max(total_pdf_pages, len(pages)), truncated