Add pipeline code

2026-10-17 15:27:14 -05:00
parent 3cd5a4c658
commit 5f0e2a9768
21 changed files with 1612 additions and 0 deletions
@@ -0,0 +1,103 @@
+"""Stage 1, ingest: copy the original PDF into its book folder (never modify it), record its sha256,
+find pages without a text layer and OCR only if there are any."""
+import hashlib
+import os
+import re
+import shutil
+import subprocess
+
+import bookio
+import pdftext
+import settings
+
+
+def sha256(path):
+    h = hashlib.sha256()
+    with open(path, "rb") as f:
+        for block in iter(lambda: f.read(1 << 20), b""):
+            h.update(block)
+    return h.hexdigest()
+
+
+def find_title(pdf):
+    out = subprocess.run(["pdfinfo", pdf], capture_output=True, text=True).stdout
+    m = re.search(r"^Title:\s+(.+)$", out, re.M)
+    return bookio.title_case(m.group(1).strip()) if m else ""
+
+
+def find_author(unit_texts):
+    """The author is the line 'By <Name>' on the first pages."""
+    for text in unit_texts:
+        m = re.search(r"^\s*By\s+([A-Z][A-Za-z.\- ]+?)\s*$", text, re.M)
+        if m:
+            return m.group(1).strip()
+    return ""
+
+
+def author_slug(author):
+    """'Glenn Clark' -> 'clark-glenn' (surname first, last word taken as surname)."""
+    parts = author.split()
+    return bookio.slug(" ".join([parts[-1]] + parts[:-1])) if parts else "unknown"
+
+
+def run(cfg, pdf, title=None, author=None):
+    """Returns the book folder path."""
+    pdf_units = pdftext.units(pdf)
+    author = author or find_author([u["text"] for u in pdf_units[:6]])
+    title = title or find_title(pdf)
+    if not author or not title:
+        raise SystemExit("could not find title/author in the PDF; pass --title and --author")
+    book_dir = os.path.join(settings.resolve(cfg, "books_dir"), author_slug(author), bookio.slug(title))
+    os.makedirs(book_dir, exist_ok=True)
+    copy = os.path.join(book_dir, os.path.basename(pdf))
+    if os.path.exists(copy) and sha256(copy) != sha256(pdf):
+        raise SystemExit("a different file with this name is already in the book folder: " + copy)
+    if not os.path.exists(copy):
+        shutil.copyfile(pdf, copy)
+    book = {}
+    if os.path.exists(os.path.join(book_dir, "book.json")):
+        book = bookio.read_book(book_dir)
+    counts = pdftext.chars_per_page(copy)
+    book.update({
+        "title": title,
+        "author": author,
+        "source_pdf": os.path.basename(pdf),
+        "sha256": sha256(copy),
+        "pdf_pages": len(counts),
+        "pages_without_text_layer": [i + 1 for i, c in enumerate(counts) if c < cfg["ocr"]["min_chars_per_page"]],
+    })
+    book.setdefault("coauthors", [])
+    book.setdefault("publisher", "")
+    book.setdefault("edition", "")
+    book.setdefault("series", "")
+    book.setdefault("summaries", {})
+    bookio.write_book(book_dir, book)
+    warnings = ocr_if_needed(cfg, book_dir, book)
+    bookio.write_warnings(book_dir, "ingest", warnings)
+    return book_dir
+
+
+def work_pdf(cfg, book_dir, book):
+    """The PDF text is read from: the OCR copy if one exists, else the untouched original."""
+    ocr_pdf = os.path.join(settings.resolve(cfg, "work_dir"), os.path.basename(book_dir), "ocr.pdf")
+    return ocr_pdf if os.path.exists(ocr_pdf) else os.path.join(book_dir, book["source_pdf"])
+
+
+def ocr_if_needed(cfg, book_dir, book):
+    """OCR a scratch copy (never the original) when pages lack a text layer. Returns warnings."""
+    missing = book["pages_without_text_layer"]
+    if not missing:
+        return []
+    where = "pdf pages " + ",".join(map(str, missing))
+    if not cfg["ocr"]["enabled"]:
+        return [("no-text-layer", where, "no text layer and OCR is disabled in settings")]
+    if shutil.which("ocrmypdf") is None:
+        return [("no-text-layer", where, "no text layer and ocrmypdf is not installed on this host; those pages have no text")]
+    out_dir = os.path.join(settings.resolve(cfg, "work_dir"), os.path.basename(book_dir))
+    os.makedirs(out_dir, exist_ok=True)
+    out = os.path.join(out_dir, "ocr.pdf")
+    src = os.path.join(book_dir, book["source_pdf"])
+    r = subprocess.run(["ocrmypdf", "--skip-text", "-l", cfg["ocr"]["language"], src, out], capture_output=True, text=True)
+    if r.returncode != 0:
+        return [("ocr-failed", where, "ocrmypdf exit %d: %s" % (r.returncode, r.stderr.strip()[-300:]))]
+    return [("ocr", where, "OCR applied to a scratch copy; original untouched")]