Grok Bot
104 lines
4.1 KiB
Python
104 lines
4.1 KiB
Python
"""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")]
|