Deploy Dataset-Maker: torn-page non-overlapping dataset generator

.github/workflows/python-app.yml

@@ -0,0 +1,39 @@
+# This workflow will install Python dependencies, run tests and lint with a single version of Python
+# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python
+
+name: Python application
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+
+permissions:
+  contents: read
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python 3.10
+      uses: actions/setup-python@v3
+      with:
+        python-version: "3.10"
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install flake8 pytest
+        if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
+    - name: Lint with flake8
+      run: |
+        # stop the build if there are Python syntax errors or undefined names
+        flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
+        # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
+        flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
+    - name: Test with pytest
+      run: |
+        pytest

.gitignore

@@ -0,0 +1,13 @@
+__pycache__/
+*.py[cod]
+*.egg-info/
+.venv/
+venv/
+.env
+.DS_Store
+*.zip
+*.pdf
+tmp/
+.gradio/
+flagged/
+.pytest_cache/

CLAUDE.md

@@ -0,0 +1,39 @@
+# Dataset-Maker - agent guide
+
+Gradio web app (HuggingFace Spaces, free tier) that tears PDF pages into
+**non-overlapping** torn fragments for image-stitching datasets.
+
+## Run / test
+```bash
+pip install -r requirements.txt
+python app.py        # serves on :7860
+pytest -q            # tests/ : partition invariants + queue ordering
+```
+`requirements.txt` is pinned for HF Spaces' Python 3.10. Locally on newer
+Pythons you may need unpinned numpy/scipy/pillow/pymupdf.
+
+## Architecture (data flow)
+`app.py` → `src/pipeline.py` → `pdf_loader` → priority `queue_manager` →
+`tearing` → `packager`.
+
+- **`src/tearing.py` is the core.** Invariant: output is a strict PARTITION of
+  the page (every pixel in exactly one piece → no overlap). Implemented as
+  nearest-seed `argmin` (Voronoi) over a value-noise **domain-warped** pixel
+  grid. Warp the *query coords*, never the partition rule, or you break the
+  no-overlap guarantee. `verify_partition()` gates this at runtime.
+- **`src/queue_manager.py`** - binary min-heap priority queue. Documented Big-O
+  in the module docstring; keep `push`/`pop` at `O(log n)`.
+- **`src/noise.py` / `src/sampling.py`** - pure NumPy, no Perlin/extra deps.
+- Per-page seed = `master_seed*1_000_003 + page_index` → randomness changes per
+  page yet stays reproducible. Don't make it global.
+
+## Conventions
+- Keep `src/` UI-free (no `gradio` imports) so it stays testable. Gradio lives
+  only in `app.py`.
+- Images are `(H, W, 3)` uint8 RGB end-to-end; pieces use black background.
+- `manifest.json` `(x, y)` offsets ARE the stitching labels - don't change the
+  schema without updating `packager.py` README.txt + tests.
+
+## Deploy
+HF Spaces reads the YAML header in `README.md` (`sdk: gradio`, `app_file`).
+Push to the Space repo; Spaces installs `requirements.txt` and runs `app.py`.

README.md

@@ -1,13 +1,121 @@
 ---
 title: Dataset Maker
-emoji: 🚀
-colorFrom: red
-colorTo: gray
+emoji: 🧩
+colorFrom: indigo
+colorTo: purple
 sdk: gradio
-sdk_version: 6.16.0
-python_version: '3.13'
+sdk_version: 4.44.1
 app_file: app.py
 pinned: false
