pathseed 0.1.0__py3-none-any.whl

pathseed/__init__.py

@@ -0,0 +1,16 @@
+"""pathseed: from an unlabeled whole-slide image to a detector for a rare target.
+
+The loop:  seed (point prompts -> masks)  ->  train a proposal model  ->  sweep the
+unlabeled pool  ->  expert review (accept / reject)  ->  retrain  ->  ...
+
+Everything is parameterised by a `TargetSpec` (what the target is: which mask
+channels, size window, a relation between channels, and how a proposal is matched
+to ground truth), so the same code runs for a tau halo, a mitotic figure, or any
+other rare object that a pathologist can point at.
+"""
+from .target import TargetSpec, Channel, Relation, Match
+from .manifest import Manifest, Item, Label
+from .budget import Budget
+
+__version__ = "0.1.0"
+__all__ = ["TargetSpec", "Channel", "Relation", "Match", "Manifest", "Item", "Label", "Budget"]

pathseed/app.py

@@ -0,0 +1,103 @@
+"""The one call a user makes:
+
+    from pathseed.app import start
+    start("/content/dataset", state="/content/state", slides_dir="/content/drive/MyDrive/svs", target="tau")
+
+It builds (or reuses) the SAM 3 segmenter, stops a previous server on the same
+port, picks a free port if needed, serves the app, and prints the URL, which on
+Colab is the kernel-proxied one that opens straight from the notebook. Restart
+as often as you like; the model stays loaded and the project bank is on disk.
+
+    stop()                  # stop every server started here
+    url()                   # the last URL again
+"""
+from __future__ import annotations
+
+import os
+import socket
+from typing import Optional
+
+from . import target as T
+from .review import serve, SERVERS
+
+_SEGMENTERS: dict = {}
+_LAST_URL: Optional[str] = None
+
+
+def _in_colab() -> bool:
+    try:
+        import google.colab  # noqa: F401
+        return True
+    except ImportError:
+        return False
+
+
+def _free_port(start: int) -> int:
+    for p in range(start, start + 50):
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            if s.connect_ex(("127.0.0.1", p)) != 0:
+                return p
+    raise RuntimeError("no free port found")
+
+
+def segmenter(kind: str = "sam3", cache_dir: Optional[str] = None, model_id: str = "facebook/sam3", device: Optional[str] = None):
+    """A cached segmenter: 'sam3' (needs the sam3 extra and the checkpoint),
+    'threshold' (no model, for dry runs) or None."""
+    if kind is None or kind == "none":
+        return None
+    key = (kind, cache_dir, model_id)
+    if key not in _SEGMENTERS:
+        if kind == "sam3":
+            from .seed import Sam3PointSegmenter
+            cache_dir = cache_dir or os.environ.get("PATHSEED_SAM3_CACHE")
+            print("loading SAM 3 …", flush=True)
+            _SEGMENTERS[key] = Sam3PointSegmenter(model_id=model_id, cache_dir=cache_dir,
+                                                  local_files_only=bool(cache_dir), device=device)
+        elif kind == "threshold":
+            from .seed import ThresholdSegmenter
+            _SEGMENTERS[key] = ThresholdSegmenter()
+        else:
+            raise ValueError(kind)
+    return _SEGMENTERS[key]
+
+
+def start(dataset: str, state: str, slides_dir: Optional[str] = None, target="tau", round_dir: Optional[str] = None,
+          port: int = 8765, seg: str = "sam3", sam3_cache: Optional[str] = None, host: str = "127.0.0.1", quiet: bool = False):
+    """Serve the app and print its URL. Returns the server."""
+    global _LAST_URL
+    tgt = T.get(target) if isinstance(target, str) else target
+    os.makedirs(state, exist_ok=True)
+    if not os.path.exists(os.path.join(state, "target.json")):
+        tgt.save(os.path.join(state, "target.json"))
+    segm = segmenter(seg, cache_dir=sam3_cache)
+    if port in SERVERS:
+        stop(port)
+    if _free_port(port) != port:
+        newp = _free_port(port + 1)
+        print(f"port {port} is busy (not ours); using {newp}")
+        port = newp
+    srv = serve(dataset, round_dir, tgt, segmenter=segm, port=port, host=host, state=state, slides_dir=slides_dir)
+    _LAST_URL = url(port)
+    if not quiet:
+        print(f"pathseed is up → {_LAST_URL}")
+    return srv
+
+
+def url(port: int = 8765) -> str:
+    if _in_colab():
+        try:
+            from google.colab.output import eval_js
+            return str(eval_js(f"google.colab.kernel.proxyPort({port})"))
+        except Exception:
+            return f"(Colab) run: from google.colab.output import eval_js; eval_js('google.colab.kernel.proxyPort({port})')"
+    return f"http://127.0.0.1:{port}/"
+
+
+def stop(port: Optional[int] = None) -> None:
+    for p in list(SERVERS):
+        if port is None or p == port:
+            s = SERVERS.pop(p)
+            try:
+                s.shutdown(); s.server_close()
+            except Exception:
+                pass

