File size: 4,772 Bytes
4dce182 e24bf15 4dce182 e24bf15 4dce182 e24bf15 4dce182 e24bf15 4dce182 e24bf15 4dce182 | 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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | """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 # RGB uint8, in the page coordinate frame
width: int
height: int
text_source: str # "ocr" | "pdf_text"
pdf_segments: list = field(default_factory=list) # [(x0, y0, x1, y1, text)] when text_source == "pdf_text"
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)
# the 1-degree search runs at half size (a quarter of the pixels); the 0.1-degree refinement at full size
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)) # page colour from a pixel sample
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) # (x, y) -> (y, w - x)
m = step @ m
w, h = h, w
return m[:2]
|