old-books/_pipeline/pdftext.py

"""Read a PDF's text layer and image positions with poppler. No guessing: coordinates in, text out."""
import os
import re
import subprocess
import tempfile
import xml.etree.ElementTree as ET
from html import unescape


def _run(cmd):
    return subprocess.run(cmd, check=True, capture_output=True, text=True).stdout


def page_geometry(pdf):
    """[(width, height)] per page as displayed (rotation applied)."""
    n = int(re.search(r"^Pages:\s+(\d+)", _run(["pdfinfo", pdf]), re.M).group(1))
    out = _run(["pdfinfo", "-f", "1", "-l", str(n), pdf])
    sizes = {int(m[0]): (float(m[1]), float(m[2])) for m in re.findall(r"^Page\s+(\d+) size:\s+([\d.]+) x ([\d.]+)", out, re.M)}
    rots = {int(m[0]): int(m[1]) for m in re.findall(r"^Page\s+(\d+) rot:\s+(\d+)", out, re.M)}
    geo = []
    for p in range(1, n + 1):
        w, h = sizes[p]
        if rots.get(p, 0) in (90, 270):
            w, h = h, w
        geo.append((w, h))
    return geo


def halves(width, height):
    """A landscape page that holds two book pages is read as two units, left then right."""
    if width > height * 1.2:
        return [("L", 0, width / 2), ("R", width / 2, width / 2)]
    return [(None, 0, width)]


def chars_per_page(pdf):
    """Number of non-space characters in the text layer of each page."""
    n = len(page_geometry(pdf))
    out = []
    for p in range(1, n + 1):
        txt = _run(["pdftotext", "-f", str(p), "-l", str(p), pdf, "-"])
        out.append(len(re.sub(r"\s", "", txt)))
    return out


_LINE_RE = re.compile(r'<line xMin="([\d.]+)" yMin="([\d.]+)" xMax="([\d.]+)" yMax="([\d.]+)">(.*?)</line>', re.S)


def page_lines(pdf, page):
    """Text lines of one PDF page with their position in points, top to bottom:
    [{'x', 'y', 'x2', 'y2', 'text'}]. Coordinates are those of the displayed page (poppler bbox-layout)."""
    xml = _run(["pdftotext", "-bbox-layout", "-f", str(page), "-l", str(page), pdf, "-"])
    out = []
    for m in _LINE_RE.finditer(xml):
        words = [unescape(w) for w in re.findall(r">([^<]+)</word>", m.group(5))]
        out.append({"x": float(m.group(1)), "y": float(m.group(2)), "x2": float(m.group(3)),
                    "y2": float(m.group(4)), "text": " ".join(words)})
    return sorted(out, key=lambda ln: (ln["y"], ln["x"]))


def _key(text):
    return re.sub(r"[^A-Za-z0-9]", "", text)


def align(raw_lines, boxes, window=8):
    """Give each line of a layout text (pdftotext -layout) the position of the same line from page_lines.
    Returns a list parallel to raw_lines: the box dict, or None for blank lines and lines that cannot be
    matched one to one (merged columns, small-caps headers). A line that only partly matches a box is left
    as None: its position would be wrong."""
    result = [None] * len(raw_lines)
    j = 0
    for i, line in enumerate(raw_lines):
        k = _key(line)
        if not k:
            continue
        for jj in range(j, min(j + window, len(boxes))):
            b = _key(boxes[jj]["text"])
            if b == k:
                result[i] = boxes[jj]
                j = jj + 1
                break
            if len(b) >= 8 and (k.endswith(b) or k.startswith(b)):
                j = jj + 1
                break
    return result


def units(pdf):
    """Reading units in order: [{'pdf_page', 'half', 'text', 'geom'}]. Text keeps the layout (indentation, hyphens).
    'geom' is parallel to text.split('\\n'): the position of each line on the page (see align)."""
    result = []
    for p, (w, h) in enumerate(page_geometry(pdf), start=1):
        boxes = page_lines(pdf, p)
        for half, x, width in halves(w, h):
            text = _run(["pdftotext", "-layout", "-f", str(p), "-l", str(p),
                         "-x", str(int(x)), "-y", "0", "-W", str(int(width)), "-H", str(int(h)), pdf, "-"])
            text = text.replace("\f", "")
            mine = [b for b in boxes if x <= b["x"] < x + width] if half else boxes
            result.append({"pdf_page": p, "half": half, "text": text, "geom": align(text.split("\n"), mine)})
    return result


def images(pdf):
    """Images with their page, half and pixel size: [{'pdf_page', 'half', 'width', 'height'}]."""
    geo = page_geometry(pdf)
    found = []
    with tempfile.TemporaryDirectory() as tmp:
        subprocess.run(["pdftohtml", "-xml", "-q", pdf, os.path.join(tmp, "o")], check=True, capture_output=True)
        root = ET.parse(os.path.join(tmp, "o.xml")).getroot()
        for page in root.iter("page"):
            p = int(page.get("number"))
            w, h = geo[p - 1]
            scale = float(page.get("width")) / w
            for img in page.iter("image"):
                cx = (float(img.get("left")) + float(img.get("width")) / 2) / scale
                half = None
                for name, x, width in halves(w, h):
                    if x <= cx < x + width:
                        half = name
                found.append({"pdf_page": p, "half": half,
                              "width": round(float(img.get("width")) / scale),
                              "height": round(float(img.get("height")) / scale)})
    return found