| """Turn a file into pages: digital-PDF text where usable, otherwise rendered/loaded images. |
| Also page-level orientation correction and deskew.""" |
|
|
| from dataclasses import dataclass, field |
| from pathlib import Path |
|
|
| import cv2 |
| import numpy as np |
|
|
| from ..charset import normalize |
| from ..geometry import merge_line_cells |
|
|
| RENDER_DPI = 200 |
| IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp", ".webp", ".gif"} |
|
|
|
|
| @dataclass |
| class Page: |
| number: int |
| image: np.ndarray | None |
| width: int |
| height: int |
| text_source: str |
| pdf_segments: list = field(default_factory=list) |
|
|
|
|
| def _usable_pdf_text(segs) -> bool: |
| text = "".join(s[4] for s in segs) |
| if len(text.strip()) < 25: |
| return False |
| good = sum(c.isalnum() or c in " .,:;-$€£%/()#&@'\"" for c in text) |
| return good / len(text) > 0.8 and "(cid:" not in text |
|
|
|
|
| def load_pages(path, dpi=RENDER_DPI, force_ocr=False): |
| path = Path(path) |
| suffix = path.suffix.lower() |
| if suffix == ".pdf": |
| import pypdfium2 as pdfium |
|
|
| doc = pdfium.PdfDocument(str(path)) |
| scale = dpi / 72 |
| for i in range(len(doc)): |
| page = doc[i] |
| pw, ph = page.get_size() |
| W, H = round(pw * scale), round(ph * scale) |
| segs = [] |
| if not force_ocr: |
| tp = page.get_textpage() |
| cells = [] |
| for r in range(tp.count_rects()): |
| l, b, rr, t = tp.get_rect(r) |
| txt = normalize(tp.get_text_bounded(l, b, rr, t)) |
| if txt: |
| cells.append((l * scale, (ph - t) * scale, rr * scale, (ph - b) * scale, txt)) |
| segs = [s for s in merge_line_cells(cells) if s[4].strip()] |
| tp.close() |
| if segs and _usable_pdf_text(segs): |
| yield Page(i + 1, None, W, H, "pdf_text", segs) |
| else: |
| img = page.render(scale=scale).to_numpy()[:, :, :3].copy() |
| yield Page(i + 1, img, img.shape[1], img.shape[0], "ocr") |
| page.close() |
| doc.close() |
| elif suffix in IMAGE_EXTS: |
| img = cv2.imread(str(path), cv2.IMREAD_COLOR) |
| if img is None: |
| from PIL import Image |
| img = np.array(Image.open(path).convert("RGB"))[:, :, ::-1] |
| img = np.ascontiguousarray(img[:, :, ::-1]) |
| yield Page(1, img, img.shape[1], img.shape[0], "ocr") |
| else: |
| raise ValueError(f"unsupported file type: {suffix} (supported: .pdf, {', '.join(sorted(IMAGE_EXTS))})") |
|
|
|
|
| def estimate_skew(img_rgb, max_angle=12.0) -> float: |
| """Projection-profile skew estimate in degrees (positive = text rotated counter-clockwise).""" |
| gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY) |
| s = 1000 / max(gray.shape) |
| if s < 1: |
| gray = cv2.resize(gray, (round(gray.shape[1] * s), round(gray.shape[0] * s)), interpolation=cv2.INTER_AREA) |
| ink = cv2.adaptiveThreshold(gray, 1, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY_INV, 25, 15).astype(np.float32) |
| |
| small = cv2.resize(ink, (ink.shape[1] // 2, ink.shape[0] // 2), interpolation=cv2.INTER_AREA) |
|
|
| def score(a, im): |
| h, w = im.shape |
| m = cv2.getRotationMatrix2D((w / 2, h / 2), -a, 1.0) |
| prof = cv2.warpAffine(im, m, (w, h)).sum(1) |
| return float(np.var(prof)) |
|
|
| coarse = np.arange(-max_angle, max_angle + 1e-6, 1.0) |
| best = coarse[int(np.argmax([score(a, small) for a in coarse]))] |
| fine = np.arange(best - 1, best + 1 + 1e-6, 0.1) |
| return float(fine[int(np.argmax([score(a, ink) for a in fine]))]) |
|
|
|
|
| def rotate_image(img, angle): |
| """Rotate by `angle` degrees (counter-clockwise), expanding the canvas. Returns (image, 2x3 matrix).""" |
| h, w = img.shape[:2] |
| m = cv2.getRotationMatrix2D((w / 2, h / 2), angle, 1.0) |
| cos, sin = abs(m[0, 0]), abs(m[0, 1]) |
| nw, nh = int(h * sin + w * cos), int(h * cos + w * sin) |
| m[0, 2] += nw / 2 - w / 2 |
| m[1, 2] += nh / 2 - h / 2 |
| border = tuple(int(v) for v in np.median(img[::8, ::8].reshape(-1, 3), 0)) |
| return cv2.warpAffine(img, m, (nw, nh), flags=cv2.INTER_LINEAR, borderValue=border), m |
|
|
|
|
| def rot90_matrix(k, w, h): |
| """2x3 matrix mapping original coords to coords after np.rot90(img, k) (counter-clockwise k times).""" |
| m = np.eye(3) |
| for _ in range(k % 4): |
| step = np.array([[0, 1, 0], [-1, 0, w], [0, 0, 1]], np.float64) |
| m = step @ m |
| w, h = h, w |
| return m[:2] |
|
|
