Spaces:
Running
Running
File size: 2,834 Bytes
a8784d9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | """PDF -> list of A4-sized RGB page images using PyMuPDF (fitz).
"Native split" handling: if a page is already ~A4 portrait we render it as-is.
If a page is much larger / a different ratio (e.g. an A3 spread, a long scan),
we slice it into A4-height bands so downstream tearing always works on A4 tiles.
"""
from __future__ import annotations
import numpy as np
from .config import A4_ASPECT, MAX_PAGES_PER_PDF, a4_pixels
def _page_pixmap(page, dpi: int) -> np.ndarray:
import fitz # PyMuPDF
zoom = dpi / 72.0
pm = page.get_pixmap(matrix=fitz.Matrix(zoom, zoom), alpha=False)
arr = np.frombuffer(pm.samples, dtype=np.uint8).reshape(pm.h, pm.w, pm.n)
if pm.n == 4: # RGBA -> RGB
arr = arr[:, :, :3]
elif pm.n == 1: # gray -> RGB
arr = np.repeat(arr, 3, axis=2)
return np.ascontiguousarray(arr)
def _fit_to_a4(img: np.ndarray, dpi: int) -> list[np.ndarray]:
"""Return one or more A4-portrait tiles covering `img`.
Tall pages are sliced into A4-height bands (each band padded to A4 width).
"""
target_w, target_h = a4_pixels(dpi)
H, W = img.shape[:2]
aspect = H / W
# Close enough to A4 portrait: letterbox-resize the whole page onto A4.
if abs(aspect - A4_ASPECT) < 0.12:
return [_letterbox(img, target_w, target_h)]
# Otherwise scale to A4 width, then slice the (now-tall) image into bands.
scale = target_w / W
new_h = max(1, int(round(H * scale)))
resized = _resize(img, target_w, new_h)
tiles = []
for top in range(0, new_h, target_h):
band = resized[top:top + target_h]
if band.shape[0] < target_h:
band = _letterbox(band, target_w, target_h)
tiles.append(band)
return tiles
def _resize(img: np.ndarray, w: int, h: int) -> np.ndarray:
from PIL import Image
return np.asarray(Image.fromarray(img).resize((w, h), Image.LANCZOS))
def _letterbox(img: np.ndarray, w: int, h: int) -> np.ndarray:
"""Resize preserving aspect, pad with white onto a w*h canvas."""
H, W = img.shape[:2]
scale = min(w / W, h / H)
nw, nh = max(1, int(W * scale)), max(1, int(H * scale))
resized = _resize(img, nw, nh)
canvas = np.full((h, w, 3), 255, dtype=np.uint8)
oy, ox = (h - nh) // 2, (w - nw) // 2
canvas[oy:oy + nh, ox:ox + nw] = resized
return canvas
def load_pdf_pages(path: str, dpi: int) -> list[np.ndarray]:
"""Render `path` into a list of A4 RGB uint8 page images."""
import fitz
doc = fitz.open(path)
try:
out: list[np.ndarray] = []
for page in doc:
raw = _page_pixmap(page, dpi)
out.extend(_fit_to_a4(raw, dpi))
if len(out) >= MAX_PAGES_PER_PDF:
return out[:MAX_PAGES_PER_PDF]
return out
finally:
doc.close()
|