+license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 🧩 Dataset-Maker
+
+Turn any PDF into a **non-overlapping torn-fragment dataset** for image
+stitching / fragment-reassembly research. Each page is rendered to A4, torn into
+irregular organic pieces on a black background, and exported as a ZIP with
+stitching ground truth.
+
+Built to run on the **HuggingFace Spaces free tier** (Gradio SDK).
+
+## The guarantee: zero overlap
+
+The pieces form a strict **partition** of the page - every pixel belongs to
+**exactly one** fragment, so pieces never overlap and together cover the page
+exactly. That makes the `(x, y)` offset of each piece an exact stitching label.
+
+How: instead of drawing blobs (which overlap), we assign each pixel to its
+**nearest seed point** (`argmin`). That is a Voronoi tessellation - already a
+perfect partition. To make edges look *torn* instead of straight, we
+**domain-warp** the pixel grid with value noise *before* the nearest-seed test.
+Warping the query points (not the rule) keeps the result a partition while
+boundaries become jagged and organic.
+
+Inspired by Voronoi-jigsaw / eroded-boundary fragment literature:
+[Eroded Boundaries](https://arxiv.org/pdf/1912.00755),
+[Pairwise Irregular Fragments](https://arxiv.org/pdf/2507.09767),
+[Deepzzle](https://arxiv.org/abs/2005.12548).
+
+## Pipeline
+
+```
+PDF ──▶ render @DPI ──▶ fit/slice to A4 ──▶ priority queue (cheap pages first)
+                                                   │
+                            per-page seed ─────────▼
+                          domain-warped Voronoi partition (no overlap)
+                                                   │
+                            crop to bbox, black bg │
+                                                   ▼
+                       PNG-optimize ──▶ ZIP + manifest.json (x,y labels)
+```
+
+## Performance notes
+
+* **Gradio queue** - `demo.queue(max_size, default_concurrency_limit)` caps load
+  for the 2-vCPU free tier; the heavy event has its own `concurrency_limit`.
+* **Priority queue** (`src/queue_manager.py`) - binary min-heap, `push`/`pop`
+  `O(log n)`, `peek` `Θ(1)`, space `Θ(n)`. Orders page jobs cheap-first so the
+  user sees early progress and tail latency drops. Best case `Ω(1)` per op; no
+  comparison heap beats `Ω(log n)` amortized under interleaved ops.
+* **Vectorized partition** - SciPy `cKDTree` nearest-seed query is
+  `O(H·W·log S)`; mask extraction is `Θ(H·W)`.
+* **Export** - tight bbox crop + PNG `optimize`/`compress_level=9`; optional
+  median-cut palette for smaller archives.
+
+## Output layout
+
+```
+pieces/page_0001/piece_000.png   # fragment on black bg
+manifest.json                    # per-piece {file, x, y, w, h} = stitching GT
+                                 # + per-page adjacency [[i, j], ...] neighbor pairs
+README.txt                       # reassembly snippet
+```
+
+Each page also lists **`adjacency`** - undirected `[i, j]` piece-index pairs that
+share a torn border (4-connectivity). Use as positive pairs for pairwise /
+graph-based stitching models (Deepzzle / PairingNet style); any unlisted pair is
+a negative. Computed in one vectorized `Θ(H·W)` pass from the partition map - no
+measurable pipeline overhead.
+
+## Run locally
+
+Use **Python 3.10–3.12**. Python 3.13/3.14 have no PyMuPDF wheel yet and fall
+back to a source build.
+
+```bash
+python3.11 -m venv .venv && source .venv/bin/activate
+pip install -r requirements.txt
+python app.py            # http://127.0.0.1:7860
+pytest -q                # invariant + queue tests
+```
+
+### Pinned web stack (don't loosen)
+
+`gradio==4.44.1` needs a matching server stack. Newer auto-resolved versions
+break it, so these are pinned in `requirements.txt`:
+
+| pin | why |
+|-----|-----|
+| `fastapi==0.112.4` / `starlette==0.38.6` | starlette ≥0.29 reordered `TemplateResponse` args → gradio passes a dict as template name → `TypeError: unhashable type: 'dict'` on every page load |
+| `huggingface_hub==0.25.2` | hub ≥1.0 removed `HfFolder` that gradio 4.44 imports |
+| `pydantic==2.10.6` | pydantic ≥2.11 emits bool `additionalProperties` → gradio_client 1.3.0 `get_api_info()` crashes |
+
+## Layout
+
+```
+app.py                 Gradio UI, queue config, theme
+src/pdf_loader.py      PDF -> A4 RGB pages (PyMuPDF), tall-page slicing
+src/sampling.py        Bridson Poisson-disk seed sampling
+src/noise.py           Vectorized value noise (domain warp)
+src/tearing.py         Voronoi partition + warp + piece extraction  ← core
+src/optimizer.py       PNG optimization / quantization
+src/queue_manager.py   Priority job queue (min-heap)
+src/packager.py        ZIP + manifest builder
+src/pipeline.py        PDF -> torn pages orchestration
+tests/                 partition no-overlap + queue ordering
+```
+
+## License
+
+MIT

app.py

@@ -0,0 +1,226 @@
+"""Dataset-Maker - Gradio web app (HuggingFace Spaces ready).
+
+Upload a PDF -> each page is rendered to A4, torn into NON-OVERLAPPING fragments
+on a black background, and packaged as a ZIP with stitching ground truth.
+
+Performance:
+  * Gradio `.queue()` caps concurrent requests for the 2-vCPU free tier.
+  * A priority queue (src/queue_manager.py) orders page jobs cheap-first.
+  * NumPy/SciPy vectorized partition; PNG-optimized export.
+"""
+from __future__ import annotations
+
+import gradio as gr
+
+from src import config, workspace
+from src.optimizer import encode_preview
+from src.packager import build_zip
+from src.pipeline import process_pdf, save_temp_pdf
+from src.tearing import verify_partition
+
+
+def _resolve_theme(name: str):
+    """Resolve a registry theme, falling back gracefully across Gradio versions.
+
+    Some themes (Ocean, Citrus) only exist in Gradio 5+. On older Gradio we fall
+    back to Default rather than crashing at startup.
+    """
+    cls_name, kwargs = config.THEME_REGISTRY.get(
+        name, config.THEME_REGISTRY[config.DEFAULT_THEME]
+    )
+    cls = getattr(gr.themes, cls_name, None) or getattr(gr.themes, "Default")
+    return cls(**kwargs)
+
+
+def available_themes() -> list[str]:
+    """Registry themes actually present in the installed Gradio build."""
+    return [
+        name for name, (cls, _) in config.THEME_REGISTRY.items()
+        if getattr(gr.themes, cls, None) is not None
+    ]
+
+
+def _pieces_gallery(pages, max_pieces: int = 60):
+    """Flatten a few torn pieces for the preview gallery (downscaled)."""
+    out = []
+    for pi, page in enumerate(pages):
+        for k, piece in enumerate(page.pieces):
+            out.append((encode_preview(piece.rgb, 256), f"p{pi+1}·{k}"))
+            if len(out) >= max_pieces:
+                return out
+    return out
+
+
+def generate(
+    pdf_file,
+    dpi: int,
+    n_pieces: int,
+    noise_strength: float,
+    noise_scale: float,
+    lossy: bool,
+    seed: int,
+    progress=gr.Progress(),
+):
+    """Main event handler: PDF -> (status, gallery, zip path)."""
+    if pdf_file is None:
+        raise gr.Error("Upload a PDF first.")
+
+    # Drop temp files from the previous run so disk stays at steady state
+    # (~1 ZIP) instead of growing every generate. HF free-tier disk is small.
+    workspace.clear_all()
+
+    progress(0.02, desc="Reading PDF…")
+    with open(pdf_file, "rb") as fh:
+        pdf_bytes = fh.read()
+    if len(pdf_bytes) > config.MAX_UPLOAD_MB * 1024 * 1024:
+        raise gr.Error(f"PDF exceeds {config.MAX_UPLOAD_MB} MB limit.")
+
+    tmp_pdf = save_temp_pdf(pdf_bytes)
+    pages = process_pdf(
+        tmp_pdf,
+        dpi=int(dpi),
+        n_pieces=int(n_pieces),
+        noise_strength=float(noise_strength),
+        noise_scale=float(noise_scale),
+        master_seed=int(seed),
+        progress=lambda f, m: progress(0.05 + 0.8 * f, desc=m),
+    )
+    # Input PDF is fully rendered into `pages` now; free it immediately.
+    workspace.discard(tmp_pdf)
+
+    # Verify the no-overlap invariant on the first page (sanity gate).
+    report = verify_partition(pages[0])
+    if not report["is_partition"]:
+        raise gr.Error(
+            f"Partition check failed: overlap={report['max_overlap']}, "
+            f"uncovered={report['uncovered_pixels']}"
+        )
+
+    progress(0.9, desc="Packaging ZIP…")
+    zip_bytes, manifest = build_zip(
+        pages,
+        source_name="upload.pdf",
+        dpi=int(dpi),
+        noise_strength=float(noise_strength),
+        noise_scale=float(noise_scale),
+        lossy=lossy,
+    )
+    out_path = workspace.new_temp(suffix="_dataset.zip")
+    with open(out_path, "wb") as fh:
+        fh.write(zip_bytes)
+
+    status = (
+        f"✅ {len(pages)} pages · {manifest['total_pieces']} pieces · "
+        f"no-overlap verified (max_overlap={report['max_overlap']}, "
+        f"uncovered={report['uncovered_pixels']})"
+    )
+    progress(1.0, desc="Done")
+    # Order: gallery, zip, status. Status is consumed by a chained .then() with
+    # progress hidden, so no progress bar paints over the status text strip.
+    return _pieces_gallery(pages), out_path, status
+
+
+def clear_all():
+    """Delete tracked temp files (PDFs + ZIPs) and reset the UI outputs."""
+    removed = workspace.clear_all()
+    status = f"🧹 Cleared {removed} temp file(s). Upload a PDF and hit **Generate**."
+    # outputs order: pdf_in, status, gallery, zip_out
+    return None, status, None, None
+
+
+# Cap the preview gallery and scroll *inside* it. Gradio 4.44's Gallery `height`
+# caps the root but the inner thumbnail grid (.grid-wrap) overflows the page
+# instead of scrolling, so force overflow on the inner container directly.
+_GALLERY_CSS = """
+#piece-gallery { max-height: 70vh; }
+#piece-gallery .grid-wrap,
+#piece-gallery .thumbnails {
+    max-height: 70vh;
+    overflow-y: auto;
+}
+"""
+
+
+def build_ui(theme_name: str = config.DEFAULT_THEME) -> gr.Blocks:
+    with gr.Blocks(
+        theme=_resolve_theme(theme_name),
+        title="Dataset-Maker · Torn-page stitching dataset",
+        css=_GALLERY_CSS,
+    ) as demo:
+        gr.Markdown(
+            "# 🧩 Dataset-Maker\n"
+            "Tear PDF pages into **non-overlapping** torn fragments for "
+            "image-stitching datasets. Every pixel lands in exactly one piece - "
+            "guaranteed by a domain-warped Voronoi partition."
+        )
+        with gr.Row():
+            with gr.Column(scale=1):
+                pdf_in = gr.File(label="PDF", file_types=[".pdf"], type="filepath")
+                n_pieces = gr.Slider(
+                    config.MIN_PIECES, config.MAX_PIECES, config.DEFAULT_PIECES,
+                    step=1, label="Pieces per page",
+                )
+                with gr.Accordion("Tearing controls", open=False):
+                    noise_strength = gr.Slider(
+                        0, 80, config.DEFAULT_NOISE_STRENGTH, step=1,
+                        label="Tear jaggedness (px)",
+                    )
+                    noise_scale = gr.Slider(
+                        8, 200, config.DEFAULT_NOISE_SCALE, step=1,
+                        label="Tear smoothness (wavelength px)",
+                    )
+                    dpi = gr.Slider(
+                        config.MIN_DPI, config.MAX_DPI, config.DEFAULT_DPI, step=1,
+                        label="Render DPI",
+                    )
+                    seed = gr.Number(value=0, precision=0, label="Master seed")
+                    lossy = gr.Checkbox(
+                        value=False, label="Lossy palette PNG (smaller ZIP)"
+                    )
+                with gr.Row():
+                    run = gr.Button("Generate dataset", variant="primary")
+                    clear = gr.Button("Clear all", variant="secondary")
+            with gr.Column(scale=2):
+                status = gr.Markdown("Upload a PDF and hit **Generate**.")
+                gallery = gr.Gallery(
+                    label="Torn pieces (preview)", columns=6, height=420,
+                    object_fit="contain", elem_id="piece-gallery",
+                )
+                zip_out = gr.File(label="Download dataset (.zip)")
+
+        # Status flows through a State, then into the Markdown via a hidden-
+        # progress .then() — keeps the progress bars on gallery + zip only,
+        # not over the thin status text (4.44 has no per-output show_progress).
+        status_state = gr.State("")
+
+        run.click(
+            generate,
+            inputs=[pdf_in, dpi, n_pieces, noise_strength, noise_scale, lossy, seed],
+            outputs=[gallery, zip_out, status_state],
+            concurrency_limit=config.WORKER_CONCURRENCY,  # heavy job throttle
+        ).then(
+            lambda s: s,
+            inputs=status_state,
+            outputs=status,
+            show_progress="hidden",
+        )
+        clear.click(
+            clear_all,
+            inputs=None,
+            outputs=[pdf_in, status, gallery, zip_out],
+        )
+        gr.Markdown(
+            "Pieces sit on black backgrounds; `manifest.json` carries each "
+            "piece's `(x, y)` offset = the stitching label."
+        )
+    return demo
+
+
+demo = build_ui()
+demo.queue(
+    max_size=config.QUEUE_MAX_SIZE,
+    default_concurrency_limit=config.WORKER_CONCURRENCY,
+)
+
+if __name__ == "__main__":
+    demo.launch(share=True)

pytest.ini

@@ -0,0 +1,5 @@
+[pytest]
+# Put the repo root on sys.path so `import src...` works under bare `pytest`
+# (CI), not only `python -m pytest` (which adds CWD implicitly).
+pythonpath = .
+testpaths = tests

requirements.txt

@@ -0,0 +1,9 @@
+gradio==4.44.1
+huggingface_hub==0.25.2
+pydantic==2.10.6
+fastapi==0.112.4
+starlette==0.38.6
+PyMuPDF==1.24.10
+numpy==1.26.4
+scipy==1.13.1
+Pillow==10.4.0

src/__init__.py

@@ -0,0 +1,3 @@
+"""Dataset-Maker: tear PDF pages into non-overlapping fragments for stitching."""
+
+__version__ = "1.0.0"

src/config.py

@@ -0,0 +1,46 @@
+"""Central configuration: A4 geometry, defaults, theme registry.
+
+All tunables live here so the UI, workers and tests share one source of truth.
+"""
+from __future__ import annotations
+
+# --- A4 geometry -----------------------------------------------------------
+# A4 = 210 x 297 mm. Pixel size scales with DPI: px = mm / 25.4 * dpi.
+A4_MM = (210.0, 297.0)
+A4_ASPECT = A4_MM[1] / A4_MM[0]  # height / width ~= 1.4142
+
+DEFAULT_DPI = 150           # 150 DPI A4 -> 1240 x 1754 px. Good detail/size trade-off.
+MIN_DPI, MAX_DPI = 72, 300
+
+# --- Tearing defaults ------------------------------------------------------
+DEFAULT_PIECES = 16         # pieces per page (n^2 grid feel -> here a flat count)
+MIN_PIECES, MAX_PIECES = 2, 256
+
+DEFAULT_NOISE_STRENGTH = 28.0   # px of boundary displacement (the "tear" jaggedness)
+DEFAULT_NOISE_SCALE = 96.0      # px wavelength of the noise (bigger = smoother tears)
+
+# --- Performance / limits --------------------------------------------------
+MAX_PAGES_PER_PDF = 60      # guardrail for the HF free tier (CPU/RAM bound)
+MAX_UPLOAD_MB = 50
+QUEUE_MAX_SIZE = 32         # Gradio request queue cap
+WORKER_CONCURRENCY = 1      # HF free tier = 2 vCPU; keep 1 heavy job at a time
+
+# --- Theme registry (Gradio built-ins, see theming guide) ------------------
+# name -> (gr.themes class name, kwargs). Resolved lazily in app.py to avoid
+# importing gradio inside worker/test code paths.
+THEME_REGISTRY = {
+    "Ocean": ("Ocean", {}),
+    "Soft": ("Soft", {}),
+    "Glass": ("Glass", {}),
+    "Monochrome": ("Monochrome", {}),
+    "Citrus": ("Citrus", {}),
+    "Default": ("Default", {}),
+}
+DEFAULT_THEME = "Soft"   # present in Gradio 4 and 5; Ocean/Citrus are 5-only
+
+
+def a4_pixels(dpi: int) -> tuple[int, int]:
+    """Return (width, height) in px for an A4 page at the given DPI."""
+    w = round(A4_MM[0] / 25.4 * dpi)
+    h = round(A4_MM[1] / 25.4 * dpi)
+    return w, h

src/noise.py

@@ -0,0 +1,61 @@
+"""Vectorized value-noise field for domain warping.
+
+We do NOT pull in a Perlin dependency: a tileable value-noise (random lattice +
+bilinear upsample, summed across octaves) is enough to warp Voronoi boundaries
+into organic "torn" edges, and it is pure NumPy.
+
+Complexity: building an (H, W) field over `octaves` is
+    Theta(H * W * octaves)  time,
+    Theta(H * W)            space.
+"""
+from __future__ import annotations
+
+import numpy as np
+
+
+def _bilinear_upsample(grid: np.ndarray, out_h: int, out_w: int) -> np.ndarray:
+    """Upsample a small lattice to (out_h, out_w) with bilinear interpolation."""
+    gh, gw = grid.shape
+    # Sample positions in lattice space.
+    ys = np.linspace(0, gh - 1, out_h)
+    xs = np.linspace(0, gw - 1, out_w)
+    y0 = np.floor(ys).astype(np.int64)
+    x0 = np.floor(xs).astype(np.int64)
+    y1 = np.minimum(y0 + 1, gh - 1)
+    x1 = np.minimum(x0 + 1, gw - 1)
+    wy = (ys - y0)[:, None]
+    wx = (xs - x0)[None, :]
+
+    top = grid[y0][:, x0] * (1 - wx) + grid[y0][:, x1] * wx
+    bot = grid[y1][:, x0] * (1 - wx) + grid[y1][:, x1] * wx
+    return top * (1 - wy) + bot * wy
+
+
+def value_noise(
+    h: int,
+    w: int,
+    scale: float,
+    rng: np.random.Generator,
+    octaves: int = 3,
+    persistence: float = 0.5,
+) -> np.ndarray:
+    """Return an (h, w) float field in roughly [-1, 1].
+
+    `scale` is the wavelength in pixels of the base octave: larger -> smoother.
+    Octaves add finer detail (fractal/fBm), persistence weights their amplitude.
+    """
+    field = np.zeros((h, w), dtype=np.float32)
+    amplitude = 1.0
+    total_amp = 0.0
+    freq_scale = max(scale, 2.0)
+
+    for _ in range(max(1, octaves)):
+        gh = max(2, int(np.ceil(h / freq_scale)) + 1)
+        gw = max(2, int(np.ceil(w / freq_scale)) + 1)
+        lattice = rng.uniform(-1.0, 1.0, size=(gh, gw)).astype(np.float32)
+        field += amplitude * _bilinear_upsample(lattice, h, w)
+        total_amp += amplitude
+        amplitude *= persistence
+        freq_scale = max(freq_scale * 0.5, 2.0)
+
+    return field / total_amp

src/optimizer.py

@@ -0,0 +1,38 @@
+"""Image optimization for the exported pieces.
+
+Torn pieces are mostly black background -> they compress extremely well as PNG.
+We:
+  * tight-crop to the fragment bbox (already done in tearing),
+  * optionally quantize the foreground to <=256 colors (palette PNG) to shrink
+    size further when `lossy` is on,
+  * always write with PNG optimize + max zlib compression.
+
+Palette quantization is Theta(P) over piece pixels (Pillow median-cut).
+"""
+from __future__ import annotations
+
+import io
+
+import numpy as np
+from PIL import Image
+
+
+def encode_piece(rgb: np.ndarray, lossy: bool = False, colors: int = 64) -> bytes:
+    """Encode an (h, w, 3) uint8 piece to optimized PNG bytes."""
+    img = Image.fromarray(rgb, mode="RGB")
+    if lossy:
+        # Median-cut palette; black background collapses to one palette entry.
+        img = img.quantize(colors=max(2, min(256, colors)), method=Image.MEDIANCUT)
+    buf = io.BytesIO()
+    img.save(buf, format="PNG", optimize=True, compress_level=9)
+    return buf.getvalue()
+
+
+def encode_preview(rgb: np.ndarray, max_side: int = 1024) -> np.ndarray:
+    """Downscale a page/preview for fast UI display (keeps aspect)."""
+    H, W = rgb.shape[:2]
+    scale = min(1.0, max_side / max(H, W))
+    if scale >= 1.0:
+        return rgb
+    nw, nh = int(W * scale), int(H * scale)
+    return np.asarray(Image.fromarray(rgb).resize((nw, nh), Image.BILINEAR))

src/packager.py

@@ -0,0 +1,93 @@
+"""Package torn pieces + stitching ground-truth into a downloadable ZIP.
+
+Layout inside the archive:
+    pieces/page_0001/piece_000.png ...
+    manifest.json          # global summary + per-piece placement (x, y, w, h)
+    README.txt             # how to reassemble
+
+The manifest IS the dataset label: each piece's (x, y) offset on its page is the
+exact stitching target. Reassembling = paste every piece at its offset.
+"""
+from __future__ import annotations
+
+import io
+import json
+import zipfile
+from datetime import datetime, timezone
+
+from .optimizer import encode_piece
+from .tearing import TornPage
+
+
+def build_zip(
+    pages: list[TornPage],
+    *,
+    source_name: str,
+    dpi: int,
+    noise_strength: float,
+    noise_scale: float,
+    lossy: bool,
+) -> tuple[bytes, dict]:
+    """Return (zip_bytes, manifest_dict) for a list of torn pages."""
+    manifest = {
+        "generator": "Dataset-Maker",
+        "created_utc": datetime.now(timezone.utc).isoformat(),
+        "source": source_name,
+        "dpi": dpi,
+        "noise_strength": noise_strength,
+        "noise_scale": noise_scale,
+        "lossy": lossy,
+        "pages": [],
+        "total_pieces": 0,
+    }
+
+    buf = io.BytesIO()
+    with zipfile.ZipFile(buf, "w", compression=zipfile.ZIP_DEFLATED) as zf:
+        for pi, page in enumerate(pages):
+            pdir = f"pieces/page_{pi + 1:04d}"
+            page_entry = {
+                "index": pi,
+                "width": page.width,
+                "height": page.height,
+                # Undirected neighbor pairs (piece-index i, j) = which fragments
+                # share a torn border. Positive pairs for pairwise/graph stitching
+                # models; non-listed pairs are negatives.
+                "adjacency": [[int(i), int(j)] for i, j in page.adjacency],
+                "pieces": [],
+            }
+            for k, piece in enumerate(page.pieces):
+                fname = f"{pdir}/piece_{k:03d}.png"
+                zf.writestr(fname, encode_piece(piece.rgb, lossy=lossy))
+                h, w = piece.mask.shape
+                page_entry["pieces"].append(
+                    {"file": fname, "x": piece.x, "y": piece.y, "w": w, "h": h}
+                )
+            manifest["total_pieces"] += len(page.pieces)
+            manifest["pages"].append(page_entry)
+
+        zf.writestr("manifest.json", json.dumps(manifest, indent=2))
+        zf.writestr("README.txt", _README)
+
+    return buf.getvalue(), manifest
+
+
+_README = """Dataset-Maker export
+=====================
+Each page was torn into NON-OVERLAPPING fragments (a strict partition: every
+pixel belongs to exactly one piece). Fragments sit on a black background.
+
+Each page also carries `adjacency`: a list of [i, j] piece-index pairs that
+share a torn border (4-connectivity, undirected, i < j). Use as positive pairs
+for pairwise/graph-based stitching models; any pair not listed is a negative.
+
+To reassemble a page (stitching ground truth):
+    import json
+    from PIL import Image
+    m = json.load(open("manifest.json"))
+    for page in m["pages"]:
+        canvas = Image.new("RGB", (page["width"], page["height"]))
+        for p in page["pieces"]:
+            piece = Image.open(p["file"])
+            canvas.paste(piece, (p["x"], p["y"]), mask=...)  # non-black pixels
+        canvas.save(f"reassembled_{page['index']}.png")
+"""

src/pdf_loader.py

@@ -0,0 +1,85 @@
+"""PDF -> list of A4-sized RGB page images using PyMuPDF (fitz).
+
+"Native split" handling: if a page is already ~A4 portrait we render it as-is.
+If a page is much larger / a different ratio (e.g. an A3 spread, a long scan),
+we slice it into A4-height bands so downstream tearing always works on A4 tiles.
+"""
+from __future__ import annotations
+
+import numpy as np
+
+from .config import A4_ASPECT, MAX_PAGES_PER_PDF, a4_pixels
+
+
+def _page_pixmap(page, dpi: int) -> np.ndarray:
+    import fitz  # PyMuPDF
+
+    zoom = dpi / 72.0
+    pm = page.get_pixmap(matrix=fitz.Matrix(zoom, zoom), alpha=False)
+    arr = np.frombuffer(pm.samples, dtype=np.uint8).reshape(pm.h, pm.w, pm.n)
+    if pm.n == 4:           # RGBA -> RGB
+        arr = arr[:, :, :3]
+    elif pm.n == 1:         # gray -> RGB
+        arr = np.repeat(arr, 3, axis=2)
+    return np.ascontiguousarray(arr)
+
+
+def _fit_to_a4(img: np.ndarray, dpi: int) -> list[np.ndarray]:
+    """Return one or more A4-portrait tiles covering `img`.
+
+    Tall pages are sliced into A4-height bands (each band padded to A4 width).
+    """
+    target_w, target_h = a4_pixels(dpi)
+    H, W = img.shape[:2]
+    aspect = H / W
+
+    # Close enough to A4 portrait: letterbox-resize the whole page onto A4.
+    if abs(aspect - A4_ASPECT) < 0.12:
+        return [_letterbox(img, target_w, target_h)]
+
+    # Otherwise scale to A4 width, then slice the (now-tall) image into bands.
+    scale = target_w / W
+    new_h = max(1, int(round(H * scale)))
+    resized = _resize(img, target_w, new_h)
+    tiles = []
+    for top in range(0, new_h, target_h):
+        band = resized[top:top + target_h]
+        if band.shape[0] < target_h:
+            band = _letterbox(band, target_w, target_h)
+        tiles.append(band)
+    return tiles
+
+
+def _resize(img: np.ndarray, w: int, h: int) -> np.ndarray:
+    from PIL import Image
+
+    return np.asarray(Image.fromarray(img).resize((w, h), Image.LANCZOS))
+
+
+def _letterbox(img: np.ndarray, w: int, h: int) -> np.ndarray:
+    """Resize preserving aspect, pad with white onto a w*h canvas."""
+    H, W = img.shape[:2]
+    scale = min(w / W, h / H)
+    nw, nh = max(1, int(W * scale)), max(1, int(H * scale))
+    resized = _resize(img, nw, nh)
+    canvas = np.full((h, w, 3), 255, dtype=np.uint8)
+    oy, ox = (h - nh) // 2, (w - nw) // 2
+    canvas[oy:oy + nh, ox:ox + nw] = resized
+    return canvas
+
+
+def load_pdf_pages(path: str, dpi: int) -> list[np.ndarray]:
+    """Render `path` into a list of A4 RGB uint8 page images."""
+    import fitz
+
+    doc = fitz.open(path)
+    try:
+        out: list[np.ndarray] = []
+        for page in doc:
+            raw = _page_pixmap(page, dpi)
+            out.extend(_fit_to_a4(raw, dpi))
+            if len(out) >= MAX_PAGES_PER_PDF:
+                return out[:MAX_PAGES_PER_PDF]
+        return out
+    finally:
+        doc.close()

src/pipeline.py

@@ -0,0 +1,64 @@
+"""Orchestration: PDF bytes -> torn pages, scheduled through the priority queue.
+
+Kept UI-free so it is unit-testable and reusable from a CLI or batch worker.
+"""
+from __future__ import annotations
+
+from typing import Callable
+
+from . import workspace
+from .pdf_loader import load_pdf_pages
+from .queue_manager import PriorityJobQueue, page_priority
+from .tearing import TornPage, tear_page
+
+
+def process_pdf(
+    pdf_path: str,
+    *,
+    dpi: int,
+    n_pieces: int,
+    noise_strength: float,
+    noise_scale: float,
+    master_seed: int = 0,
+    progress: Callable[[float, str], None] | None = None,
+) -> list[TornPage]:
+    """Render + tear every page, ordered by the priority queue (cheap first)."""
+    pages = load_pdf_pages(pdf_path, dpi)
+    if not pages:
+        raise ValueError("No renderable pages found in PDF.")
+
+    queue = PriorityJobQueue()
+    for idx, page_img in enumerate(pages):
+        queue.push(page_priority(n_pieces, idx), payload=(idx, page_img))
+
+    total = len(pages)
+    results: dict[int, TornPage] = {}
+    done = 0
+    while True:
+        job = queue.pop()
+        if job is None:
+            break
+        idx, page_img = job.payload
+        # Per-page seed -> randomness changes page by page, yet reproducible.
+        seed = (master_seed * 1_000_003 + idx) & 0x7FFFFFFF
+        results[idx] = tear_page(
+            page_img,
+            n_pieces=n_pieces,
+            seed=seed,
+            noise_strength=noise_strength,
+            noise_scale=noise_scale,
+        )
+        done += 1
+        if progress:
+            progress(done / total, f"Torn page {done}/{total}")
+
+    # Return in document order for a coherent manifest.
+    return [results[i] for i in sorted(results)]
+
+
+def save_temp_pdf(file_bytes: bytes) -> str:
+    """Persist uploaded bytes to a tracked temp file PyMuPDF can open."""
+    path = workspace.new_temp(suffix=".pdf")
+    with open(path, "wb") as fh:
+        fh.write(file_bytes)
+    return path

src/queue_manager.py

@@ -0,0 +1,84 @@
+"""Priority job queue for batch dataset generation.
+
+Why a custom queue when Gradio already queues HTTP requests? Gradio's queue is
+FIFO over *requests*. Here a single request can enqueue many *page jobs*, and we
+want cheap jobs (few pieces / small pages) to clear first so the user sees early
+progress and tail latency drops. That ordering is a priority queue's job.
+
+Data structure: binary min-heap (heapq) keyed by a (priority, seq) tuple.
+`seq` is a monotonic counter that breaks ties FIFO and keeps ordering stable
+without comparing the payloads.
+
+Complexity:
+    push  : O(log n)          worst & average
+    pop   : O(log n)          worst & average
+    peek  : Theta(1)
+    space : Theta(n)
+Big-Omega: push/pop are Omega(1) in the best case (no sift needed); the heap can
+never beat Omega(log n) amortized when both ops interleave under arbitrary keys.
+"""
+from __future__ import annotations
+
+import heapq
+import itertools
+import threading
+from dataclasses import dataclass, field
+from typing import Any, Callable
+
+
+@dataclass(order=False)
+class Job:
+    priority: float
+    payload: Any
+    seq: int = field(default=0)
+
+
+class PriorityJobQueue:
+    """Thread-safe min-heap priority queue. Lower priority value = runs first."""
+
+    def __init__(self) -> None:
+        self._heap: list[tuple[float, int, Job]] = []
+        self._counter = itertools.count()
+        self._lock = threading.Lock()
+
+    def __len__(self) -> int:
+        with self._lock:
+            return len(self._heap)
+
+    def push(self, priority: float, payload: Any) -> Job:
+        """O(log n). Smaller `priority` is dequeued earlier; ties are FIFO."""
+        with self._lock:
+            seq = next(self._counter)
+            job = Job(priority=priority, payload=payload, seq=seq)
+            heapq.heappush(self._heap, (priority, seq, job))
+            return job
+
+    def pop(self) -> Job | None:
+        """O(log n). Returns the highest-priority job, or None if empty."""
+        with self._lock:
+            if not self._heap:
+                return None
+            return heapq.heappop(self._heap)[2]
+
+    def peek(self) -> Job | None:
+        """Theta(1). Look at the next job without removing it."""
+        with self._lock:
+            return self._heap[0][2] if self._heap else None
+
+    def drain(self, handler: Callable[[Job], Any]) -> list[Any]:
+        """Pop every job in priority order, applying `handler`. Returns results."""
+        results = []
+        while True:
+            job = self.pop()
+            if job is None:
+                break
+            results.append(handler(job))
+        return results
+
+
+def page_priority(piece_count: int, page_index: int) -> float:
+    """Cheaper pages first; page_index breaks ties to keep document order-ish.
+
+    Cost grows with piece_count (more masks to extract), so use it as the key.
+    """
+    return float(piece_count) * 1000.0 + float(page_index)

src/sampling.py

@@ -0,0 +1,99 @@
+"""Seed-point sampling for the Voronoi partition.
+
+Bridson Poisson-disk sampling gives blue-noise seeds: random but evenly spread,
+so piece areas stay within sane bounds (no slivers, no giant blobs). We target
+a piece count, derive the min-distance r from area, then top up / trim to hit
+the exact count.
+
+Bridson is Theta(k) in the number of accepted samples (grid-accelerated
+neighbour test), Theta(grid cells) = Theta(W*H / r^2) space.
+"""
+from __future__ import annotations
+
+import math
+
+import numpy as np
+
+
+def poisson_disk(
+    width: int,
+    height: int,
+    radius: float,
+    rng: np.random.Generator,
+    k: int = 30,
+) -> np.ndarray:
+    """Bridson's algorithm. Returns (n, 2) float array of (x, y) samples."""
+    cell = radius / math.sqrt(2)
+    gw = int(math.ceil(width / cell))
+    gh = int(math.ceil(height / cell))
+    grid = -np.ones((gh, gw), dtype=np.int64)
+    samples: list[tuple[float, float]] = []
+    active: list[int] = []
+
+    def grid_xy(p):
+        return int(p[0] / cell), int(p[1] / cell)
+
+    first = (rng.uniform(0, width), rng.uniform(0, height))
+    samples.append(first)
+    gx, gy = grid_xy(first)
+    grid[gy, gx] = 0
+    active.append(0)
+
+    while active:
+        idx = active[rng.integers(len(active))]
+        base = samples[idx]
+        placed = False
+        for _ in range(k):
+            ang = rng.uniform(0, 2 * math.pi)
+            rad = rng.uniform(radius, 2 * radius)
+            cand = (base[0] + rad * math.cos(ang), base[1] + rad * math.sin(ang))
+            if not (0 <= cand[0] < width and 0 <= cand[1] < height):
+                continue
+            cgx, cgy = grid_xy(cand)
+            ok = True
+            for ny in range(max(0, cgy - 2), min(gh, cgy + 3)):
+                for nx in range(max(0, cgx - 2), min(gw, cgx + 3)):
+                    j = grid[ny, nx]
+                    if j >= 0:
+                        d = math.dist(cand, samples[j])
+                        if d < radius:
+                            ok = False
+                            break
+                if not ok:
+                    break
+            if ok:
+                samples.append(cand)
+                grid[cgy, cgx] = len(samples) - 1
+                active.append(len(samples) - 1)
+                placed = True
+                break
+        if not placed:
+            active.remove(idx)
+
+    return np.asarray(samples, dtype=np.float32)
+
+
+def sample_seeds(
+    width: int,
+    height: int,
+    n_pieces: int,
+    rng: np.random.Generator,
+) -> np.ndarray:
+    """Return exactly `n_pieces` blue-noise seeds (x, y) inside the page."""
+    n_pieces = max(2, int(n_pieces))
+    # r from target area per piece: area ~ W*H/n, packing factor ~0.7.
+    area_per = (width * height) / n_pieces
+    radius = math.sqrt(area_per) * 0.66
+
+    pts = poisson_disk(width, height, radius, rng)
+
+    # Top up with uniform random points if Poisson under-shot.
+    if len(pts) < n_pieces:
+        extra = rng.uniform([0, 0], [width, height], size=(n_pieces - len(pts), 2))
+        pts = np.vstack([pts, extra.astype(np.float32)])
+    # Trim deterministically if it over-shot.
+    if len(pts) > n_pieces:
+        keep = rng.choice(len(pts), size=n_pieces, replace=False)
+        pts = pts[keep]
+
+    return pts.astype(np.float32)

src/tearing.py

@@ -0,0 +1,156 @@
+"""Core: tear a page image into non-overlapping torn pieces.
+
+Guarantee (the whole point of the dataset): the output is a PARTITION of the
+page. Every pixel is assigned to exactly one piece via nearest-seed argmin, so
+pieces never overlap and together cover the page exactly -> perfect ground
+truth for image stitching.
+
+The "torn" look comes from DOMAIN WARPING: before computing the nearest seed we
+displace each pixel's coordinate by a value-noise vector. Warping the query
+points (not the partition rule) keeps the result a strict partition while making
+boundaries jagged/organic instead of straight Voronoi edges.
+
+Complexity (H*W pixels, S seeds, kd-tree):
+    nearest-seed query : O(H*W * log S)   time
+    label / mask pass  : Theta(H*W)       time, Theta(H*W) space
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+from scipy.spatial import cKDTree
+
+from .noise import value_noise
+from .sampling import sample_seeds
+
+
+@dataclass
+class Piece:
+    """One torn fragment + its placement for reassembly ground truth."""
+    label: int
+    x: int          # left offset on the original page canvas
+    y: int          # top offset
+    rgb: np.ndarray      # (h, w, 3) uint8, black where outside the fragment
+    mask: np.ndarray     # (h, w) bool, True inside the fragment
+
+
+@dataclass
+class TornPage:
+    width: int
+    height: int
+    pieces: list[Piece]
+    labels: np.ndarray   # (H, W) int32 partition map (for verification / GT)
+    adjacency: list[tuple[int, int]]   # undirected (i, j) piece-index neighbor pairs
+
+
+def _adjacency_pairs(labels: np.ndarray) -> np.ndarray:
+    """Return unique unordered raw-label neighbor pairs from a partition map.
+
+    4-connectivity: two pieces are neighbors iff they touch horizontally or
+    vertically. Vectorized: compare each pixel to its right/down neighbor, keep
+    label pairs that differ. Cost Theta(H*W) - a few ms even at 150 DPI, dwarfed
+    by the kd-tree query, so no measurable pipeline slowdown.
+    """
+    h_a, h_b = labels[:, :-1], labels[:, 1:]
+    v_a, v_b = labels[:-1, :], labels[1:, :]
+    hd, vd = h_a != h_b, v_a != v_b
+    pairs = np.concatenate(
+        [
+            np.stack([h_a[hd], h_b[hd]], axis=1),
+            np.stack([v_a[vd], v_b[vd]], axis=1),
+        ],
+        axis=0,
+    )
+    if pairs.size == 0:                       # single-piece page
+        return pairs.reshape(0, 2)
+    pairs.sort(axis=1)                        # (min, max) -> undirected
+    return np.unique(pairs, axis=0)
+
+
+def compute_adjacency(
+    labels: np.ndarray, label_to_idx: dict[int, int]
+) -> list[tuple[int, int]]:
+    """Map raw-label neighbor pairs to manifest piece indices, sorted."""
+    out = []
+    for a, b in _adjacency_pairs(labels):
+        ia, ib = label_to_idx.get(int(a)), label_to_idx.get(int(b))
+        if ia is not None and ib is not None and ia != ib:
+            out.append((ia, ib) if ia < ib else (ib, ia))
+    return sorted(set(out))
+
+
+def tear_page(
+    page_rgb: np.ndarray,
+    n_pieces: int,
+    seed: int,
+    noise_strength: float,
+    noise_scale: float,
+) -> TornPage:
+    """Partition `page_rgb` (H, W, 3 uint8) into `n_pieces` torn fragments.
+
+    `seed` makes a page reproducible; pass a different seed per page so the
+    randomness changes page by page.
+    """
+    if page_rgb.ndim != 3 or page_rgb.shape[2] != 3:
+        raise ValueError("page_rgb must be (H, W, 3) uint8")
+    H, W = page_rgb.shape[:2]
+    rng = np.random.default_rng(seed)
+
+    seeds = sample_seeds(W, H, n_pieces, rng)             # (S, 2) -> (x, y)
+
+    # Domain-warp the pixel grid with two independent noise fields.
+    ys, xs = np.mgrid[0:H, 0:W]
+    wx = value_noise(H, W, noise_scale, rng) * noise_strength
+    wy = value_noise(H, W, noise_scale, rng) * noise_strength
+    qx = (xs + wx).ravel()
+    qy = (ys + wy).ravel()
+    query = np.stack([qx, qy], axis=1).astype(np.float32)
+
+    tree = cKDTree(seeds)
+    _, flat_labels = tree.query(query, k=1, workers=-1)
+    labels = flat_labels.reshape(H, W).astype(np.int32)
+
+    pieces: list[Piece] = []
+    for lbl in np.unique(labels):
+        mask = labels == lbl
+        ys_idx, xs_idx = np.nonzero(mask)
+        if ys_idx.size == 0:
+            continue
+        y0, y1 = int(ys_idx.min()), int(ys_idx.max()) + 1
+        x0, x1 = int(xs_idx.min()), int(xs_idx.max()) + 1
+        sub_mask = mask[y0:y1, x0:x1]
+        rgb = np.zeros((y1 - y0, x1 - x0, 3), dtype=np.uint8)   # black background
+        rgb[sub_mask] = page_rgb[y0:y1, x0:x1][sub_mask]
+        pieces.append(
+            Piece(label=int(lbl), x=x0, y=y0, rgb=rgb, mask=sub_mask)
+        )
+
+    # Piece-index <-> raw-label map from the pieces we actually emitted, so
+    # adjacency indices line up exactly with the manifest's piece ordering.
+    label_to_idx = {p.label: i for i, p in enumerate(pieces)}
+    adjacency = compute_adjacency(labels, label_to_idx)
+
+    return TornPage(
+        width=W, height=H, pieces=pieces, labels=labels, adjacency=adjacency
+    )
+
+
+def verify_partition(torn: TornPage) -> dict:
+    """Assert the no-overlap / full-coverage invariants. Returns a report.
+
+    Reassembles a coverage counter from piece masks at their offsets and checks
+    every pixel is covered exactly once.
+    """
+    cover = np.zeros((torn.height, torn.width), dtype=np.int32)
+    for p in torn.pieces:
+        h, w = p.mask.shape
+        cover[p.y:p.y + h, p.x:p.x + w] += p.mask
+    max_cover = int(cover.max())
+    min_cover = int(cover.min())
+    return {
+        "pieces": len(torn.pieces),
+        "max_overlap": max_cover,       # must be 1 -> no overlap
+        "uncovered_pixels": int((cover == 0).sum()),  # must be 0 -> full cover
+        "is_partition": max_cover == 1 and min_cover == 1,
+    }

src/workspace.py

@@ -0,0 +1,72 @@
+"""Temp-file registry so 'Clear all' genuinely frees disk, not just the UI.
+
+Every PDF/ZIP scratch file goes through `new_temp`, which records the path. A
+later `clear_all` unlinks every recorded file. Thread-safe (a min worker pool
+on HF still shares this process). UI-free so `src/` stays testable.
+
+Note: this clears the *file cache* we create. Gradio's own request queue is
+per-request and transient (a handler can't flush other users' pending events),
+and the priority queue in `queue_manager` is built and drained within a single
+`process_pdf` call - neither leaves persistent state to clear.
+"""
+from __future__ import annotations
+
+import os
+import tempfile
+import threading
+
+_lock = threading.Lock()
+_tracked: set[str] = set()
+
+
+def new_temp(suffix: str = "") -> str:
+    """Create a tracked temp file and return its path (handle closed)."""
+    fd, path = tempfile.mkstemp(suffix=suffix)
+    os.close(fd)
+    with _lock:
+        _tracked.add(path)
+    return path
+
+
+def register(path: str) -> None:
+    """Track an externally created path so clear_all() will remove it."""
+    with _lock:
+        _tracked.add(path)
+
+
+def discard(path: str) -> bool:
+    """Unlink one tracked file early (e.g. an input PDF after it's loaded).
+
+    Returns True if the file was removed. Untracks regardless so a vanished
+    file doesn't linger in the registry.
+    """
+    with _lock:
+        _tracked.discard(path)
+    try:
+        os.remove(path)
+        return True
+    except OSError:
+        return False
+
+
+def clear_all() -> int:
+    """Unlink every tracked temp file. Returns count actually removed."""
+    removed = 0
+    with _lock:
+        # list() snapshot is required: we mutate _tracked (discard) in-loop.
+        # Iterating the set directly -> "Set changed size during iteration".
+        for path in list(_tracked):  # NOSONAR python:S7504 false positive
+            try:
+                os.remove(path)
+                removed += 1
+            except FileNotFoundError:
+                pass
+            except OSError:
+                continue  # leave it tracked; retry on next clear
+            _tracked.discard(path)
+    return removed
+
+
+def tracked_count() -> int:
+    with _lock:
+        return len(_tracked)

tests/test_partition.py

@@ -0,0 +1,73 @@
+"""No-overlap / full-coverage invariant tests - the dataset's core guarantee."""
+import numpy as np
+
+from src.tearing import compute_adjacency, tear_page, verify_partition
+
+
+def _page(h=400, w=300):
+    rng = np.random.default_rng(0)
+    return rng.integers(0, 255, size=(h, w, 3), dtype=np.uint8)
+
+
+def test_partition_no_overlap_full_cover():
+    torn = tear_page(_page(), n_pieces=12, seed=1,
+                     noise_strength=20, noise_scale=60)
+    rep = verify_partition(torn)
+    assert rep["is_partition"], rep
+    assert rep["max_overlap"] == 1
+    assert rep["uncovered_pixels"] == 0
+
+
+def test_piece_count_close_to_request():
+    torn = tear_page(_page(), n_pieces=16, seed=2,
+                     noise_strength=25, noise_scale=80)
+    # Some seeds can lose their cell after warping; allow a small shortfall.
+    assert 10 <= len(torn.pieces) <= 16
+
+
+def test_per_page_randomness_changes():
+    p = _page()
+    a = tear_page(p, 12, seed=1, noise_strength=20, noise_scale=60).labels
+    b = tear_page(p, 12, seed=2, noise_strength=20, noise_scale=60).labels
+    assert not np.array_equal(a, b)
+
+
+def test_reproducible_same_seed():
+    p = _page()
+    a = tear_page(p, 12, seed=5, noise_strength=20, noise_scale=60).labels
+    b = tear_page(p, 12, seed=5, noise_strength=20, noise_scale=60).labels
+    assert np.array_equal(a, b)
+
+
+def test_adjacency_known_grid():
+    # 3x3 label map: pieces 0,1,2 all mutually touch.
+    labels = np.array([[0, 0, 1],
+                       [0, 2, 1],
+                       [2, 2, 1]], dtype=np.int32)
+    adj = compute_adjacency(labels, {0: 0, 1: 1, 2: 2})
+    assert adj == [(0, 1), (0, 2), (1, 2)]
+
+
+def test_adjacency_invariants_on_real_page():
+    torn = tear_page(_page(), n_pieces=16, seed=3,
+                     noise_strength=25, noise_scale=80)
+    n = len(torn.pieces)
+    adj = torn.adjacency
+    # undirected, deduped, sorted, in-range, no self-loops
+    assert adj == sorted(set(adj))
+    for i, j in adj:
+        assert 0 <= i < j < n
+    # partition over a connected page -> adjacency graph is connected
+    seen = set()
+    stack = [0]
+    nbrs = {k: set() for k in range(n)}
+    for i, j in adj:
+        nbrs[i].add(j)
+        nbrs[j].add(i)
+    while stack:
+        u = stack.pop()
+        if u in seen:
+            continue
+        seen.add(u)
+        stack.extend(nbrs[u] - seen)
+    assert len(seen) == n          # every piece reachable

tests/test_queue.py

@@ -0,0 +1,31 @@
+"""Priority queue ordering + complexity-contract tests."""
+from src.queue_manager import PriorityJobQueue, page_priority
+
+
+def test_priority_order_cheap_first():
+    q = PriorityJobQueue()
+    q.push(page_priority(64, 0), "expensive")
+    q.push(page_priority(4, 0), "cheap")
+    q.push(page_priority(16, 0), "mid")
+    order = [q.pop().payload for _ in range(3)]
+    assert order == ["cheap", "mid", "expensive"]
+
+
+def test_fifo_tie_break():
+    q = PriorityJobQueue()
+    for i in range(5):
+        q.push(1.0, i)          # identical priority
+    assert [q.pop().payload for _ in range(5)] == [0, 1, 2, 3, 4]
+
+
+def test_pop_empty_returns_none():
+    assert PriorityJobQueue().pop() is None
+
+
+def test_len_and_peek():
+    q = PriorityJobQueue()
+    q.push(2.0, "b")
+    q.push(1.0, "a")
+    assert len(q) == 2
+    assert q.peek().payload == "a"
+    assert len(q) == 2          # peek does not consume

tests/test_workspace.py

@@ -0,0 +1,34 @@
+"""Temp-file registry: 'Clear all' must genuinely unlink tracked files."""
+import os
+
+from src import workspace
+
+
+def test_new_temp_tracked_then_cleared():
+    before = workspace.tracked_count()
+    p1 = workspace.new_temp(suffix=".pdf")
+    p2 = workspace.new_temp(suffix=".zip")
+    assert os.path.exists(p1) and os.path.exists(p2)
+    assert workspace.tracked_count() == before + 2
+
+    removed = workspace.clear_all()
+    assert removed >= 2
+    assert not os.path.exists(p1) and not os.path.exists(p2)
+    assert workspace.tracked_count() == 0
+
+
+def test_discard_unlinks_and_untracks():
+    p = workspace.new_temp(suffix=".pdf")
+    before = workspace.tracked_count()
+    assert workspace.discard(p) is True
+    assert not os.path.exists(p)
+    assert workspace.tracked_count() == before - 1
+    # idempotent: discarding again is harmless
+    assert workspace.discard(p) is False
+
+
+def test_clear_all_tolerates_missing_file():
+    p = workspace.new_temp(suffix=".tmp")
+    os.remove(p)                      # vanish underneath the registry
+    workspace.clear_all()             # must not raise
+    assert workspace.tracked_count() == 0