Automatic page splitting

copienator/fast_page_split.py

@@ -0,0 +1,276 @@
+"""Standard coarse-to-fine splitting: cheap positive content gate, precise fold.
+
+The whole-page gate only establishes obvious writing. All uncertain/sparse pages
+use the reference blank detector. Full-resolution fold processing is restricted
+to a central strip; edge-reaching overflow and ambiguous ownership fall back to
+the complete reference analysis. Full-spread analysis remains an internal fallback.
+"""
+
+from __future__ import annotations
+import cv2
+import numpy as np
+import pymupdf
+from copienator.auto_page_split import analyse_ink, _analyse_page_full, plan_masks
+
+
+def quick_page_content(rgb, dpi, target_dpi=75, *, native=True, _fallback=None):
+    """Positive-only coarse gate. Ambiguous/sparse pages use full detection."""
+    h, w = rgb.shape[:2]
+    small = cv2.resize(
+        rgb,
+        (round(w * target_dpi / dpi), round(h * target_dpi / dpi)),
+        interpolation=cv2.INTER_AREA,
+    )
+    px = target_dpi / 25.4
+    lo = small.min(axis=2)
+    ch = small.max(axis=2) - lo
+    gray = cv2.cvtColor(small, cv2.COLOR_RGB2GRAY)
+    strong = (ch > 100) | (gray < 95)
+    n, lab, stats, _ = cv2.connectedComponentsWithStats(strong.astype("uint8"), 8)
+    ok = np.zeros(n, bool)
+    x, y, bw, bh, area = stats[1:].T
+    # Only compact interior pen components can establish an obvious written page.
+    ok[1:] = (
+        (area >= 0.25 * px * px)
+        & (bw >= 0.35 * px)
+        & (bh >= 0.5 * px)
+        & (bw < 20 * px)
+        & (bh < 15 * px)
+        & (x > 0.09 * small.shape[1])
+        & (x + bw < 0.91 * small.shape[1])
+    )
+    selected = ok[lab]
+    grouped = cv2.dilate(
+        selected.astype("uint8"),
+        np.ones((max(1, round(px)), max(1, round(3 * px))), np.uint8),
+    )
+    ng, groups, gs, _ = cv2.connectedComponentsWithStats(grouped, 8)
+    comp_group = np.zeros(n, np.int32)
+    ys, xs = np.nonzero(selected)
+    comp_group[lab[ys, xs]] = groups[ys, xs]
+    counts = np.bincount(comp_group[ok], minlength=ng)
+    amounts = np.bincount(groups[selected], minlength=ng)
+    good = (
+        (counts >= 4)
+        & (amounts >= 3 * px * px)
+        & (gs[:, 2] >= 8 * px)
+        & (gs[:, 3] >= 2 * px)
+    )
+    good[0] = False
+    if good.any():
+        # Retain all strong ink for ownership support; the gate alone is interior-only.
+        return (
+            strong
+            if native
+            else cv2.resize(
+                strong.astype("uint8"), (w, h), interpolation=cv2.INTER_NEAREST
+            )
+            > 0
+        ), True
+    return (_fallback() if _fallback is not None else analyse_ink(rgb, dpi)[0]), False
+
+
+def _a4_plan(mask, gate, h, w, dpi, gate_dpi):
+    count = int(mask.sum())
+    return dict(
+        format="a4",
+        size_px=[w, h],
+        outputs=[
+            dict(
+                side="whole",
+                rect_px=[0, 0, w, h],
+                state="content" if count else "blank",
+                keep=bool(count),
+                ink_pixels=count,
+                content_dpi=gate_dpi if gate else dpi,
+            )
+        ],
+        crossings=[],
+        review_reasons=[],
+        extensions_mm=[0, 0],
+        analysis=dict(
+            mode="fast",
+            early_written=[gate],
+            fold_fraction=0,
+            gate_dpi=gate_dpi,
+            fallback=None,
+        ),
+    )
+
+
+def analyse_page_fast(rgb, dpi, kind, *, fold_fraction=0.1, gate_dpi=75):
+    """Fast analysis for an existing raster image."""
+    if kind not in {"a3", "a4"}:
+        return _analyse_page_full(rgb, dpi, kind)
+    gate_dpi = min(gate_dpi, dpi)
+    h, w = rgb.shape[:2]
+    if kind == "a4":
+        mask, gate = quick_page_content(rgb, dpi, gate_dpi, native=True)
+        return _a4_plan(mask, gate, h, w, dpi, gate_dpi)
+    middle = w // 2
+    halves = []
+    early = []
+    for a, b in ((0, middle), (middle, w)):
+        mask, gate = quick_page_content(rgb[:, a:b], dpi, gate_dpi, native=True)
+        halves.append(mask)
+        early.append(gate)
+    return _assemble_fast(
+        halves,
+        early,
+        h,
+        w,
+        dpi,
+        fold_fraction,
+        gate_dpi,
+        lambda x0, x1: rgb[:, x0:x1],
+        lambda: _analyse_page_full(rgb, dpi, kind),
+    )
+
+
+def _assemble_fast(
+    halves, early, h, w, dpi, fold_fraction, gate_dpi, render_roi, reference
+):
+    middle = w // 2
+    half = max(1, round(w * fold_fraction / 2))
+    x0 = max(0, middle - half)
+    x1 = min(w, middle + half)
+    roi = render_roi(x0, x1)
+    content = np.zeros(roi.shape[:2], bool)
+    for i, (a, b) in enumerate(((0, middle), (middle, w))):
+        lo = max(a, x0)
+        hi = min(b, x1)
+        mask = halves[i]
+        ys = np.arange(h) * mask.shape[0] // h
+        xs = (np.arange(lo, hi) - a) * mask.shape[1] // (b - a)
+        content[:, lo - x0 : hi - x0] = mask[ys[:, None], xs[None, :]]
+    ink, _ = analyse_ink(roi, dpi, spread=True)
+    remaining = {}
+    plan = plan_masks(
+        ink, dpi, "a3", content_ink=content, rgb=roi, _remaining_content=remaining
+    )
+    # If a proposed overflow reaches the strip's edge or ownership remains
+    # ambiguous, use the complete reference analysis rather than truncate it.
+    margin = round(2 * dpi / 25.4)
+    if (
+        plan["review_reasons"]
+        or plan["outputs"][0]["rect_px"][2] >= x1 - x0 - margin
+        or plan["outputs"][1]["rect_px"][0] <= margin
+    ):
+        plan = reference()
+        plan["analysis"] = dict(mode="fast", fallback="full-spread")
+        return plan
+    for i, (a, b) in enumerate(((0, middle), (middle, w))):
+        mask = halves[i]
+        side = "left" if i == 0 else "right"
+        if not early[i]:
+            lo = max(a, x0)
+            hi = min(b, x1)
+            mask[:, lo - a : hi - a] = remaining[side]
+            count = int(mask.sum())
+        else:
+            # Exact count of the implied nearest-neighbour full-resolution mask,
+            # without allocating it. Coarse masks establish obvious content only.
+            wy = np.diff(
+                np.ceil(np.arange(mask.shape[0] + 1) * h / mask.shape[0]).astype(
+                    "int64"
+                )
+            )
+            wx = np.diff(
+                np.ceil(np.arange(mask.shape[1] + 1) * (b - a) / mask.shape[1]).astype(
+                    "int64"
+                )
+            )
+            count = int(np.sum(mask * wy[:, None] * wx[None, :]))
+        o = plan["outputs"][i]
+        o["ink_pixels"] = count
+        o["keep"] = count > 0
+        o["state"] = "content" if count else "blank"
+        o["rect_px"] = (
+            [0, 0, o["rect_px"][2] + x0, h]
+            if i == 0
+            else [o["rect_px"][0] + x0, 0, w, h]
+        )
+    for c in plan["crossings"]:
+        c["box"][0] += x0
+        c["box"][2] += x0
+        for box in c.get("anchor_boxes", []):
+            if box is not None:
+                box[0] += x0
+                box[2] += x0
+    plan["size_px"] = [w, h]
+    plan["analysis"] = dict(
+        mode="fast",
+        gate_dpi=gate_dpi,
+        fold_fraction=fold_fraction,
+        early_written=early,
+        fallback=None,
+    )
+    return plan
+
+
+def analyse_pdf_page_fast(page, dpi, kind, *, fold_fraction=0.1, gate_dpi=75):
+    """Render a small overview and the precise fold directly from the PDF.
+
+    Full-resolution halves are rendered only when the positive overview gate
+    cannot establish writing. Raster clips are aligned to the same integer
+    pixel grid as a complete reference render, including rotated CropBoxes.
+    """
+    gate_dpi = min(gate_dpi, dpi)
+    bounds = (page.rect * pymupdf.Matrix(dpi / 72, dpi / 72)).irect
+    h, w = bounds.height, bounds.width
+    coarse_pix = page.get_pixmap(
+        dpi=round(gate_dpi), colorspace=pymupdf.csRGB, alpha=False
+    )
+    coarse = np.frombuffer(coarse_pix.samples, np.uint8).reshape(
+        coarse_pix.height, coarse_pix.width, 3
+    )
+
+    def render(x0, x1):
+        rect = pymupdf.Rect(x0 * 72 / dpi, 0, x1 * 72 / dpi, page.rect.height)
+        pix = page.get_pixmap(
+            dpi=round(dpi), clip=rect, colorspace=pymupdf.csRGB, alpha=False
+        )
+        data = np.frombuffer(pix.samples, np.uint8).reshape(pix.height, pix.width, 3)
+        out = np.full((h, x1 - x0, 3), 255, np.uint8)
+        left = max(x0, pix.x)
+        top = max(0, pix.y)
+        right = min(x1, pix.x + pix.width)
+        bottom = min(h, pix.y + pix.height)
+        out[top:bottom, left - x0 : right - x0] = data[
+            top - pix.y : bottom - pix.y, left - pix.x : right - pix.x
+        ]
+        return out
+
+    def reference():
+        return _analyse_page_full(render(0, w), dpi, kind)
+
+    if kind not in {"a3", "a4"}:
+        return reference(), coarse
+    if kind == "a4":
+        mask, gate = quick_page_content(
+            coarse,
+            gate_dpi,
+            gate_dpi,
+            native=True,
+            _fallback=lambda: analyse_ink(render(0, w), dpi)[0],
+        )
+        return _a4_plan(mask, gate, h, w, dpi, gate_dpi), coarse
+    middle = w // 2
+    cm = coarse.shape[1] // 2
+    halves = []
+    early = []
+    for a, b, ca, cb in ((0, middle, 0, cm), (middle, w, cm, coarse.shape[1])):
+        mask, gate = quick_page_content(
+            coarse[:, ca:cb],
+            gate_dpi,
+            gate_dpi,
+            native=True,
+            _fallback=lambda a=a, b=b: analyse_ink(render(a, b), dpi)[0],
+        )
+        halves.append(mask)
+        early.append(gate)
+    plan = _assemble_fast(
+        halves, early, h, w, dpi, fold_fraction, gate_dpi, render, reference
+    )
+    plan["analysis"]["rendering"] = "direct-pdf-roi"
+    return plan, coarse