de921e308f
Merge in RR/cv-analysis from rotation-logic-refactor to master
Squashed commit of the following:
commit 684dd140cbfc9fbebe9beb8c13b52a2d131c9932
Author: Julius Unverfehrt <julius.unverfehrt@iqser.com>
Date: Wed Jul 27 14:22:58 2022 +0200
move rotation logic to before cv-analysis, so that cv-analysis only needs to operate on portrait images and matrix rotation logic can be dropped
48 lines
1.3 KiB
Python
48 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)
|
|
page.set_rotation(0)
|
|
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
|
|
}
|