pathseed/budget.py

@@ -0,0 +1,52 @@
+"""Annotator cost accounting. One unit = one annotator action: placing one point
+prompt, or making one accept/reject decision. A wall-clock estimate is derived
+with per-action seconds that can be re-fitted from a timed session.
+
+Measured on the tau seed session (slide 20): 64 boxes, 968 points -> ~15 points
+per instance. Review decisions in the grid UI are about 2 s each."""
+from __future__ import annotations
+
+from dataclasses import dataclass, field, asdict
+
+
+@dataclass
+class Budget:
+    seed_points: int = 0          # point prompts placed
+    seed_instances: int = 0       # instances masked from those prompts
+    seed_negatives: int = 0       # empty tiles picked as negatives
+    decisions: int = 0            # accept / reject decisions on proposals
+    accepted: int = 0
+    rejected: int = 0
+    sec_per_point: float = 3.0
+    sec_per_negative: float = 4.0
+    sec_per_decision: float = 2.0
+    log: list = field(default_factory=list)
+
+    @property
+    def actions(self) -> int:
+        return self.seed_points + self.seed_negatives + self.decisions
+
+    @property
+    def minutes(self) -> float:
+        return (self.seed_points * self.sec_per_point + self.seed_negatives * self.sec_per_negative
+                + self.decisions * self.sec_per_decision) / 60.0
+
+    def add_seed(self, instances: int, points_per_instance: float, negatives: int = 0, note: str = "") -> None:
+        pts = int(round(instances * points_per_instance))
+        self.seed_points += pts
+        self.seed_instances += instances
+        self.seed_negatives += negatives
+        self.log.append({"step": "seed", "instances": instances, "points": pts, "negatives": negatives, "note": note})
+
+    def add_review(self, accepted: int, rejected: int, note: str = "") -> None:
+        self.decisions += accepted + rejected
+        self.accepted += accepted
+        self.rejected += rejected
+        self.log.append({"step": "review", "accepted": accepted, "rejected": rejected, "note": note})
+
+    def snapshot(self) -> dict:
+        d = asdict(self)
+        d.pop("log")
+        d["actions"] = self.actions
+        d["minutes"] = round(self.minutes, 1)
+        return d

pathseed/cli.py

