cv-analysis-service/cv_analysis/utils/pdf2image.py

from dataclasses import dataclass
from functools import partial
from typing import Iterator

import fitz
import numpy as np


@dataclass
class ImageMetadataPair:
    image: np.ndarray
    metadata: dict


def pdf_to_image_metadata_pairs(pdf: bytes, index=None, dpi=200) -> Iterator[ImageMetadataPair]:
    """Streams PDF as pairs of image (matrix) and metadata.
    Note: If Index is not given or evaluates to None, the whole PDF will be processed.
    """
    convert_fn = partial(page_to_image_metadata_pair, dpi=dpi)
    yield from map(convert_fn, stream_pages(pdf, index))


def page_to_image_metadata_pair(page: fitz.Page, dpi):
    metadata = get_page_info(page)
    pixmap = page.get_pixmap(dpi=dpi)
    array = np.frombuffer(pixmap.samples, dtype=np.uint8).reshape(pixmap.h, pixmap.w, pixmap.n)

    return ImageMetadataPair(array, metadata)


def stream_pages(pdf: bytes, index=None) -> Iterator[fitz.Page]:
    with fitz.open(stream=pdf) as pdf_handle:
        if not index:
            yield from pdf_handle
        else:
            for i in index:
                yield pdf_handle[i]


def get_page_info(page):
    return {
        "index": page.number,
        "rotation": page.rotation,
        "width": page.rect.width,  # rotated page width in inches
        "height": page.rect.height,  # rotated page height in inches
    }