parisaocr 0.2.0__py3-none-any.whl

parisaocr/__init__.py

@@ -0,0 +1,27 @@
+"""ParisaOCR: open-source OCR for printed Persian.
+
+A text-line detector (PP-OCRv6, run with ONNX Runtime) finds the lines of a
+page; a Kraken recognition model trained for Persian print reads each line; the
+lines are ordered into right-to-left columns and written as text, hOCR (with
+line and word boxes on the original page) or JSONL. Both models ship with the
+package, so it runs offline, on a CPU or a GPU.
+
+    parisaocr ocr page.png                      # print the text
+    parisaocr ocr book.pdf --out out            # out/txt, out/hocr
+    parisaocr ocr book.pdf --out out --format pdf   # out/pdf/book.pdf, searchable
+
+    from parisaocr import OCR
+    print(OCR().text("page.png"))
+
+The recognition model and the detector are replaceable (`--model`, `--detector`):
+Tesseract models and the Surya detector are supported for comparison, and
+`parisaocr cut` cuts the lines of scanned books for building training data.
+"""
+__version__ = "0.2.0"
+
+
+def __getattr__(name):  # the interface imports torch and Kraken; only load them when asked for
+    if name == "OCR":
+        from .api import OCR
+        return OCR
+    raise AttributeError(name)

parisaocr/__main__.py

@@ -0,0 +1,3 @@
+from .cli import main
+
+main()

parisaocr/api.py

@@ -0,0 +1,54 @@
+"""Python interface: read a page image and get its text or its lines.
+
+    from parisaocr import OCR
+    ocr = OCR()                       # loads the detector and the recognition model once
+    print(ocr.text("page.png"))       # the page's text, lines in reading order
+    for line in ocr.lines("page.png"):
+        print(line.bbox, line.conf, line.text)
+
+`OCR(model=..., detector=..., device=...)` takes the same model and detector
+specifications as the command line (`parisaocr ocr --help`).
+"""
+import pathlib
+import tempfile
+from types import SimpleNamespace
+
+from PIL import Image
+
+from . import output
+from .cli import build_engine
+from .order import order
+from .pipeline import recognize_batch
+
+
+class OCR:
+    def __init__(self, model="default", detector="ppocr:v6-small:1.3", cpu=False, line_order="rtl"):
+        opts = SimpleNamespace(model=model, detector=detector, cpu=cpu, min_width=1600, max_width=2000, batch=1,
+                               pad=6, pad_frac=0.4, min_height=40, block_tall=True, fallback=True, jobs=8)
+        self.detector, self.reader, self.options = build_engine(opts)
+        self.line_order = line_order
+
+    def columns(self, image):
+        """Lines of one page (a path or a PIL image) as reading-order columns of `pipeline.Line`."""
+        with tempfile.TemporaryDirectory(prefix="parisaocr-") as tmp:
+            tmp = pathlib.Path(tmp)
+            if not isinstance(image, (str, pathlib.Path)):
+                path = tmp / "page.png"
+                image.save(path)
+            else:
+                path = pathlib.Path(image)
+            with Image.open(path) as page:
+                im, f = self.detector.prepare(page)
+                rec = {"page": "page", "image": str(path), "width": im.width, "height": im.height, "scale": f,
+                       "orig_width": page.width, "orig_height": page.height}
+            rec["lines"] = self.detector.detect([im])[0]
+            lines = recognize_batch([rec], self.reader, self.options, tmp)["page"]
+        return order(lines, self.line_order)
+
+    def lines(self, image):
+        """The lines of one page in reading order (text, bbox on the page, mean confidence, words)."""
+        return [line for col in self.columns(image) for line in col]
+
+    def text(self, image):
+        """The text of one page: one line per printed line, columns right to left."""
+        return output.text(self.columns(image))

parisaocr/bidi.py