@@ -0,0 +1,218 @@
+"""Command line.
+
+    pathseed folds     --dataset D --out folds.json [--k 5] [--tau-fixed]
+    pathseed simulate  --dataset D --target tau --folds folds.json --test-fold 0 --out runs/
+                       [--model prototype|unet] [--seed-pos 30] [--seed-neg 30] [--rounds 2]
+                       [--review-cap 300] [--seeds 0,1,2] [--random] [--full] [--opt k=v ...]
+    pathseed summarise --runs runs/ [--csv out.csv]
+    pathseed sweep     --source slide.svs --model-path round_0/model.pt --target tau --out pool/ [--polys r.json]
+    pathseed review    --dataset D --round round_0 --target tau [--sam3] [--port 8765]
+    pathseed bench     --dataset D --target tau --folds folds.json --model-path m.pt --out results/
+"""
+from __future__ import annotations
+
+import argparse
+import glob
+import json
+import os
+import sys
+
+
+def _opts(pairs):
+    out = {}
+    for p in pairs or []:
+        k, v = p.split("=", 1)
+        try:
+            v = json.loads(v)
+        except json.JSONDecodeError:
+            pass
+        out[k] = v
+    return out
+
+
+def cmd_folds(a):
+    from .manifest import Manifest
+    from .folds import make_folds, save_folds, TAU_FIXED
+    man = Manifest.load(a.dataset)
+    folds = make_folds(man, a.k, TAU_FIXED if a.tau_fixed else None, by=a.by)
+    save_folds(a.out, folds)
+    for f in folds:
+        pos = sum(1 for i in f["test_ids"] if man.get(i).kind == "pos")
+        where = f.get("test_domains") or f["test_slides"]
+        print(f"fold {f['fold']}: {where if len(where) < 8 else str(len(where)) + ' slides'} tiles {len(f['test_ids'])} pos {pos}")
+
+
+def _split(a):
+    from .manifest import Manifest
+    from .folds import load_folds
+    man = Manifest.load(a.dataset)
+    if getattr(a, "test_split", None):                     # manifest's own split column (e.g. MIDOG++ test)
+        test_ids = [it.id for it in man if it.split == a.test_split]
+    else:
+        folds = load_folds(a.folds)
+        test = next(f for f in folds if int(f["fold"]) == a.test_fold)
+        test_ids = list(test["test_ids"])
+    ts = set(test_ids)
+    pool_ids = [it.id for it in man if it.id not in ts]
+    return man, pool_ids, test_ids
+
+
+def cmd_simulate(a):
+    from . import target as T
+    from .simulate import SimConfig, run, random_baseline, upper_bound
+    man, pool, test = _split(a)
+    tgt = T.get(a.target)
+    opts = _opts(a.opt)
+    seeds = [int(s) for s in a.seeds.split(",")]
+    for s in seeds:
+        cfg = SimConfig(model=a.model, model_opts=opts, seed_pos=a.seed_pos, seed_neg=a.seed_neg, rounds=a.rounds,
+                        review_cap=a.review_cap, review_mix=a.review_mix, seed=s, pool_limit=a.pool_limit,
+                        test_limit=a.test_limit, save_models=a.save_models)
+        res = run(man, tgt, cfg, pool, test, out_dir=a.out)
+        if a.random:
+            for rec in res["records"]:
+                random_baseline(man, tgt, cfg, pool, test, actions=rec["budget"]["actions"], out_dir=a.out)
+    if a.full:
+        cfg = SimConfig(model=a.model, model_opts=opts, seed=seeds[0], test_limit=a.test_limit)
+        upper_bound(man, tgt, cfg, pool, test, out_dir=a.out)
+
+
+def cmd_summarise(a):
+    from .simulate import summarise
+    results = []
+    for p in sorted(glob.glob(os.path.join(a.runs, "*.json"))):
+        with open(p) as f:
+            results.append(json.load(f))
+    rows = summarise(results)
+    hdr = ["kind", "model", "seed_pos", "review_cap", "round", "n", "actions", "object_f1", "object_sd", "object_ap",
+           "tile_f1", "tile_sd", "tile_ap", "found", "positives", "false_alarms"]
+    print("\t".join(hdr))
+    for r in rows:
+        print("\t".join(f"{r[h]:.3f}" if isinstance(r[h], float) else str(r[h]) for h in hdr))
+    if a.csv:
+        import csv
+        with open(a.csv, "w", newline="") as f:
+            w = csv.DictWriter(f, fieldnames=hdr); w.writeheader()
+            for r in rows:
+                w.writerow({h: r[h] for h in hdr})
+    if a.plot:
+        from .plot import plot_curves
+        plot_curves(rows, a.plot)
+        print(f"plot -> {a.plot}")
+
+
+def _load_model(path):
+    if path.endswith(".pt"):
+        from .models.unet import UNetModel
+        return UNetModel.load(path)
+    from .models.prototype import PrototypeModel
+    return PrototypeModel.load(path)
+
+
+def cmd_sweep(a):
+    from . import target as T
+    from .sweep import open_source, sweep, load_polys
+    tgt = T.get(a.target)
+    model = _load_model(a.model_path)
+    src = open_source(a.source)
+    sid = a.slide_id or os.path.splitext(os.path.basename(a.source))[0]
+    polys = load_polys(a.polys) if a.polys else None
+    sweep(src, model, tgt, sid, a.out, polys=polys, batch=a.batch, max_tiles=a.max_tiles)
+
+
+def cmd_review(a):
+    from . import target as T
+    from .review import serve
+    tgt = T.get(a.target)
+    seg = None
+    if a.sam3:
+        from .seed import Sam3PointSegmenter
+        seg = Sam3PointSegmenter(cache_dir=a.sam3_cache, local_files_only=bool(a.sam3_cache))
+    elif a.threshold_segmenter:
+        from .seed import ThresholdSegmenter
+        seg = ThresholdSegmenter()
+    serve(a.dataset, a.round, tgt, segmenter=seg, port=a.port, host=a.host, block=True, state=a.state, slides_dir=a.slides)
+
+
+def cmd_encoders(a):
+    from .models.encoders import describe
+    for spec in a.specs:
+        try:
+            d = describe(spec, a.size)
+            print(f"{spec}: {d['params_M']} M params; " + ", ".join(f"/{l['stride']}:{l['channels']}ch {l['shape'][0]}x{l['shape'][1]}" for l in d["levels"]))
+        except Exception as e:
+            print(f"{spec}: FAILED {type(e).__name__}: {e}")
+
+
+def cmd_bench(a):
+    from . import target as T
+    from .folds import load_folds
+    from .manifest import Manifest
+    from .metrics import score, fold_record
+    from .simulate import _predict_proposals
+    tgt = T.get(a.target)
+    man = Manifest.load(a.dataset)
+    model = _load_model(a.model_path)
+    os.makedirs(a.out, exist_ok=True)
+    for f in load_folds(a.folds):
+        if a.fold is not None and int(f["fold"]) != a.fold:
+            continue
+        props = _predict_proposals(model, man, tgt, f["test_ids"], a.batch, keep_masks=True)
+        sc = score(props, man, f["test_ids"], tgt)
+        rec = fold_record(a.name, f["fold"], f["test_slides"], {"target": tgt.to_dict()}, sc)
+        with open(os.path.join(a.out, f"fold{f['fold']}.json"), "w") as fh:
+            json.dump(rec, fh, indent=2)
+        print(f"fold {f['fold']}: object {sc['object']}  tile {sc['tile']}")
+
+
+def main(argv=None):
+    ap = argparse.ArgumentParser(prog="pathseed")
+    sp = ap.add_subparsers(dest="cmd", required=True)
+
+    p = sp.add_parser("folds"); p.add_argument("--dataset", required=True); p.add_argument("--out", required=True)
+    p.add_argument("--k", type=int, default=5); p.add_argument("--tau-fixed", action="store_true")
+    p.add_argument("--by", choices=["slide", "domain"], default="slide", help="domain = all slides of a tumor type / scanner in one fold")
+    p.set_defaults(fn=cmd_folds)
+
+    p = sp.add_parser("simulate")
+    p.add_argument("--dataset", required=True); p.add_argument("--target", default="tau")
+    p.add_argument("--folds"); p.add_argument("--test-fold", type=int, default=0)
+    p.add_argument("--test-split", help="use manifest items with this split value as the test set instead of a fold")
+    p.add_argument("--out", required=True); p.add_argument("--model", default="prototype")
+    p.add_argument("--seed-pos", type=int, default=30); p.add_argument("--seed-neg", type=int, default=30)
+    p.add_argument("--rounds", type=int, default=2); p.add_argument("--review-cap", type=int, default=300)
+    p.add_argument("--review-mix", type=float, default=1.0, help="fraction of the review queue from the top of the ranking; rest from the uncertain bottom")
+    p.add_argument("--seeds", default="0"); p.add_argument("--random", action="store_true")
+    p.add_argument("--full", action="store_true"); p.add_argument("--pool-limit", type=int, default=0)
+    p.add_argument("--test-limit", type=int, default=0); p.add_argument("--save-models", action="store_true")
+    p.add_argument("--opt", nargs="*", help="model options k=v (json values), e.g. epochs=20 encoder=fastvit_t8 freeze_encoder=true")
+    p.set_defaults(fn=cmd_simulate)
+
+    p = sp.add_parser("summarise"); p.add_argument("--runs", required=True); p.add_argument("--csv")
+    p.add_argument("--plot"); p.set_defaults(fn=cmd_summarise)
+
+    p = sp.add_parser("sweep"); p.add_argument("--source", required=True); p.add_argument("--model-path", required=True)
+    p.add_argument("--target", default="tau"); p.add_argument("--out", required=True); p.add_argument("--polys")
+    p.add_argument("--slide-id"); p.add_argument("--batch", type=int, default=64); p.add_argument("--max-tiles", type=int, default=0)
+    p.set_defaults(fn=cmd_sweep)
+
+    p = sp.add_parser("review"); p.add_argument("--dataset", required=True); p.add_argument("--round")
+    p.add_argument("--target", default="tau"); p.add_argument("--sam3", action="store_true"); p.add_argument("--sam3-cache")
+    p.add_argument("--threshold-segmenter", action="store_true"); p.add_argument("--port", type=int, default=8765)
+    p.add_argument("--host", default="127.0.0.1"); p.add_argument("--state", help="project state dir (project.db)")
+    p.add_argument("--slides", help="directory of slides the server may open"); p.set_defaults(fn=cmd_review)
+
+    p = sp.add_parser("encoders", help="show the pyramid each encoder gives the decoder")
+    p.add_argument("specs", nargs="+"); p.add_argument("--size", type=int, default=256); p.set_defaults(fn=cmd_encoders)
+
+    p = sp.add_parser("bench"); p.add_argument("--dataset", required=True); p.add_argument("--target", default="tau")
+    p.add_argument("--folds", required=True); p.add_argument("--model-path", required=True); p.add_argument("--out", required=True)
+    p.add_argument("--name", default="model"); p.add_argument("--fold", type=int); p.add_argument("--batch", type=int, default=64)
+    p.set_defaults(fn=cmd_bench)
+
+    a = ap.parse_args(argv)
+    a.fn(a)
+
+
+if __name__ == "__main__":
+    sys.exit(main())

