47 lines
1.3 KiB
Python
47 lines
1.3 KiB
Python
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
|
|
}
|