@@ -0,0 +1,37 @@
+"""Reading order <-> display order for one text line, with each line's own base direction.
+
+Recognition models see a line image from left to right, so they are trained on
+the display order of the text and their output is turned back into reading
+order. Kraken can do both itself, but (mittagessen/kraken#809) its bidi code
+deletes ZWNJ, the Persian half-space, and it gives every line the same base
+direction, so a Latin-only line such as "1. Vercingétorix" is trained as
+"Vercingétorix .1". Here the Unicode bidi algorithm of kraken.lib.bidi runs
+with ZWNJ protected (swapped for U+00A6, a neutral that stays in place between
+two letters) and with the base direction taken from the line itself.
+"""
+import unicodedata
+
+ZWNJ, PLACEHOLDER = "‌", "¦"
+
+
+def base_direction(text):
+    """"L" for a line with more Latin than Arabic-script letters, else "R"."""
+    latin = sum(1 for c in text if c.isalpha() and unicodedata.bidirectional(c) == "L")
+    arabic = sum(1 for c in text if unicodedata.bidirectional(c) in ("R", "AL"))
+    return "L" if latin > arabic else "R"
+
+
+def to_display(text, base=None):
+    """Display (left-to-right visual) order of a line in reading order; ZWNJ kept."""
+    from kraken.lib.bidi import get_display
+    base = base or base_direction(text)
+    return get_display(text.replace(ZWNJ, PLACEHOLDER), base_dir=base).replace(PLACEHOLDER, ZWNJ)
+
+
+def to_logical(display, base=None):
+    """Reading order of a line in display order, and for each character of the result the index
+    of the display character it came from (to carry character positions along)."""
+    from kraken.lib.bidi import get_display_map
+    base = base or base_direction(display)
+    text, order = get_display_map(display.replace(ZWNJ, PLACEHOLDER), base_dir=base)
+    return text.replace(PLACEHOLDER, ZWNJ), order

parisaocr/cli.py