pathseed/folds.py

@@ -0,0 +1,75 @@
+"""Slide-held-out folds: every slide lives wholly in one fold. Greedy assignment
+by positive count then total count keeps the folds roughly balanced even when a
+single slide carries a third of the positives (as slide 17 does for tau)."""
+from __future__ import annotations
+
+import json
+from typing import Optional
+
+from .manifest import Manifest
+
+
+def slide_fold_map(manifest: Manifest, k: int = 5, fixed: Optional[dict] = None) -> dict[str, int]:
+    counts: dict[str, dict] = {}
+    for it in manifest:
+        c = counts.setdefault(it.slide, {"pos": 0, "bg": 0})
+        c[it.kind] += 1
+    if fixed and set(fixed) == set(counts):
+        return {s: int(f) for s, f in fixed.items()}
+    fold = [{"pos": 0, "tot": 0} for _ in range(k)]
+    out = {}
+    for s, c in sorted(counts.items(), key=lambda kv: (kv[1]["pos"], kv[1]["pos"] + kv[1]["bg"]), reverse=True):
+        i = min(range(k), key=lambda j: (fold[j]["pos"], fold[j]["tot"]))
+        out[s] = i
+        fold[i]["pos"] += c["pos"]
+        fold[i]["tot"] += c["pos"] + c["bg"]
+    return out
+
+
+def domain_fold_map(manifest: Manifest, k: Optional[int] = None) -> dict[str, int]:
+    """Every slide of a domain in the same fold. k=None -> one fold per domain
+    (leave-one-domain-out); otherwise domains are packed greedily into k folds."""
+    dom_of = {}; counts: dict[str, dict] = {}
+    for it in manifest:
+        d = it.domain or "?"
+        dom_of[it.slide] = d
+        c = counts.setdefault(d, {"pos": 0, "tot": 0}); c["pos"] += it.kind == "pos"; c["tot"] += 1
+    doms = sorted(counts, key=lambda d: -counts[d]["pos"])
+    if k is None or k >= len(doms):
+        dfold = {d: i for i, d in enumerate(doms)}
+    else:
+        load = [0] * k; dfold = {}
+        for d in doms:
+            i = min(range(k), key=lambda j: load[j]); dfold[d] = i; load[i] += counts[d]["pos"]
+    return {s: dfold[d] for s, d in dom_of.items()}
+
+
+def make_folds(manifest: Manifest, k: int = 5, fixed: Optional[dict] = None, by: str = "slide") -> list[dict]:
+    if by == "domain":
+        fm = domain_fold_map(manifest, k)
+        k = max(fm.values()) + 1
+    else:
+        fm = slide_fold_map(manifest, k, fixed)
+    out = []
+    for f in range(k):
+        slides = sorted([s for s, ff in fm.items() if ff == f], key=lambda s: (0, int(s)) if s.isdigit() else (1, s))
+        ids = [it.id for it in manifest if fm[it.slide] == f]
+        doms = sorted({it.domain for it in manifest if fm[it.slide] == f and it.domain})
+        out.append({"fold": f, "test_slides": slides, "test_ids": ids, **({"test_domains": doms} if doms else {})})
+    return out
+
+
+def save_folds(path: str, folds: list[dict]) -> None:
+    with open(path, "w") as f:
+        json.dump({"folds": folds}, f)
+
+
+def load_folds(path: str) -> list[dict]:
+    with open(path) as f:
+        d = json.load(f)
+    return d["folds"] if isinstance(d, dict) else d
+
+
+# the published 17-slide tau assignment, so results stay comparable
+TAU_FIXED = {"3": 0, "9": 0, "17": 0, "1": 1, "2": 1, "5": 1, "10": 1, "4": 2, "8": 2, "11": 2,
+             "14": 3, "16": 3, "20": 3, "6": 4, "7": 4, "13": 4, "18": 4}

