36 lines
1.3 KiB
Python
36 lines
1.3 KiB
Python
import fitz # PyMuPDF
|
|
import numpy as np
|
|
from PIL import Image
|
|
from app.config import settings
|
|
|
|
class PDFReader:
|
|
def __init__(self, path: str):
|
|
self.doc = fitz.open(path)
|
|
|
|
def is_searchable(self, page) -> bool:
|
|
"""Detect if a page has a usable native text layer."""
|
|
text = page.get_text("text").strip()
|
|
return len(text) > 100 # heuristic
|
|
|
|
def page_to_image(self, page) -> np.ndarray:
|
|
zoom = settings.DPI / 72
|
|
mat = fitz.Matrix(zoom, zoom)
|
|
pix = page.get_pixmap(matrix=mat, alpha=False)
|
|
img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
|
|
return np.array(img)
|
|
|
|
def native_text_blocks(self, page):
|
|
"""Returns word boxes from searchable PDF: [(x0,y0,x1,y1,text)]"""
|
|
words = page.get_text("words")
|
|
scale = settings.DPI / 72
|
|
return [(w[0]*scale, w[1]*scale, w[2]*scale, w[3]*scale, w[4]) for w in words]
|
|
|
|
def iter_pages(self):
|
|
for i, page in enumerate(self.doc):
|
|
yield {
|
|
"page_number": i + 1,
|
|
"image": self.page_to_image(page),
|
|
"searchable": self.is_searchable(page),
|
|
"native_words": self.native_text_blocks(page) if self.is_searchable(page) else [],
|
|
}
|