@@ -0,0 +1,256 @@
+"""Command line: `parisaocr ocr` and `parisaocr cut`."""
+import argparse
+import os
+import pathlib
+import shlex
+import sys
+import tempfile
+import time
+from concurrent.futures import ThreadPoolExecutor
+
+from . import __version__, output
+from .detect import detect_pages
+from .detectors import make_detector
+from .order import order
+from .pages import collect
+from .pipeline import Options, recognize_batch
+from .reader import TesseractReader
+
+ROOT = pathlib.Path(__file__).resolve().parent.parent
+BUNDLED_MODEL = pathlib.Path(__file__).resolve().parent / "models" / "parisaocr-fa-0.1.safetensors"
+DEFAULT_TESSDATA = os.environ.get("PARISAOCR_TESSDATA", str(ROOT / "dist/tessdata_contrib/fas_print/best"))
+DEFAULT_LANG = os.environ.get("PARISAOCR_LANG", "fas_print")
+SMALL_TEXT_PX = 14  # median line height (original page pixels) below which accuracy drops clearly
+
+
+def engine_args(p, default_model="default"):
+    g = p.add_argument_group("engine")
+    g.add_argument("--model", default=default_model,
+                   help="recognition model: default (the ParisaOCR model shipped with the package), "
+                        "kraken:MODEL.safetensors (another Kraken model), or TESSDATA_DIR:LANG[:EXTRA ARGS] for Tesseract "
+                        f"(the cut command defaults to the fas_print Tesseract model, {DEFAULT_TESSDATA}:{DEFAULT_LANG})")
+    g.add_argument("--detector", default="ppocr:v6-small:1.3",
+                   help="line detector: ppocr[:VERSION-SIZE[:UNCLIP]] (PaddleOCR weights via RapidOCR on ONNX Runtime, "
+                        "Apache-2.0; default ppocr:v6-small:1.3), surya (needs the surya extra; its weights are free "
+                        "only for research, personal use and small companies), kraken (Kraken's default segmenter)")
+    g.add_argument("--min-width", type=int, default=1600, help="upscale narrower pages to this width before detection")
+    g.add_argument("--max-width", type=int, default=2000, help="reduce wider pages to this width for detection only (0 = never)")
+    g.add_argument("--batch", type=int, default=8, help="pages per detector batch")
+    g.add_argument("--cpu", action="store_true", help="run the detector on the CPU")
+    g.add_argument("--pad", type=int, default=6, help="least margin around a line crop, in working pixels (vertical: half)")
+    g.add_argument("--pad-frac", type=float, default=0.4,
+                   help="margin as a fraction of the page's median line height; the larger of --pad and this is used")
+    g.add_argument("--min-height", type=int, default=40, help="enlarge crops shorter than this many pixels")
+    g.add_argument("--no-block-tall", dest="block_tall", action="store_false", help="read tall boxes in raw-line mode too")
+    g.add_argument("--no-fallback", dest="fallback", action="store_false", help="no single-line retry of near-empty results")
+    g.add_argument("--jobs", type=int, default=os.cpu_count(), help="parallel Tesseract processes")
+    g.add_argument("--redo", action="store_true", help="ignore cached detection and existing output")
+
+
+def build_engine(opts):
+    if opts.model == "default" or opts.model.startswith("kraken:"):
+        import torch
+        from .kraken_reader import KrakenReader
+        path = BUNDLED_MODEL if opts.model == "default" else opts.model.split(":", 1)[1]
+        reader = KrakenReader(path, "cuda:0" if torch.cuda.is_available() and not opts.cpu else "cpu")
+        fallback, short_conf = False, 80.0  # the single-line retry is a Tesseract remedy
+    else:
+        tessdata, lang, *extra = opts.model.split(":", 2)
+        reader = TesseractReader(tessdata, lang, shlex.split(extra[0]) if extra else [])
+        fallback, short_conf = opts.fallback, 0.0
+    detector = make_detector(opts.detector, opts.min_width, opts.max_width, opts.batch, opts.cpu)
+    options = Options(opts.pad, opts.pad_frac, opts.min_height, opts.block_tall, fallback, opts.jobs, short_conf)
+    return detector, reader, options
+
+
+def small_text_warning(rec):
+    """A note on stderr when a page's text is so small that accuracy drops (median line height in page pixels)."""
+    heights = sorted((l["bbox"][3] - l["bbox"][1]) / rec.get("scale", 1.0) for l in rec["lines"])
+    if heights and heights[len(heights) // 2] < SMALL_TEXT_PX:
+        print(f"parisaocr: {rec['page']}: text lines are only ~{heights[len(heights) // 2]:.0f} px high; "
+              "accuracy drops on such small text (a 300 dpi scan of the page helps)", file=sys.stderr, flush=True)
+
+
+def cmd_ocr(opts):
+    if opts.out is None:  # no output directory: read into a temporary one and print the text
+        with tempfile.TemporaryDirectory(prefix="parisaocr-out-") as tmp:
+            opts.out, opts.format, opts.quiet = tmp, "txt", True
+            pages = cmd_ocr(opts)
+            for i, (pid, _) in enumerate(pages):
+                if len(pages) > 1:
+                    print(("\n" if i else "") + f"# {pid}")
+                sys.stdout.write((pathlib.Path(tmp) / "txt" / f"{pid}.txt").read_text(encoding="utf-8"))
+        return
+    log = (lambda *a, **k: None) if getattr(opts, "quiet", False) else (lambda *a, **k: print(*a, **k, flush=True))
+    out = pathlib.Path(opts.out)
+    formats = [f.strip() for f in opts.format.split(",") if f.strip()]
+    bad = set(formats) - {"txt", "hocr", "jsonl", "pdf"}
+    if bad:
+        sys.exit(f"parisaocr: unknown format {', '.join(bad)}")
+    pages = collect(opts.input, out, opts.pdf, opts.dpi, opts.first, opts.last, opts.redo)
+    if not pages:
+        sys.exit("parisaocr: no pages to read")
+    want_pdf = "pdf" in formats
+    formats = [f for f in formats if f != "pdf"]
+    if want_pdf and "jsonl" not in formats:
+        formats.append("jsonl")  # the searchable PDF is built from the per-page lines and boxes
+    for f in formats:
+        (out / f).mkdir(parents=True, exist_ok=True)
+    todo = [(pid, p) for pid, p in pages if opts.redo or not all((out / f / f"{pid}.{f}").exists() for f in formats)]
+    log(f"{len(pages)} pages, {len(todo)} to read -> {out}/{{{','.join(formats + ['pdf'] * want_pdf)}}}")
+    if not todo:
+        if want_pdf:
+            write_pdfs(opts, out, pages, log)
+        return pages
+    detector, reader, options = build_engine(opts)
+    started, n_lines, done = time.time(), 0, 0
+
+    def finish(records, lines_of):
+        n = 0
+        for rec in records:
+            pid = rec["page"]
+            small_text_warning(rec)
+            cols = order(lines_of[pid], opts.order)
+            n += sum(len(c) for c in cols)
+            if "txt" in formats:
+                (out / "txt" / f"{pid}.txt").write_text(output.text(cols), encoding="utf-8")
+            if "hocr" in formats:
+                (out / "hocr" / f"{pid}.hocr").write_text(
+                    output.hocr(pid, pathlib.Path(rec["image"]).name, rec.get("orig_width", rec["width"]),
+                                rec.get("orig_height", rec["height"]), cols, reader.describe()), encoding="utf-8")
+            if "jsonl" in formats:
+                (out / "jsonl" / f"{pid}.jsonl").write_text(output.jsonl(pid, cols), encoding="utf-8")
+        return n
+
+    with tempfile.TemporaryDirectory(prefix="parisaocr-") as tmp, ThreadPoolExecutor(1) as cpu:
+        tmp = pathlib.Path(tmp)
+        pending = None
+        # Detection (GPU) of one batch overlaps with recognition (CPU) of the previous one.
+        for records in detect_pages(detector, todo, out / "det", opts.redo):
+            if pending:
+                n_lines += pending.result()
+                done += opts.batch
+                log(f"  {min(done, len(todo))}/{len(todo)} pages, {n_lines} lines, {time.time() - started:.0f} s")
+            pending = cpu.submit(lambda recs=records: finish(recs, recognize_batch(recs, reader, options, tmp)))
+        if pending:
+            n_lines += pending.result()
+    log(f"{len(todo)} pages, {n_lines} lines in {time.time() - started:.0f} s")
+    if want_pdf:
+        write_pdfs(opts, out, pages, log)
+    return pages
+
+
+def write_pdfs(opts, out, pages, log):
+    """Searchable PDFs in OUT/pdf: the input PDF with a text layer, and/or the image inputs as PDFs."""
+    from PIL import Image
+    from . import pdfout
+    (out / "pdf").mkdir(exist_ok=True)
+    pdf_inputs = [pathlib.Path(p) for p in opts.input if pathlib.Path(p).suffix.lower() == ".pdf"]
+    pages_dir = (out / "pages").resolve()
+    from_pdf = [(pid, p) for pid, p in pages if pathlib.Path(p).resolve().parent == pages_dir]
+    images = [(pid, p) for pid, p in pages if (pid, p) not in from_pdf]
+
+    def lines_of(pid):
+        f = out / "jsonl" / f"{pid}.jsonl"
+        return pdfout.read_jsonl(f) if f.exists() else []
+
+    if pdf_inputs and from_pdf:
+        src = pdf_inputs[0]
+        layers = {}
+        for pid, image in from_pdf:
+            with Image.open(image) as im:
+                layers[int(pid.split("-")[1])] = (im.width, im.height, lines_of(pid))
+        dst = out / "pdf" / src.name
+        done = pdfout.overlay_pdf(src, layers, dst, keep_text=opts.pdf_text == "skip")
+        skipped = len(layers) - len(done)
+        log(f"searchable PDF -> {dst} (text layer on {len(done)} pages"
+            + (f"; {skipped} already had text and were left as they are (--pdf-text add to add ours)" if skipped else "")
+            + ")")
+    if images:
+        if opts.merge_pdf:
+            dst = out / "pdf" / (opts.merge_pdf if opts.merge_pdf.lower().endswith(".pdf") else opts.merge_pdf + ".pdf")
+            pdfout.images_to_pdf([(p, lines_of(pid)) for pid, p in images], dst)
+            log(f"searchable PDF -> {dst} ({len(images)} pages)")
+        else:
+            for pid, p in images:
+                pdfout.images_to_pdf([(p, lines_of(pid))], out / "pdf" / f"{pid}.pdf")
+            log(f"searchable PDFs -> {out / 'pdf'} ({len(images)} files)")
+
+
+def cmd_pages(opts):
+    """Only turn a PDF into page images (OUT/p-NNN.png), e.g. to prepare a book for `parisaocr cut`."""
+    from .pages import pdf_pages
+    out = pathlib.Path(opts.out)
+    got = pdf_pages(pathlib.Path(opts.pdf), out, opts.pdf_mode, opts.dpi, opts.first, opts.last, opts.redo)
+    print(f"{len(got)} pages in {out}")
+
+
+def cmd_cut(opts):
+    from .cut import cut
+    detector, reader, options = build_engine(opts)
+    with tempfile.TemporaryDirectory(prefix="parisaocr-") as tmp:
+        cut(opts.book_dir, opts.out or pathlib.Path(opts.book_dir) / "lines", detector, reader, options,
+            test=opts.test, train=opts.train, exclude=opts.exclude, seed=opts.seed, min_lines=opts.min_lines,
+            min_page_conf=opts.min_page_conf, min_conf=opts.min_conf, pad=opts.crop_pad, redo=opts.redo, tmp=pathlib.Path(tmp),
+            test_pages=opts.test_pages, train_pages=opts.train_pages)
+
+
+def main(argv=None):
+    ap = argparse.ArgumentParser(prog="parisaocr", description=__doc__)
+    ap.add_argument("--version", action="version", version=f"parisaocr {__version__}")
+    sub = ap.add_subparsers(dest="cmd", required=True)
+
+    o = sub.add_parser("ocr", help="read pages (images, a directory of images, or one PDF)",
+                       description="Writes OUT/txt/PAGE.txt, OUT/hocr/PAGE.hocr and/or OUT/jsonl/PAGE.jsonl; "
+                                   "PDF pages are rendered to OUT/pages first; detection is cached in OUT/det.")
+    o.add_argument("input", nargs="+")
+    o.add_argument("--out", help="output directory; without it the text is printed")
+    o.add_argument("--format", default="txt,hocr",
+                   help="comma-separated: txt, hocr, jsonl, pdf (with --out). pdf: searchable PDF, the input PDF with "
+                        "an invisible text layer, or one PDF per input image (see --merge-pdf)")
+    o.add_argument("--merge-pdf", metavar="NAME", help="with --format pdf: put all input images into one PDF, OUT/pdf/NAME.pdf")
+    o.add_argument("--pdf-text", choices=["skip", "add"], default="skip",
+                   help="with --format pdf and a PDF input: leave pages that already have a text layer alone (skip), "
+                        "or add ours to them too (add; e.g. over a poor earlier OCR layer)")
+    o.add_argument("--order", choices=["rtl", "raster"], default="rtl", help="line order: right-to-left columns, or top to bottom")
+    o.add_argument("--pdf", choices=["auto", "extract", "render"], default="auto",
+                   help="PDF pages: extract the embedded scan images, render at --dpi, or decide per file")
+    o.add_argument("--dpi", type=int, default=300)
+    o.add_argument("--first", type=int, help="first PDF page")
+    o.add_argument("--last", type=int, help="last PDF page")
+    engine_args(o)
+    o.set_defaults(func=cmd_ocr)
+
+    p = sub.add_parser("pages", help="extract or render a PDF's pages to OUT/p-NNN.png (no OCR)")
+    p.add_argument("pdf")
+    p.add_argument("--out", required=True)
+    p.add_argument("--pdf-mode", choices=["auto", "extract", "render"], default="auto")
+    p.add_argument("--dpi", type=int, default=300)
+    p.add_argument("--first", type=int)
+    p.add_argument("--last", type=int)
+    p.add_argument("--redo", action="store_true")
+    p.set_defaults(func=cmd_pages)
+
+    c = sub.add_parser("cut", help="cut the lines of a book (BOOK_DIR/p-NNN.png) for labelling",
+                       description="Writes LINES/test|train/p-NNN_LL.png crops and LINES/manifest.json (default LINES = BOOK_DIR/lines).")
+    c.add_argument("book_dir")
+    c.add_argument("--out", help="lines directory (default BOOK_DIR/lines)")
+    c.add_argument("--test", type=int, default=12, help="test pages")
+    c.add_argument("--train", type=int, default=24, help="train pages")
+    c.add_argument("--exclude", default="", help="page ranges to leave out, e.g. 1-11,93-97")
+    c.add_argument("--test-pages", default="", help="use exactly these pages as test, e.g. 30,45 (no random choice, no filters)")
+    c.add_argument("--train-pages", default="", help="use exactly these pages as train, e.g. 18,50-52")
+    c.add_argument("--seed", type=int, default=0)
+    c.add_argument("--min-lines", type=int, default=15, help="pages with fewer detected lines are not candidates")
+    c.add_argument("--min-page-conf", type=float, default=80, help="skip pages whose mean word confidence is lower (0 = keep all)")
+    c.add_argument("--min-conf", type=float, default=60, help="draft words below this confidence are listed for checking")
+    c.add_argument("--crop-pad", type=int, default=10, help="margin of the saved line crops, in page pixels (vertical: half)")
+    engine_args(c, default_model=f"{DEFAULT_TESSDATA}:{DEFAULT_LANG}")  # page filters are calibrated on Tesseract confidences
+    c.set_defaults(func=cmd_cut)
+
+    opts = ap.parse_args(argv)
+    opts.func(opts)
+
+
+if __name__ == "__main__":
+    main()

parisaocr/cut.py

@@ -0,0 +1,113 @@
+"""Cut the text lines of a scanned book for labelling (replaces `correct_lines.py cut`).
+
+Pages are BOOK_DIR/p-NNN.png. All pages are detected (cached in BOOK_DIR/det);
+pages with too few lines or in --exclude are dropped, the rest are shuffled
+(--seed) and read one batch at a time until --test and --train pages are
+accepted; a page whose mean word confidence is below --min-page-conf (tables,
+maps, foreign text) is skipped. With --test-pages/--train-pages the given
+pages are used as they are. Every detected line of an accepted page is
+cropped from the original page at native resolution with a small margin into
+LINES/SPLIT/p-NNN_LL.png (LL = line index in detection order) and described in
+LINES/manifest.json with the draft reading and its low-confidence words, in the
+format scripts/correct_lines.py serve, line_sheets.py and dataset.py read.
+"""
+import json
+import pathlib
+import random
+import statistics
+import sys
+
+from PIL import Image
+
+from .detect import detect_pages
+from .pipeline import expand, line_crop, margins, recognize_batch
+
+
+def parse_ranges(spec):
+    pages = set()
+    for part in filter(None, spec.split(",")):
+        a, _, b = part.partition("-")
+        pages.update(range(int(a), int(b or a) + 1))
+    return pages
+
+
+def cut(book_dir, out, detector, reader, opts, test=12, train=24, exclude="", seed=0,
+        min_lines=15, min_page_conf=80.0, min_conf=60.0, pad=10, redo=False, tmp=None,
+        test_pages="", train_pages=""):
+    book_dir, out = pathlib.Path(book_dir), pathlib.Path(out)
+    pages = [(p.stem, p) for p in sorted(book_dir.glob("p-*.png"))]
+    if not pages:
+        sys.exit(f"parisaocr cut: no p-NNN.png pages in {book_dir}")
+    # Explicit pages (--test-pages/--train-pages) are taken as given, without the line-count
+    # and confidence filters; otherwise pages are drawn at random until the quotas are met.
+    explicit = {f"p-{p:03d}": split for split, spec in (("test", test_pages), ("train", train_pages)) for p in parse_ranges(spec)}
+    missing = [p for p in explicit if p not in dict(pages)]
+    if missing:
+        sys.exit(f"parisaocr cut: no such pages in {book_dir}: {' '.join(missing)}")
+    if explicit:
+        pages = [(pid, path) for pid, path in pages if pid in explicit]
+    records = {}
+    for batch in detect_pages(detector, pages, book_dir / "det", redo):
+        records.update({r["page"]: r for r in batch})
+        print(f"detected {len(records)}/{len(pages)} pages", end="\r", flush=True)
+    print()
+    if explicit:
+        order = [pid for pid, _ in pages]
+        print(f"{len(order)} pages given explicitly")
+    else:
+        excluded = parse_ranges(exclude)
+        order = [pid for pid, _ in pages if len(records[pid]["lines"]) >= min_lines and int(pid.split("-")[1]) not in excluded]
+        random.Random(seed).shuffle(order)
+        print(f"{len(order)} of {len(pages)} pages have at least {min_lines} lines and are not excluded")
+
+    manifest, chosen, skipped = [], {}, []
+    quota = {"test": test, "train": train}
+    i = 0
+    while i < len(order) and (explicit or sum(quota.values()) > 0):
+        batch = order[i:i + detector.batch]
+        i += len(batch)
+        lines_of = recognize_batch([records[p] for p in batch], reader, opts, tmp)
+        for pid in batch:
+            lines = lines_of[pid]
+            if explicit:
+                split = explicit[pid]
+            else:
+                if sum(quota.values()) == 0:
+                    break  # the rest of the batch was read for nothing; a batch is at most 8 pages
+                words = [w for l in lines for w in l.words]
+                mean = sum(w.conf for w in words) / len(words) if words else 0.0
+                if mean < min_page_conf:
+                    skipped.append((pid, round(mean)))
+                    continue
+                split = "test" if quota["test"] > 0 else "train"
+                quota[split] -= 1
+            chosen[pid] = split
+            (out / split).mkdir(parents=True, exist_ok=True)
+            with Image.open(records[pid]["image"]) as page:
+                im = page.convert("L") if page.mode not in ("L", "1") else page.copy()
+                # The saved crop gets the same margin rule as the recognizer's crop, in page pixels,
+                # so training lines look like what the model sees at inference.
+                heights = [l.bbox[3] - l.bbox[1] for l in lines]
+                px, py = margins(pad, opts.pad_frac, statistics.median(heights) if heights else 0)
+                boxes = [l.bbox for l in lines]
+                rects = expand(boxes, px, py, im.width, im.height)
+                if im.mode == "1":
+                    im = im.convert("L")  # so neighbours can be painted over
+                for k, (line, rect) in enumerate(zip(lines, rects)):  # detection order; a block box read as several lines gives several names
+                    name = f"{pid}_{k:02d}"
+                    line_crop(im, rect, boxes[k], boxes).save(out / split / f"{name}.png")
+                    manifest.append({"name": name, "split": split, "page": pid, "draft": line.text,
+                                     "check": [[w.text, None] for w in line.words if w.conf < min_conf],
+                                     "bbox": list(line.bbox), "conf": line.conf})
+    manifest.sort(key=lambda m: (m["split"] != "test", m["name"]))
+    out.mkdir(parents=True, exist_ok=True)
+    (out / "manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=0), encoding="utf-8")
+    for split in ("test", "train"):
+        ps = sorted(p for p, s in chosen.items() if s == split)
+        print(f"{split}: {len(ps)} pages, {sum(m['split'] == split for m in manifest)} lines: " + " ".join(p[2:] for p in ps))
+    if skipped:
+        print(f"skipped {len(skipped)} pages below {min_page_conf:.0f} mean confidence: "
+              + " ".join(f"{p[2:]}({c})" for p, c in skipped[:20]) + (" ..." if len(skipped) > 20 else ""))
+    if not explicit and sum(quota.values()):
+        print(f"short of {quota['test']} test and {quota['train']} train pages", file=sys.stderr)
+    return manifest

parisaocr/detect.py

@@ -0,0 +1,87 @@
+"""Text-line detection. `LineDetector` wraps Surya; replace it to try another detector.
+
+A detection record (cached as JSON) holds the page image path, the size of the
+image the detector saw (`width`, `height`), the factor it was scaled by
+(`scale`, 1.0 when the page was left alone) and the line boxes in those
+detection coordinates. Small pages are upscaled to `min_width` first: the
+detector misses lines on low-resolution scans otherwise; pages wider than
+`max_width` are reduced for detection only.
+"""
+import json
+import os
+import pathlib
+
+from PIL import Image
+
+
+class LineDetector:
+    def __init__(self, min_width=1600, max_width=2000, batch=8, cpu=False):
+        self.min_width, self.max_width, self.batch, self.cpu = min_width, max_width, batch, cpu
+        if cpu:
+            os.environ["TORCH_DEVICE"] = "cpu"
+        self._model = None
+
+    def model(self):
+        if self._model is None:
+            from surya.detection import DetectionPredictor  # slow import, only when detection is needed
+            # local(): this process owns the model. The default constructor talks to a shared
+            # server subprocess that outlives the client (it was found still holding GPU memory
+            # minutes after a run) and that crashed on batches of large pages.
+            self._model = DetectionPredictor.local()
+        return self._model
+
+    def prepare(self, image):
+        """The RGB image the detector will see and the factor it was scaled by.
+
+        Small pages are enlarged (the detector misses lines on them), very large
+        scans are reduced (the detector works at a fixed internal size anyway,
+        and a batch of 400-dpi pages overflows its server).
+        """
+        im = image.convert("RGB")
+        f = 1.0
+        if 0 < im.width < self.min_width:
+            f = self.min_width / im.width
+        elif self.max_width and im.width > self.max_width:
+            f = self.max_width / im.width
+        if f != 1.0:
+            im = im.resize((round(im.width * f), round(im.height * f)), Image.LANCZOS)
+        return im, f
+
+    def detect(self, images):
+        """Line boxes for each RGB image, as dicts with bbox, polygon, confidence (detection coordinates)."""
+        out = []
+        for i in range(0, len(images), self.batch):
+            for res in self.model()(images[i:i + self.batch]):
+                out.append([{"bbox": [round(v) for v in b.bbox],
+                             "polygon": [[round(x), round(y)] for x, y in b.polygon],
+                             "confidence": round(float(b.confidence), 3)} for b in res.bboxes])
+        return out
+
+
+def detect_pages(detector, pages, cache_dir=None, redo=False):
+    """Detection records for (page id, image path) pairs, batch by batch, using CACHE_DIR/PID.json when present."""
+    if cache_dir:
+        cache_dir = pathlib.Path(cache_dir)
+        cache_dir.mkdir(parents=True, exist_ok=True)
+    for i in range(0, len(pages), detector.batch):
+        chunk, records, todo = pages[i:i + detector.batch], {}, []
+        for pid, path in chunk:
+            cached = cache_dir / f"{pid}.json" if cache_dir else None
+            if cached and cached.exists() and not redo:
+                records[pid] = json.loads(cached.read_text(encoding="utf-8"))
+            else:
+                todo.append((pid, path))
+        if todo:
+            images, recs = [], []
+            for pid, path in todo:
+                with Image.open(path) as page:
+                    im, f = detector.prepare(page)
+                    recs.append({"page": pid, "image": str(path), "width": im.width, "height": im.height, "scale": f,
+                                 "orig_width": page.width, "orig_height": page.height})
+                images.append(im)
+            for rec, lines in zip(recs, detector.detect(images)):
+                rec["lines"] = lines
+                records[rec["page"]] = rec
+                if cache_dir:
+                    (cache_dir / f"{rec['page']}.json").write_text(json.dumps(rec, ensure_ascii=False), encoding="utf-8")
+        yield [records[pid] for pid, _ in chunk]