pathseed/loop.py

@@ -0,0 +1,198 @@
+"""The real loop, with a person in the review seat. State lives in one directory:
+
+    <state>/target.json          the TargetSpec
+    <state>/labels.json          the labelled set (grows every round)
+    <state>/budget.json          annotator cost so far
+    <state>/round_<k>/model.*    proposal model trained on the labelled set
+    <state>/round_<k>/proposals.json    what the sweep found on the pool
+    <state>/round_<k>/review.json       the reviewer's decisions (written by the UI)
+
+Tiles (images + masks) live in a dataset directory with the layout of
+`pathseed.manifest`; the seeding step and the sweep both append to it, so the
+labelled set only ever refers to tile ids.
+
+    loop = Loop.init(state, dataset, target)
+    loop.add_labels(seed_labels)          # from pathseed.seed
+    loop.train(k)                         # -> round_k/model
+    loop.propose(k, pool_ids)             # -> round_k/proposals.json
+    ... review in the browser (pathseed.review.serve) ...
+    loop.ingest_review(k)                 # decisions -> labels.json, budget.json
+"""
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import asdict
+from typing import Iterable, Optional
+
+from .budget import Budget
+from .manifest import Manifest, Label, save_labels, load_labels
+from .models import make as make_model
+from .propose import proposals_from_probs
+from .target import TargetSpec
+
+
+class Loop:
+    def __init__(self, state: str, dataset: str, target: TargetSpec):
+        self.state, self.dataset_dir, self.target = state, dataset, target
+        self.manifest = Manifest.load(dataset)
+        self.labels: list[Label] = load_labels(self.labels_path) if os.path.exists(self.labels_path) else []
+        self.budget = self._load_budget()
+
+    # ------------------------------------------------------------------ paths
+    @property
+    def labels_path(self):
+        return os.path.join(self.state, "labels.json")
+
+    @property
+    def budget_path(self):
+        return os.path.join(self.state, "budget.json")
+
+    def round_dir(self, k: int) -> str:
+        d = os.path.join(self.state, f"round_{k}")
+        os.makedirs(d, exist_ok=True)
+        return d
+
+    # ------------------------------------------------------------------ init/io
+    @classmethod
+    def init(cls, state: str, dataset: str, target: TargetSpec) -> "Loop":
+        os.makedirs(state, exist_ok=True)
+        target.save(os.path.join(state, "target.json"))
+        if not os.path.exists(os.path.join(dataset, "manifest.json")):
+            Manifest(dataset, [], tile=target.tile).save()
+        return cls(state, dataset, target)
+
+    @classmethod
+    def open(cls, state: str, dataset: str) -> "Loop":
+        return cls(state, dataset, TargetSpec.load(os.path.join(state, "target.json")))
+
+    def _load_budget(self) -> Budget:
+        if os.path.exists(self.budget_path):
+            with open(self.budget_path) as f:
+                d = json.load(f)
+            b = Budget(**{k: v for k, v in d.items() if k in Budget.__dataclass_fields__})
+            return b
+        return Budget()
+
+    def save(self) -> None:
+        save_labels(self.labels_path, self.labels)
+        with open(self.budget_path, "w") as f:
+            json.dump(asdict(self.budget), f, indent=1)
+
+    def reload_manifest(self) -> None:
+        self.manifest = Manifest.load(self.dataset_dir)
+
+    # ------------------------------------------------------------------ project bank
+    @property
+    def project(self):
+        from .project import Project
+        if not hasattr(self, "_project"):
+            self._project = Project(os.path.join(self.state, "project.db"))
+        return self._project
+
+    def sync_project(self) -> dict:
+        """Pull the labelled set from the project bank (what the browser
+        recorded: seed tiles, review decisions) and its budget ledger."""
+        self.reload_manifest()
+        labs = self.project.labels(self.dataset_dir, self.target.channels)
+        have = {l.id for l in self.labels}
+        new = [l for l in labs if l.id not in have]
+        self.labels.extend(new)
+        b = self.project.budget()
+        self.budget.seed_points = b["seed_points"]; self.budget.seed_instances = b["seed_instances"]
+        self.budget.seed_negatives = b["seed_negatives"]; self.budget.decisions = b["decisions"]
+        self.budget.accepted = b["accepted"]; self.budget.rejected = b["rejected"]
+        self.save()
+        return {"added": len(new), "labelled": len(self.labels), "budget": self.budget.snapshot()}
+
+    # ------------------------------------------------------------------ steps
+    def add_labels(self, labels: Iterable[Label], seed_instances: int = 0, seed_negatives: int = 0) -> None:
+        have = {l.id for l in self.labels}
+        new = [l for l in labels if l.id not in have]
+        self.labels.extend(new)
+        if seed_instances or seed_negatives:
+            self.budget.add_seed(seed_instances, self.target.points_per_instance, seed_negatives, note="seed")
+        self.save()
+
+    def train(self, k: int, model: str = "unet", log=print, **opts):
+        if log is None:
+            log = lambda *a: None  # noqa: E731
+        self.reload_manifest()
+        n_pos = sum(1 for l in self.labels if l.kind == "pos")
+        if os.path.exists(os.path.join(self.state, "project.db")):
+            self.project.start_round(k, model, len(self.labels), n_pos)
+        m = make_model(model, **opts).fit(self.manifest, self.target, self.labels, log=log)
+        m.save(os.path.join(self.round_dir(k), "model." + ("pt" if model in ("unet", "dualmask") else "json")))
+        if os.path.exists(os.path.join(self.state, "project.db")):
+            self.project.finish_round(k, {})
+        with open(os.path.join(self.round_dir(k), "train.json"), "w") as f:
+            json.dump({"model": model, "opts": opts, "labelled": len(self.labels),
+                       "labelled_pos": sum(1 for l in self.labels if l.kind == "pos")}, f, indent=1)
+        return m
+
+    def load_model(self, k: int):
+        d = self.round_dir(k)
+        if os.path.exists(os.path.join(d, "model.pt")):
+            from .models.unet import UNetModel
+            return UNetModel.load(os.path.join(d, "model.pt"))
+        from .models.prototype import PrototypeModel
+        return PrototypeModel.load(os.path.join(d, "model.json"))
+
+    def propose(self, k: int, pool_ids: Optional[Iterable[str]] = None, model=None, cap: int = 0,
+                batch: int = 64, mix: float = 1.0, log=print) -> dict:
+        """Score the pool (default: every dataset tile not yet labelled), keep tiles
+        that fire, rank by best proposal score, write round_k/proposals.json."""
+        self.reload_manifest()
+        model = model or self.load_model(k)
+        labelled = {l.id for l in self.labels}
+        pool = [i for i in (pool_ids if pool_ids is not None else (it.id for it in self.manifest)) if i not in labelled]
+        tiles = {}
+        for tid, probs in model.predict(self.manifest, pool, batch=batch):
+            props = proposals_from_probs(tid, probs, self.target, keep_masks=False)
+            if props:
+                tiles[tid] = [p.to_dict() for p in props]
+        from .simulate import select_for_review
+        ranked = sorted(tiles, key=lambda t: -max(p["score"] for p in tiles[t]))
+        if cap:
+            ranked = [t for _, t in select_for_review([(max(p["score"] for p in tiles[t]), t) for t in ranked], cap, mix)]
+        out = {"round": k, "pool": len(pool), "fired": len(tiles), "shown": len(ranked),
+               "tiles": [{"id": t, "score": max(p["score"] for p in tiles[t]), "proposals": tiles[t]} for t in ranked]}
+        with open(os.path.join(self.round_dir(k), "proposals.json"), "w") as f:
+            json.dump(out, f)
+        log(f"round {k}: pool {len(pool)} tiles, {len(tiles)} fired, {len(ranked)} queued for review")
+        return out
+
+    def ingest_review(self, k: int, accepted_source: str = "gt") -> dict:
+        """Read round_k/review.json ({id: "accept"|"reject"}) into the labelled set.
+        Accepted tiles take masks from `accepted_source`: "gt" when the dataset
+        holds (corrected) masks for them, else the proposal masks saved by the
+        review server under round_k/masks/<channel>/<id>.png."""
+        d = self.round_dir(k)
+        self.reload_manifest()
+        with open(os.path.join(d, "review.json")) as f:
+            rev = json.load(f)
+        have = {l.id for l in self.labels}
+        acc = rej = 0
+        for tid, dec in rev.items():
+            if tid in have:
+                continue
+            if dec == "accept":
+                if accepted_source == "gt":
+                    self.labels.append(Label(tid, "pos", "gt", origin="review", round=k + 1))
+                else:
+                    # the proposal masks: written by the sweep into the dataset's channel dirs
+                    # (or, when a reviewer edited them, under round_k/masks/)
+                    files = {}
+                    for c in self.target.channels:
+                        edited = os.path.join(d, "masks", c.dir, f"{tid}.png")
+                        files[c.name] = edited if os.path.exists(edited) else self.manifest.mask_path(c.dir, tid)
+                        if not os.path.exists(files[c.name]):
+                            raise FileNotFoundError(f"accepted tile {tid} has no {c.name} mask at {files[c.name]}")
+                    self.labels.append(Label(tid, "pos", "files", files, origin="review", round=k + 1))
+                acc += 1
+            elif dec == "reject":
+                self.labels.append(Label(tid, "bg", "empty", origin="review", round=k + 1))
+                rej += 1
+        self.budget.add_review(acc, rej, note=f"round {k} review")
+        self.save()
+        return {"accepted": acc, "rejected": rej, "labelled": len(self.labels), "budget": self.budget.snapshot()}