"""Browse, verify, fetch and benchmark spaCR segmentation and classification models.
Scope
-----
spaCR can run a stock Cellpose model, a fine-tuned Cellpose model, a
classifier checkpoint from a training run, or a checkpoint downloaded from
Hugging Face. Three questions follow for any machine: which of those models
are on it, what each one was trained on, and whether the copy on disk is the
file the author published.
This module answers those three questions and nothing else:
* **what is here** — :func:`discover_local` walks folders and returns the
checkpoints, classified as Cellpose or classifier, with whatever provenance
is recoverable from the settings snapshots spaCR already writes;
* **is it the right bytes** — :func:`sha256_file` / :func:`verify` /
:func:`fetch`, which downloads atomically, checksums what arrived, and
refuses to install a mismatch;
* **what does it do on my data** — :func:`benchmark`, which is
:mod:`spacr.model_compare`'s "three fields" harness pointed at one model
instead of two.
Nothing here imports torch or cellpose at module import time. Browsing the zoo,
reading provenance, checksumming a file and rendering the table all work on a
machine with neither installed; only :func:`benchmark` (through
:func:`spacr.model_compare.segment_with_cellpose`) needs them, and only when it
is called. ``tests/test_model_zoo.py`` asserts that.
Safeguards
----------
**Every download is checksummed, and verification happens before use.**
A checkpoint truncated by a dropped connection, or swapped for a different
one at the same URL, still loads. It does not raise; it produces silently
different masks, and the run that used it looks exactly like the run that
did not. So :func:`fetch` hashes what arrived and compares it to the hash
the catalogue published: a mismatch deletes the file and raises
:class:`ChecksumMismatch` — the entry is never registered. The hash that
was actually computed is stored on the returned entry, so "verified"
means "these bytes", not "this filename".
A catalogue entry with no published hash cannot be verified at all, and
that is refused by default (``require_checksum=True``) rather than quietly
treated as fine. Callers that knowingly accept an unverifiable source pass
``require_checksum=False``; the resulting entry carries
``verified=False`` and says so in :func:`format_zoo`.
**Every download is atomic, and never overwrites.**
Bytes stream into a temporary file *in the destination directory* (so the
final ``os.replace`` is a same-filesystem rename, which is atomic) and the
rename happens only after the checksum passes. An interrupted or cancelled
download therefore leaves nothing behind that looks like a model — the
failure mode where half a checkpoint sits at the real filename, loads, and
segments badly, cannot happen.
The destination is versioned rather than overwritten
(:func:`versioned_path`: ``foo.CP_model``, ``foo_v2.CP_model``, …). Two
models with the same filename are a normal thing to have; losing the first
one to the second is not.
**Provenance is recorded, and "unknown" is written out in full.**
A Cellpose model fine-tuned on 60x confluent HeLa is not interchangeable
with one trained on 20x sparse fibroblasts, whatever their benchmark
scores. :class:`ModelEntry` therefore carries
:attr:`~ModelEntry.trained_on` and :attr:`~ModelEntry.trained_by`, both
recovered from the settings snapshots spaCR already writes beside its
models (``<file>_settings.csv`` for a Cellpose model,
``<dst>/settings.csv`` for a classifier — read through
:func:`spacr.train_compare.load_run`, which already knows where to look).
Where it could not be recovered the field reads ``'unknown'``, never ``''``.
A blank cell in a provenance table reads as "no constraints"; that is the
opposite of what it means.
**Benchmarks are only comparable inside one field set.**
A model's score on your three fields does not predict its score on
somebody else's, and a table that sorts the two together produces a
ranking of two unrelated numbers. So every :class:`BenchmarkResult`
records a :func:`fieldset_id` — a hash of the actual pixels, not the folder
name — and :func:`rank` **raises** :class:`IncomparableBenchmarks` when
handed results from more than one field set. :func:`rank_groups` and
:func:`format_benchmarks` are the supported alternative: they group by
field set, rank within each group, and label the groups.
A model file that is missing, empty, or not a torch
checkpoint at all fails in :func:`inspect_checkpoint` with a message naming the
file, before anything tries to load it. The default failure — a ``KeyError`` on
a state-dict key from deep inside torch — names nothing the user chose.
What the benchmark can and cannot say
-------------------------------------
There is no ground truth here. :func:`benchmark` runs one model over N fields
and reports what came out: object counts, timings and the
:mod:`spacr.seg_qc` verdict per field (fused? shattered? empty? all on the
border?). That is a *quality-control* score, not an accuracy — it can tell you
a model collapsed on your data, it cannot tell you which of two plausible
segmentations is right. :data:`RANK_KEYS` therefore offers exactly two keys,
``'qc'`` and ``'seconds'``, and no key that would read as accuracy. To compare
two models against each other, use :func:`compare_entries`, which hands both to
:func:`spacr.model_compare.compare_models` — the A/B harness that is explicit
about neither side being the truth.
Cellpose 4 accepts and ignores ``model_type``, ``diam_mean``, ``nchan``,
``channels`` and ``rescale``; only ``diameter`` at ``eval`` still changes the
masks. :class:`BenchmarkResult` carries the resolved ``honoured`` and
``ignored`` parameter dicts straight from
:class:`spacr.model_compare.ModelConfig` so a benchmark cannot silently be a
benchmark of settings nothing read.
Example::
from spacr import model_zoo as zoo
entries = zoo.catalogue() + zoo.discover_local('/data/screen1')
print(zoo.format_zoo(entries))
entry = zoo.resolve('cpsam_plaque_r3', entries)
result = zoo.benchmark(entry, source='/data/screen1/plate1/1', n_fields=3)
print(zoo.format_benchmarks([result]))
See Also:
:mod:`spacr.model_compare` — the A/B harness this module reuses for
segmentation and for the two-model comparison.
:mod:`spacr.train_compare` — run discovery and settings recovery, reused
wholesale for the classifier half of the zoo.
:func:`spacr.utils.download_models` — the legacy bulk pull of the bundled
Hugging Face model pack, wrapped by :func:`download_bundled_models`.
"""
from __future__ import annotations
import hashlib
import html
import json
import logging
import os
import re
import threading
import time
from dataclasses import dataclass, replace
from dataclasses import field as _dc_field
from pathlib import Path
from typing import (
Any,
Callable,
Dict,
Iterable,
Iterator,
List,
Mapping,
Optional,
Sequence,
Tuple,
)
import numpy as np
from ._segmentation_backends import _SPECS as _BACKEND_SPECS
from ._segmentation_backends import _BACKEND_NAMES as _SEGMENTATION_BACKEND_NAMES
from .logging_util import _spacr_home
LOG = logging.getLogger(__name__)
__all__ = [
"BUNDLED_REMOTE_MODELS",
"BenchmarkResult",
"CATALOGUE_ENV_VAR",
"RETIRED_MODEL_NAMES",
"REMOTE_CATALOGUE_URI",
"shared_catalogue",
"shared_catalogue_is_stale",
"publish_model",
"CLASSIFIER_SUFFIXES",
"CELLPOSE_SUFFIXES",
"ChecksumMismatch",
"ZOO_SOURCES",
"DEFAULT_ZOO_SOURCES",
"MITONET_KEY",
"THIRD_PARTY_MODEL_KEYS",
"is_third_party",
"mitonet_instances",
"mitonet_weights",
"segment_mitonet",
"source_of",
"group_by_source",
"entries_from_sources",
"DEFAULT_N_FIELDS",
"DEFAULT_SCAN_DEPTH",
"DownloadCancelled",
"FieldBenchmark",
"HF_MODELS_REPO",
"IncomparableBenchmarks",
"KINDS",
"ModelEntry",
"ModelUnreadable",
"ModelZooError",
"RANK_KEYS",
"UNKNOWN",
"benchmark",
"catalogue",
"classify_kind",
"compare_entries",
"default_local_roots",
"discover_local",
"download_bundled_models",
"entry_from_file",
"fetch",
"fieldset_id",
"format_benchmarks",
"format_zoo",
"group_by_fieldset",
"hf_uri",
"inspect_checkpoint",
"install",
"load_catalogue_file",
"open_uri",
"package_model_root",
"rank",
"rank_groups",
"resolve",
"sha256_file",
"verify",
"versioned_path",
]
#: What an unrecoverable provenance field says. Never ``''``: a blank cell in a
#: provenance table reads as "no constraints", which is the opposite of "we do
#: not know what this model was trained on".
UNKNOWN = "unknown"
#: The kinds of model spaCR runs.
#:
#: ``detector`` was added for the YOLO well detector, which is neither of the
#: first two: it does not segment and it does not classify a crop, it locates
#: regions so something else can. :class:`ModelEntry` VALIDATES against this
#: tuple and raises on anything else, so an entry naming a kind that is not
#: here fails at construction rather than being quietly filed as a Cellpose
#: model and handed to CellposeModel later.
#:
#: ``cellpose3`` is a model that runs through the Cellpose 3 backend -- its
#: cyto3, cyto2, cyto and nuclei models, and Cellpose-format checkpoints from
#: bioimage.io. It is a kind of its own because spaCR's Cellpose 4 loads such
#: a checkpoint without complaint and then segments nonsense with it.
#:
#: ``cellpose_dino`` is a Cellpose-DINO checkpoint, which runs through the
#: Cellpose-DINO backend (item 525): spaCR's own Cellpose 4 has no DINOv3,
#: and without it Cellpose 4 cannot build the network the weights are for.
#:
#: Each backend an object chooses by a model-setting prefix of its own --
#: StarDist, InstanSeg, Omnipose (items 551-553) -- is a kind named after
#: it (:data:`PREFIXED_KINDS`): its models run in that backend and nowhere
#: else.
PREFIXED_KINDS = tuple(_name for _name, _spec in _BACKEND_SPECS.items()
if _spec.prefix)
KINDS = ("cellpose", "classifier", "detector", "encoder", "backend",
"cellpose3", "cellpose_dino", "mitonet", "timepose") + PREFIXED_KINDS
#: "backend" is not a checkpoint: it is a segmentation PACKAGE the zoo
#: lists so a user learns it exists and can install it from inside spaCR.
#: ``encoder`` is 386's kind: a self-supervised backbone that produces an
#: EMBEDDING rather than a mask or a label. It belongs in the zoo for the
#: reason 386 gives -- "an embedding that ships without one is a black box
#: twice over" -- and it is a separate kind because nothing that consumes a
#: classifier can consume one. Its weights are resolved by `timm` from the
#: HuggingFace hub rather than shipped, so its entry describes a download
#: that already happened; see :func:`spacr.embeddings.encoder_entry`.
#: Filename endings that mark a Cellpose checkpoint. ``.CP_model`` is what
#: :func:`spacr.submodules.train_cellpose` names its output.
CELLPOSE_SUFFIXES = (".cp_model", ".cpmodel")
#: Filename endings that mark a torch checkpoint. A ``.pth`` inside a Cellpose
#: folder is still a Cellpose model — see :func:`classify_kind`.
CLASSIFIER_SUFFIXES = (".pth", ".pt")
#: Directory names that mean "the files in here are Cellpose checkpoints".
#: ``cellpose_model`` is :func:`spacr.submodules.train_cellpose`'s output
#: folder; ``cp`` is where :func:`spacr.utils.download_models` lands the
#: bundled pack; ``models`` is what ``cellpose.train.train_seg`` creates under
#: whatever ``save_path`` it is given.
CELLPOSE_DIR_NAMES = ("cellpose_model", "cp", "models")
#: How deep :func:`discover_local` walks below each root.
DEFAULT_SCAN_DEPTH = 6
#: Ceiling on how many files one :func:`discover_local` call will *examine*, so
#: pointing it at ``/`` is slow rather than fatal.
DEFAULT_SCAN_LIMIT = 20000
#: Fields a benchmark uses by default — the number a human actually looks at.
DEFAULT_N_FIELDS = 3
#: Seconds before a download gives up on the server.
DEFAULT_TIMEOUT = 30
#: Bytes per chunk while streaming a download.
DEFAULT_CHUNK = 1 << 16
#: The Hugging Face dataset repo :func:`spacr.utils.download_models` pulls the
#: bundled model pack from. Same repo, same URL scheme — this module adds the
#: checksum, the atomic write and the versioned destination that one lacks.
HF_MODELS_REPO = "einarolafsson/models"
#: Environment variable naming a JSON catalogue of remote models. See
#: :func:`load_catalogue_file` for the format.
CATALOGUE_ENV_VAR = "SPACR_MODEL_CATALOGUE"
#: The SHARED catalogue, fetched at runtime rather than shipped.
#:
#: WHY THIS IS REMOTE. Until this existed, a model reached other spaCR users
#: only by someone editing :data:`BUNDLED_REMOTE_MODELS` in this file and
#: cutting a release -- so contributing a model meant contributing to spaCR,
#: waiting for a version, and every user upgrading. That is a high price for a
#: row of metadata, and it is why the zoo had exactly one entry.
#:
#: A contributor now uploads to THEIR OWN Hugging Face account -- they keep
#: ownership, and nobody has to hand out write access to anyone else's -- and
#: adds one row here. See :func:`publish_model`, which does the upload and
#: prints the row.
REMOTE_CATALOGUE_URI = (
"https://huggingface.co/datasets/einarolafsson/models/resolve/main/"
"catalogue.json"
)
#: How long a fetched shared catalogue is reused before being re-fetched.
#: Long enough that opening a module repeatedly is not a repeated request,
#: short enough that a newly contributed model appears the same day.
CATALOGUE_CACHE_SECONDS = 3600
#: Remote entries spaCR knows about out of the box.
#:
#: EVERY ENTRY HERE CARRIES A REAL sha256, and that is now the rule rather
#: than an aspiration. The retired ``toxo_plaque_cyto`` entry published none,
#: so :func:`fetch` refused to install it -- correctly, since a truncated or
#: substituted checkpoint could not be told from the real one -- which meant
#: it appeared in the model zoo as a row whose Download button could never
#: succeed. An entry without a hash is not a conservative entry; it is one
#: nobody can install.
BUNDLED_REMOTE_MODELS: Tuple[Dict[str, Any], ...] = (
{
"key": "toxoplasma_pv_v1",
"name": "cpsam_v2_toxo_r2",
"kind": "cellpose",
"repo_id": "einarolafsson/toxoplasma-pv-segmentation-cpsam",
"repo_type": "model",
"uri": None,
"sha256":
"182d8cf6b32c7b9ef2917c85870d188486e5e119f05e9c5c1f07652f6859f2d0",
"metrics": {'n_train': '229', 'train_objects': 'not recorded', 'n_test': '11 wells', 'test_objects': 'not recorded', 'cv': 'no', 'f1': '0.8640', 'aji': '0.8090', 'dice': 'not recorded', 'stock_f1': '0.7130', 'stock_aji': '0.4260', 'stock_dice': 'not recorded', 'train_loss': 'not recorded', 'val_loss': 'not recorded', 'best_epoch': '100 / 100'},
"display_name": "Toxoplasma PV v1",
"architecture": "Cellpose-SAM (cpsam_v2)",
"dataset": "anti-Toxoplasma-biotin and DsRed PV lumen; 229 images "
"from 2 datasets, 104 round-1 and 125 newly curated",
"versus_stock": "F1 0.864 against 0.713 for stock cpsam on 11 "
"held-out in-house wells, at IoU 0.5; literature "
"hold-out pending",
"trained_on": (
"Toxoplasma tachyzoite parasitophorous vacuoles stained with goat "
"anti-Toxoplasma-biotin, and tachyzoites expressing DsRed in the "
"PV lumen. Round 2: 229 training images (round 1's 104 plus 125 "
"newly curated RH and ME49 fields), 100 epochs, base cpsam_v2"
),
"trained_by": "einarolafsson",
"notes": (
"F1 0.864 at IoU 0.5 against 0.713 for stock cpsam on the 11 "
"wells round 1 also held out (round 1 scored 0.867); AJI 0.809 "
"against 0.426",
"accuracy falls sharply above IoU 0.8 -- suited to counting and "
"area rather than precise morphometry",
"the held-out literature scorecard is pending a stock-seeded "
"re-curation; on the current literature set, whose truth leans "
"toward this model's lineage, it ties stock Cellpose-SAM on "
"detection (F1 0.403 against 0.400)",
'Re-scored 2026-10-11 on curated ground truth (spaCR pipeline call, '
'defaults; receipt '
'features/data/model_rescore_curated_gt_2026-10-11/receipt.json): '
'strict hold-out, fields this model never saw, F1 0.793 at IoU 0.5 on n'
' = 486 of 655 fields (n differs between models); common subset unseen '
'by every pv model, n = 163, F1 0.783 (95% bootstrap CI 0.710 to '
'0.830). Without the 98 paged round-3 fields (n = 65): F1 0.732. '
'Earlier numbers in this row are superseded by this re-score where they'
' differ. Default unchanged.',
),
},
{
"key": "toxoplasma_plaque_v1",
"name": "cpsam_plaque_r3",
"kind": "cellpose",
"repo_id": "einarolafsson/toxoplasma-plaque-segmentation-cpsam",
"repo_type": "model",
"uri": None,
"sha256":
"eeecd2d6cd5cbb4dddee71564d5f460d26bb07ac125e0b494b7502fea4292d5d",
"metrics": {'n_train': '184 wells', 'train_objects': 'not recorded', 'n_test': 'per fold', 'test_objects': 'not recorded', 'cv': '3-fold', 'f1': '0.8560 in-domain / 0.8060 literature', 'aji': 'not recorded', 'dice': 'not recorded', 'stock_f1': 'not recorded', 'stock_aji': 'not recorded', 'stock_dice': 'not recorded', 'train_loss': 'not recorded', 'val_loss': 'not recorded', 'best_epoch': 'not recorded'},
"display_name": "Toxoplasma Plaque v1",
"architecture": "Cellpose-SAM (cpsam)",
"dataset": "crystal violet plaque wells; 184 wells from 3 datasets, "
"95 in-house and 89 literature",
"versus_stock": "F1 0.856 in-domain; 0.806 on literature "
"(3-fold cross-validated, SD 0.020)",
"trained_on": (
"Toxoplasma gondii plaque assays; round 3, evaluated in-domain "
"(NAS) and against a literature generalisation set"
),
"trained_by": "einarolafsson",
"notes": (
"F1 0.856 in-domain and 0.806 on the literature set (3-fold "
"cross-validated, SD 0.020), against 0.718 for round 1",
"round 3 trades precision (0.939 down to 0.858) for recall "
"(0.631 up to 0.811) on the literature set, which is the right "
"direction for a counting assay",
"PREFER THIS ONE FOR MICROSCOPE-ONLY WORK. On the round-5 test "
"split it scores 0.836 on PFA-fixed wells against round 5's 0.808, "
"at precision 0.93 against 0.81. For mixed sources, or any phone-"
"camera image, use toxoplasma_plaque_v2, which round 3 cannot "
"handle at all (0.249 there)",
'Re-scored 2026-10-11 on curated ground truth (spaCR pipeline call, '
'defaults; receipt '
'features/data/model_rescore_curated_gt_2026-10-11/receipt.json): '
'strict hold-out, fields this model never saw, F1 0.707 at IoU 0.5 on n'
' = 498 of 562 fields (n differs between models); common subset unseen '
'by every plaque model, n = 114, F1 0.402 (95% bootstrap CI 0.300 to '
'0.504). Without the v3-seeded Gel Doc truth (n = 81): F1 0.472. '
'Earlier numbers in this row are superseded by this re-score where they'
' differ. 263-field set seen counts: v1 54, v2 178, v3 137, v4 0; GRA8 '
'Plate1 (hold-out of v3) is the same physical plate as ReimagedPlate1 '
'(training data of v4). Default unchanged.',
),
},
{
"key": "toxoplasma_plaque_v2",
"name": "cpsam_plaque_r5",
"kind": "cellpose",
"repo_id": "einarolafsson/toxoplasma-plaque-segmentation-cpsam-r5",
"repo_type": "model",
"uri": "https://huggingface.co/einarolafsson/"
"toxoplasma-plaque-segmentation-cpsam-r5/resolve/main/"
"weights/cpsam_plaque_r5?download=true",
"sha256":
"0927023a745ac6a19bae0ec72c89b7b864a4ff8d047a41f3f1e9767e1a4d0600",
"metrics": {'n_train': '332 fields, 4 domains', 'train_objects': '18532', 'n_test': '81 fields', 'test_objects': '4294', 'cv': 'no (grouped train/valid/test)', 'f1': '0.819 literature / 0.876 bigbean / 0.808 patrick / 0.415 malnio', 'aji': 'not recorded', 'dice': 'not recorded', 'stock_f1': 'not measured; scored against round 3 (0.820 literature) and round 4 (0.819)', 'stock_aji': 'not recorded', 'stock_dice': 'not recorded', 'train_loss': '0.1822', 'val_loss': '0.2069', 'best_epoch': '100 / 100'},
"display_name": "Toxoplasma Plaque v2 (round 5)",
"architecture": "Cellpose-SAM (cpsam_v2)",
"dataset": "488 curated fields across four domains -- 298 wells cropped "
"from published figures, 96 phone-camera wells, 67 PFA and 27 "
"methanol-fixed whole-well microscope scans; 27,582 plaques",
"versus_stock": "not scored against stock; on 81 held-out fields it ties "
"round 3 on literature (0.819 vs 0.820) and beats it by "
"0.166 on phone-camera wells (0.415 vs 0.249)",
"trained_on": (
"Toxoplasma plaque assays stained with crystal violet, from three "
"microscopes and from published figures. Round 5: 332 training fields "
"grouped by figure and by plate so none straddles the split, 100 "
"epochs, base cpsam_v2, empty wells kept as negatives"
),
"trained_by": "einarolafsson",
"notes": (
"no longer the default: toxoplasma_plaque_v3 replaced it on "
"2026-10-10 (F1 0.844 against 0.799 on 263 held-out fields; this "
"model scores 0.013 on Gel Doc wells), and toxoplasma_plaque_v4 "
"on 2026-10-11; kept so recorded runs reproduce",
"COMPLEMENTS toxoplasma_plaque_v1 rather than replacing it: prefer v1 "
"(round 3) for microscope-only work, where it scores 0.836 against "
"this model's 0.808 on PFA-fixed wells and is far more precise; "
"prefer this one for mixed or unknown sources",
"the only plaque model trained on phone-camera wells -- F1 0.415 "
"against round 3's 0.249, though recall there is 0.296, so it still "
"misses most plaques on phone images and is not yet a counting tool",
"it did NOT clear the promotion bar of 0.02 literature F1 fixed before "
"the run (it came in at -0.001), so round 3 remains production",
"balanced precision/recall (0.81/0.83) where round 3 is lopsided "
"(0.93/0.73): round 3's low recall systematically UNDERCOUNTS, which "
"matters more than F1 for a counting assay",
"hallucinates 2 objects across 6 blank-lawn wells where round 3 "
"hallucinates 19",
"first plaque model on cpsam_v2; rounds 1-4 used cpsam v1, so base and "
"data changed together and the gap to round 3 is not attributable to "
"the extra curation alone",
"training data: https://huggingface.co/datasets/einarolafsson/"
"toxoplasma-plaque-dataset",
'Re-scored 2026-10-11 on curated ground truth (spaCR pipeline call, '
'defaults; receipt '
'features/data/model_rescore_curated_gt_2026-10-11/receipt.json): '
'strict hold-out, fields this model never saw, F1 0.455 at IoU 0.5 on n'
' = 282 of 562 fields (n differs between models); common subset unseen '
'by every plaque model, n = 114, F1 0.360 (95% bootstrap CI 0.249 to '
'0.472). Without the v3-seeded Gel Doc truth (n = 81): F1 0.406. '
'Earlier numbers in this row are superseded by this re-score where they'
' differ. 263-field set seen counts: v1 54, v2 178, v3 137, v4 0; GRA8 '
'Plate1 (hold-out of v3) is the same physical plate as ReimagedPlate1 '
'(training data of v4). Default unchanged.',
),
},
{
"key": "toxoplasma_well_detector_v1",
"name": "yolo_welldetect_v3.pt",
"kind": "detector",
"repo_id": "einarolafsson/toxoplasma-plaque-well-detector-yolo11",
"repo_type": "model",
"uri": None,
"sha256":
"b826058754fb5d4df36c3a7283aac049015cbb044b5ef096c55d19f37172a50c",
"metrics": {'n_train': '562 images', 'train_objects': 'not recorded', 'n_test': 'held-out split', 'test_objects': 'not recorded', 'cv': 'no', 'f1': 'mAP50 0.9930 on v3\'s own split; 0.8838 on the shared test set', 'aji': 'mAP50-95 0.8860 own split; 0.7630 shared', 'dice': 'P/R 0.9870 own split; P 0.8613 / R 0.9085 shared', 'stock_f1': 'not recorded', 'stock_aji': 'not recorded', 'stock_dice': 'not recorded', 'train_loss': 'not recorded', 'val_loss': 'not recorded', 'best_epoch': '150 / 150'},
"display_name": "Toxoplasma Plaque Well Detector v1",
"architecture": "YOLO11n",
"dataset": "whole-plate and multi-well crystal violet images; 562 "
"images from 1 dataset, 190 of them with no well in them",
"versus_stock": "mAP50 0.993 on its own held-out split; on the test "
"set shared with v2 it scores mAP50 0.8838, against "
"v2's 0.9457",
"trained_on": (
"whole-plate and multi-well Toxoplasma plaque-assay images; "
"yolo11n base, 150 epochs, batch 16, imgsz 640"
),
"trained_by": "einarolafsson",
"notes": (
"the 0.993 is measured on v3's OWN split, which is easier than "
"the set v2 is measured on; on that shared set this model scores "
"mAP50 0.8838 against v2's 0.9457, so v2 is the better detector",
"kept because the published plaque corpus was measured with these "
"weights, so results in the paper trace back to this row",
"locates WELLS, not plaques; it is the front half of a two-stage "
"pipeline with toxoplasma_plaque_v1, and the well it finds also "
"gives the diameter that makes areas comparable across "
"microscopes",
),
},
{
"key": "toxoplasma_well_detector_v2",
"name": "yolo_welldetect_v4.pt",
"kind": "detector",
"repo_id": "einarolafsson/toxoplasma-plaque-well-detector-yolo26",
"repo_type": "model",
"uri": "https://huggingface.co/einarolafsson/"
"toxoplasma-plaque-well-detector-yolo26/resolve/main/"
"weights/best.pt",
"remote_name": "best.pt",
"sha256":
"f2a1e1110f09b2a1d5ef5545adaba7c57f1158669d0bfc50d8fabe9f86da30c7",
"metrics": {'n_train': '1,070 images', 'train_objects': '2,455 boxes', 'n_test': '129 images, 84 of them with no well', 'test_objects': '297 boxes', 'cv': 'no', 'f1': 'mAP50 0.9457 (shared test set)', 'aji': 'mAP50-95 0.8341', 'dice': 'P 0.8912 / R 0.9440', 'stock_f1': 'v1 scores 0.8838 on this set', 'stock_aji': 'v1 scores 0.7630', 'stock_dice': 'v1 P 0.8613 / R 0.9085', 'train_loss': 'not recorded', 'val_loss': 'not recorded', 'best_epoch': '28'},
"display_name": "Toxoplasma Plaque Well Detector v2",
"architecture": "YOLO26n (ultralytics 8.4.155)",
"dataset": "plate images and literature figures; 1,070 train / 254 "
"val / 129 test, split by PMC article so no paper is in two "
"sets; training data at einarolafsson/"
"toxoplasma-plaque-well-detector-dataset",
"versus_stock": "mAP50 0.9457 and mAP50-95 0.8341 against v1's "
"(yolo_welldetect_v3.pt) 0.8838 and 0.7630 on the SAME "
"test set; stock YOLO has no plaque-well class, so v1 "
"is the baseline",
"trained_on": (
"whole-plate and multi-well Toxoplasma plaque-assay images plus 939 "
"newly reviewed PMC figures, accepted boxes and confirmed negatives "
"alike; yolo26n base, best validation mAP50-95 at epoch 28"
),
"trained_by": "einarolafsson",
"notes": (
"on the shared test set it beats v1 (the v3 weights) on every "
"measure, and cuts false boxes on no-well figures from 152 to 49",
"84 of the 129 test images contain no well at all, which is what "
"the false-box count is measured on",
"locates WELLS, not plaques; the front half of a two-stage pipeline "
"with the plaque segmentation model",
"the repository publishes this weight as weights/best.pt; spaCR "
"saves it under the name above so two detectors cannot both land "
"as best.pt",
),
},
{
"key": "toxoplasma_from_cellmask_v1",
"name": "toxoplasma_from_cellmask_pv",
"kind": "cellpose",
"repo_id": "einarolafsson/toxoplasma-from-cellmask-cpsam",
"repo_type": "model",
"uri": None,
"sha256":
"481dfccc1a68cc594aafcb71088efc25b5f5c6a6240e52902c0089759b3149ab",
"metrics": {'n_train': '2567', 'train_objects': 'not recorded', 'n_test': '463', 'test_objects': '6116', 'cv': 'no', 'f1': '0.6058', 'aji': '0.4939', 'dice': '0.6096', 'stock_f1': '0.0215', 'stock_aji': '0.0080', 'stock_dice': '0.0201', 'train_loss': '0.0075', 'val_loss': '0.0100', 'best_epoch': '100 / 100'},
"display_name": "Toxoplasma from Cell Mask (cross-channel)",
"architecture": "Cellpose-SAM (cpsam_v2)",
"dataset": "Toxoplasma PV masks predicted from the HOST CELL MASK channel "
"alone; 2567 training and 463 held-out fields, split by well, "
"hosts HFF/HeLa/THP1",
"versus_stock": "F1 0.606 against 0.021 for stock cpsam_v2 on 463 "
"well-grouped held-out fields, at IoU 0.5",
"trained_on": (
"cross-channel: given the host cell image, predicts where the "
"Toxoplasma parasitophorous vacuoles are, with no parasite stain. "
"100 epochs, base cpsam_v2, AdamW lr 1e-5, targets are "
"PV-regenerated masks"
),
"trained_by": "einarolafsson",
"notes": (
"F1 0.606, AJI 0.494, Dice 0.610 at IoU 0.5 against stock cpsam_v2's "
"0.021/0.008/0.020 -- stock cannot do this task at all",
"per host: HeLa 0.711, HFF 0.557, THP1 0.465; THP1 is the weak case",
"the held-out split selects the checkpoint, so it is validation data "
"rather than an independent test set",
"accuracy falls above IoU 0.8 -- suited to counting, occupancy and "
"area rather than precise morphometry",
),
},
{
"key": "toxoplasma_pv_v2",
"name": "cpsam_v2_toxo_r5",
"kind": "cellpose",
"repo_id": "einarolafsson/toxoplasma-pv-segmentation-cpsam-r5",
"repo_type": "model",
"uri": None,
"sha256":
"17c689e3b117745561e20a885c2a2a998ed360fa97cac8c0446316ae5905c10f",
"metrics": {'n_train': '556', 'train_objects': 'not recorded', 'n_test': '619 pairs', 'test_objects': 'not recorded', 'cv': '5-fold', 'f1': '0.8170 ± 0.036', 'aji': '0.7144 ± 0.107', 'dice': '0.8024 ± 0.118', 'stock_f1': '0.7130', 'stock_aji': '0.4260', 'stock_dice': 'not recorded', 'train_loss': '0.0476', 'val_loss': 'not recorded', 'best_epoch': '100 / 100'},
"display_name": "Toxoplasma PV v2 (round 5)",
"architecture": "Cellpose-SAM (cpsam_v2)",
"dataset": "anti-Toxoplasma-biotin and DsRed PV lumen; 556 curated "
"images accumulated over five rounds",
"versus_stock": "F1 0.817 +/- 0.036 by 5-fold cross-validation over 619 "
"pairs; ~0.86 against 0.713 for stock on the 11 in-house "
"held-out wells",
"trained_on": (
"Toxoplasma tachyzoite parasitophorous vacuoles stained with goat "
"anti-Toxoplasma-biotin, and tachyzoites expressing DsRed in the PV "
"lumen (RH and ME49). Round 5: 556 images, 100 epochs, base cpsam_v2"
),
"trained_by": "einarolafsson",
"notes": (
"supersedes toxoplasma_pv_v1 (round 2, 229 images): more than twice "
"the training data and cross-validated rather than single-split",
"5-fold CV over 619 pairs: F1 0.817 (SD 0.036), AJI 0.714, Dice 0.802",
"per-dataset variance is real -- F1 ranges ~0.74 to ~0.93 by screen",
"accuracy falls above IoU 0.8 -- suited to counting and area rather "
"than precise morphometry",
'Re-scored 2026-10-11 on curated ground truth (spaCR pipeline call, '
'defaults; receipt '
'features/data/model_rescore_curated_gt_2026-10-11/receipt.json): '
'strict hold-out, fields this model never saw, F1 0.713 at IoU 0.5 on n'
' = 258 of 655 fields (n differs between models); common subset unseen '
'by every pv model, n = 163, F1 0.761 (95% bootstrap CI 0.691 to '
'0.806). Without the 98 paged round-3 fields (n = 65): F1 0.742. '
'Earlier numbers in this row are superseded by this re-score where they'
' differ. Default unchanged.',
),
},
{
"key": "toxoplasma_pv_v3",
"name": "cpsam_v2_toxo_r6",
"kind": "cellpose",
"repo_id": "einarolafsson/toxoplasma-pv-segmentation-cpsam-r6",
"repo_type": "model",
"uri": "https://huggingface.co/einarolafsson/"
"toxoplasma-pv-segmentation-cpsam-r6/resolve/main/"
"weights/cpsam_v2_toxo_r6",
"sha256":
"146ef269979b1d1ab45c11039b0ab164f68001adaa8f73f1f8f18be6fcfd060e",
"metrics": {'n_train': '437 fields', 'train_objects': '15,550', 'n_test': '11 anchor wells', 'test_objects': '683', 'cv': '5-fold, grouped by source', 'f1': '0.8602', 'aji': '0.8026', 'dice': '0.9059', 'stock_f1': '0.7648', 'stock_aji': '0.5050', 'stock_dice': '0.6431', 'train_loss': 'not recorded', 'val_loss': '0.0864', 'best_epoch': 'final epoch 100 / 100 (shipped weights; recorded 2026-10-12)'},
"display_name": "Toxoplasma PV v3 (round 6)",
"architecture": "Cellpose-SAM (cpsam_v2)",
"dataset": "the 556 curated PV fields of round 5, split 437 train / "
"108 validation / 11 test; training data at "
"einarolafsson/toxoplasma-pv-segmentation-dataset",
"versus_stock": "F1 0.860 against stock cpsam_v2's 0.765 on the 11 "
"anchor wells at IoU 0.5; AJI 0.803 against 0.505",
"trained_on": (
"Toxoplasma tachyzoite parasitophorous vacuoles stained with goat "
"anti-Toxoplasma-biotin, and tachyzoites expressing DsRed in the PV "
"lumen (RH and ME49). Round 6 retrains round 5's data with cellpose "
"4.2.1.1, 100 epochs, base cpsam_v2"
),
"trained_by": "einarolafsson",
"notes": (
"NEWEST IS NOT BEST HERE: round 6 does not beat round 2 on the "
"anchor wells -- 0.8602 against 0.8648 -- and the PV project still "
"promotes round 5",
"the shipped weights are the FINAL epoch 100, not a validation-chosen "
"checkpoint: the training records show no checkpoint selection "
"(recorded 2026-10-12)",
"5-fold cross-validation, grouped by source: F1 0.8168 +/- 0.028, "
"AJI 0.7516, Dice 0.8424",
"the 11 anchor wells have been held out since round 1, so they are "
"the only fields no PV round has ever trained on",
'Re-scored 2026-10-11 on curated ground truth (spaCR pipeline call, '
'defaults; receipt '
'features/data/model_rescore_curated_gt_2026-10-11/receipt.json): '
'strict hold-out, fields this model never saw, F1 0.750 at IoU 0.5 on n'
' = 269 of 655 fields (n differs between models); common subset unseen '
'by every pv model, n = 163, F1 0.790 (95% bootstrap CI 0.726 to '
'0.835). Without the 98 paged round-3 fields (n = 65): F1 0.754. '
'Earlier numbers in this row are superseded by this re-score where they'
' differ. Default unchanged.',
),
},
{
"key": "toxoplasma_pv_v4",
"name": "cpsam_v2_toxo_r7",
"kind": "cellpose",
"repo_id": "einarolafsson/toxoplasma-pv-segmentation-cpsam-r7",
"repo_type": "model",
"uri": "https://huggingface.co/einarolafsson/"
"toxoplasma-pv-segmentation-cpsam-r7/resolve/main/"
"weights/cpsam_v2_toxo_r7",
"sha256":
"621a475c9bfb6865be7de12c5c4638f4509d734e0c4d39adb15ea541e48961d2",
"metrics": {'n_train': '502 fields',
'train_objects': '23,340',
'n_test': '11 anchor wells; 10 nonempty scored',
'test_objects': '683',
'cv': '5-fold, grouped by source',
'f1': {'finetuned': 0.8536, 'vanilla': 0.7648, 'delta': 0.08879999999999999},
'aji': {'finetuned': 0.7755, 'vanilla': 0.505, 'delta': 0.27049999999999996},
'dice': {'finetuned': 0.8905, 'vanilla': 0.6431, 'delta': 0.24739999999999995},
'stock_f1': '0.7648',
'stock_aji': '0.5050',
'stock_dice': '0.6431',
'train_loss': '0.0424',
'val_loss': '0.0796',
'best_epoch': 'final epoch 100 / 100 (shipped weights; recorded 2026-10-12)',
'precision': {'finetuned': 0.8493, 'vanilla': 0.7539, 'delta': 0.09540000000000004},
'recall': {'finetuned': 0.858, 'vanilla': 0.776, 'delta': 0.08199999999999996},
'pixel_iou': {'finetuned': 0.8055, 'vanilla': 0.5239, 'delta': 0.28159999999999996},
'mAP': {'finetuned': 0.4921, 'vanilla': 0.362, 'delta': 0.1301},
'holdout': 'pv_r7_test',
'holdout_version': 'r7-published-truth',
'n_fields': 10,
'n_objects': 683},
"display_name": "Toxoplasma PV v4 (round 7)",
"architecture": "Cellpose-SAM (cpsam_v2)",
"dataset": "round 6's 556 curated fields plus 80 hand-curated fields of "
"a new plate (Anu revision, Replication09182026 plate 1); "
"502 train / 123 validation / 11 test; training data at "
"einarolafsson/toxoplasma-pv-segmentation-dataset-r7",
"versus_stock": "F1 0.854 against stock cpsam_v2's 0.765 on the 11 "
"anchor wells at IoU 0.5; AJI 0.776 against 0.505",
"trained_on": (
"Toxoplasma tachyzoite parasitophorous vacuoles: round 6's DsRed and "
"anti-Toxoplasma-biotin fields plus the Toxoplasma channel of a new "
"384-well plate. Round 7, cellpose 4.2.1.1, 100 epochs, base cpsam_v2"
),
"trained_by": "einarolafsson",
"notes": (
"Scorecard metrics are means over nonempty fields, not pooled "
"object counts; see aggregation.json.",
"NEWEST IS NOT BEST HERE: on 20 held-out fields of the new plate it "
"was trained on, round 7 scores F1 0.748 against 0.899 for round 6 "
"(recall 0.636 against 0.829) -- prefer toxoplasma_pv_v3 on that plate",
"on the 11 anchor wells it ties round 6 (0.8536 against 0.8602); "
"5-fold cross-validation F1 0.8142 +/- 0.015 against round 6's 0.8168",
"it traces the vacuoles it finds more tightly than round 6 (AJI 0.827 "
"against 0.786 on the new plate) but finds fewer of them",
"the drop on the new plate is under investigation: round 7 also "
"under-fits that plate's own training fields, which points at how "
"those fields entered training rather than at generalisation",
'Re-scored 2026-10-11 on curated ground truth (spaCR pipeline call, '
'defaults; receipt '
'features/data/model_rescore_curated_gt_2026-10-11/receipt.json): '
'strict hold-out, fields this model never saw, F1 0.761 at IoU 0.5 on n'
' = 209 of 655 fields (n differs between models); common subset unseen '
'by every pv model, n = 163, F1 0.781 (95% bootstrap CI 0.725 to '
'0.817). Without the 98 paged round-3 fields (n = 65): F1 0.757. '
'Earlier numbers in this row are superseded by this re-score where they'
' differ. Default unchanged.',
),
},
{
"key": "toxoplasma_pv_v5",
"name": "cpsam_v2_toxo_r8",
"kind": "cellpose",
"repo_id": "einarolafsson/toxoplasma-pv-segmentation-cpsam-r8",
"repo_type": "model",
"uri": "https://huggingface.co/einarolafsson/"
"toxoplasma-pv-segmentation-cpsam-r8/resolve/main/"
"weights/cpsam_v2_toxo_r8",
"sha256":
"a1d7f1a978cf4925ec60db486a00eea33e87e6fc5a6b0d0f34c6db5f132d5c1c",
"size_bytes": 1218639811,
"metrics": {'n_train': '202 fields',
'train_objects': '9,276',
'n_test': '50 fields (whole acquisitions/papers held out)',
'test_objects': '3,530',
'cv': 'no (grouped train/validation/test)',
'f1': {'finetuned': 0.7921, 'vanilla': 0.5586, 'delta': 0.2335},
'precision': {'finetuned': 0.7531, 'vanilla': 0.4635, 'delta': 0.2896},
'recall': {'finetuned': 0.8354, 'vanilla': 0.7028, 'delta': 0.1326},
'aji': 'not recorded',
'dice': 'not recorded',
'stock_f1': '0.5586',
'stock_aji': 'not recorded',
'stock_dice': 'not recorded',
'train_loss': 'see training/epoch_history.json',
'val_loss': 'see training/epoch_history.json',
'best_epoch': '20 / 100',
'holdout': 'retrain_20261009_pv_test',
'holdout_version': '2026-10-09',
'n_fields': 50,
'n_objects': 3530},
"display_name": "Toxoplasma PV v5 (round 8, alternative)",
"architecture": "Cellpose-SAM (cpsam_v2, Cellpose 4.0.9)",
"dataset": "285 curated PV fields (lab fields and published-figure "
"crops), 13,814 vacuoles; 202 train / 33 validation / 50 test "
"grouped by acquisition and paper",
"versus_stock": "F1 0.792 against stock cpsam_v2's 0.559 on "
"the 50 held-out fields at IoU 0.5",
"trained_on": (
"Toxoplasma parasitophorous vacuoles from lab immunofluorescence and "
"DsRed fields and published-figure crops. Round 8, cellpose 4.0.9, "
"best of 100 epochs (epoch 20), base cpsam_v2"
),
"trained_by": "einarolafsson",
"notes": (
"ALTERNATIVE, NOT BETTER: on the 15 test fields round 7 never saw it "
"scores F1 0.536 against toxoplasma_pv_v4's "
"0.525 -- a tie; prefer toxoplasma_pv_v4",
"35 of the 50 test fields were in round 7's training or validation "
"set, so the all-field comparison (0.792 against 0.851) favours "
"toxoplasma_pv_v4 and is not a fair test",
"weak on published-figure crops (F1 0.440 on 14 literature fields), "
"as is every PV model",
'Re-scored 2026-10-11 on curated ground truth (spaCR pipeline call, '
'defaults; receipt '
'features/data/model_rescore_curated_gt_2026-10-11/receipt.json): '
'strict hold-out, fields this model never saw, F1 0.733 at IoU 0.5 on n'
' = 481 of 655 fields (n differs between models); common subset unseen '
'by every pv model, n = 163, F1 0.633 (95% bootstrap CI 0.544 to '
'0.695). Without the 98 paged round-3 fields (n = 65): F1 0.665. '
'Earlier numbers in this row are superseded by this re-score where they'
' differ. LEAKED: this model trained or validated on 29 of 39 curated '
'PV literature hold-out fields and 3 of 11 anchor wells, so its earlier'
' hold-out scores are not hold-out scores; the re-score uses only the '
'fields it never saw. Default unchanged.',
),
},
{
"key": "live_cell_v1",
"name": "live_cell_v1",
"kind": "cellpose",
"repo_id": "einarolafsson/live-cell-segmentation-cpsam",
"repo_type": "model",
"uri": "https://huggingface.co/einarolafsson/"
"live-cell-segmentation-cpsam/resolve/main/"
"weights/live_cell_v1",
"sha256":
"7ade69377093fe81830ddc7c52ba8618bef1fefe7d1243c01b9c1beed7fcb090",
"metrics": {'n_train': '6,778 fields, 14 datasets', 'train_objects': 'not recorded', 'n_test': '2,199 fields', 'test_objects': 'not recorded', 'cv': 'no (acquisition-grouped train/valid/test)', 'f1': '0.694 all / 0.960 on datasets stock never saw', 'aji': 'not recorded', 'dice': 'not recorded', 'stock_f1': '0.738 all / 0.885 on datasets stock never saw', 'stock_aji': 'not recorded', 'stock_dice': 'not recorded', 'train_loss': 'not recorded', 'val_loss': 'not recorded', 'best_epoch': '37 (stopped by the maintainer)'},
"display_name": "Live cell v1 (phase, brightfield, DIC; alternative)",
"architecture": "Cellpose-SAM (cpsam_v2)",
"dataset": "11,007 transmitted-light fields from 14 public datasets, "
"split by acquisition 6,778 train / 2,030 validation / "
"2,199 test; training data at "
"einarolafsson/live-cell-segmentation-dataset",
"versus_stock": "on the datasets stock cpsam_v2 never trained on, F1 "
"0.960 against 0.885 at IoU 0.5; over all 2,199 test "
"fields, 0.694 against 0.738, because stock trained "
"on LIVECell and YeaZ and wins on LIVECell",
"trained_on": (
"unstained cells in phase contrast, brightfield and DIC: "
"LIVECell, DeepSea, YeaZ, yeast microstructures, five Cell "
"Tracking Challenge sets, BBBC009, BBBC030, QPI and Revvity. "
"Base cpsam_v2, cellpose 4.2.1.1, lr 1e-5, batch 4; stopped at "
"epoch 37 of 100"
),
"trained_by": "einarolafsson",
"notes": (
"an optional alternative: stock Cellpose-SAM (cpsam_v2) stays the "
"default for label-free and live-cell segmentation",
"TWO STOCK COMPARISONS, NOT ONE: stock cpsam_v2 trained on LIVECell "
"and YeaZ, so its score there is partly memorisation. On the "
"datasets it never saw (DeepSea, the CTC sets, BBBC009, BBBC030, "
"QPI, Revvity, yeast microstructures) this model scores F1 0.960 "
"against 0.885",
"it does NOT replace stock on LIVECell-style Incucyte phase: "
"0.671 against 0.724 there, and it missed its own pre-registered "
"promotion bar",
"by modality at IoU 0.5: brightfield 0.964 (stock 0.912), DIC "
"+0.026 over stock, phase 0.689 (stock 0.735)",
"F1 0.865 on a train sample, 0.696 on validation and 0.694 on "
"test; no per-epoch loss was recorded",
),
},
{
"key": "nuclei_from_cellmask_v1",
"name": "nuclei_from_cellmask_best",
"kind": "cellpose",
"repo_id": "einarolafsson/cross-channel-nuclei-from-cellmask-cpsam",
"repo_type": "model",
"uri": "https://huggingface.co/einarolafsson/"
"cross-channel-nuclei-from-cellmask-cpsam/resolve/main/"
"weights/nuclei_from_cellmask_best",
"sha256":
"2675553a46e97a7bc4bd2bfe3e954954194fe71ca4e94261e752a02bf0b6eb47",
"metrics": {'n_train': 'not recorded', 'train_objects': 'not recorded', 'n_test': '453', 'test_objects': 'not recorded', 'cv': 'no', 'f1': '0.8881', 'aji': '0.7916', 'dice': '0.8774', 'stock_f1': '0.201', 'stock_aji': '0.286', 'stock_dice': '0.449', 'train_loss': 'not recorded', 'val_loss': 'not recorded', 'best_epoch': 'not recorded'},
"display_name": "Cross-channel nuclei-from-cellmask",
"architecture": "Cellpose-SAM (cpsam_v2)",
"dataset": "nuclei predicted from the HOST CELL MASK channel alone; "
"453 well-grouped held-out fields, hosts HFF/HeLa/THP1",
"versus_stock": "F1 0.888 against 0.201 for stock cpsam_v2 on "
"453 well-grouped held-out fields, at IoU 0.5",
"trained_on": (
"cross-channel: given the cell image, predicts where the nuclei "
"are, with no nuclear stain -- which frees the DAPI/Hoechst "
"channel for another marker. 100 epochs, base cpsam_v2"
),
"trained_by": "einarolafsson",
"notes": (
"F1 0.888, AJI 0.792, Dice 0.877 at IoU 0.5 against stock "
"cpsam_v2's 0.201/0.286/0.449",
"per host: HFF 0.932, HeLa 0.860, THP1 0.861",
"the held-out split selects the checkpoint, so it is validation "
"data rather than an independent test set",
"predicts nuclei from cell morphology -- expect degraded accuracy "
"on unusual or highly confluent morphologies",
),
},
{
"key": "cell_from_hoechst_v1",
"name": "cell_from_hoechst_best",
"kind": "cellpose",
"repo_id": "einarolafsson/cross-channel-cell-from-hoechst-cpsam",
"repo_type": "model",
"uri": "https://huggingface.co/einarolafsson/"
"cross-channel-cell-from-hoechst-cpsam/resolve/main/"
"weights/cell_from_hoechst_best",
"sha256":
"d1992433b4f2f291f73738830bb953198165fdc10719e54dae9b8c2bd430e0eb",
"size_bytes": 1218647799,
"metrics": {'n_train': '2,578 fields', 'train_objects': '237,957', 'n_test': '451 fields', 'test_objects': '45,098', 'cv': 'no, split by well', 'f1': '0.8697', 'aji': '0.7991', 'dice': '0.8948', 'stock_f1': '0.3012', 'stock_aji': '0.3506', 'stock_dice': '0.5235', 'train_loss': 'not recorded', 'val_loss': 'not recorded', 'best_epoch': '70 / 100'},
"display_name": "Cross-channel cell-from-hoechst",
"architecture": "Cellpose-SAM (cpsam_v2)",
"dataset": "the HOST CELL outline predicted from the Hoechst (nuclear) "
"channel alone; 2,578 training fields and 451 held-out test "
"fields, split by well so no well is on both sides",
"versus_stock": "F1 0.870 against stock cpsam_v2's 0.301 on 451 "
"held-out fields at IoU 0.5 -- a delta of 0.569",
"trained_on": (
"Hoechst-stained nuclei paired with curated host-cell masks; "
"fine-tuned from stock cpsam_v2, 100 epochs, best epoch 70"
),
"trained_by": "einarolafsson",
"notes": (
"the counterpart of nuclei_from_cellmask_v1: that one predicts "
"nuclei from the cell mask, this one predicts the cell from the "
"nucleus",
"precision 0.944 against recall 0.806 -- it misses cells rather "
"than inventing them, which is the safer direction for counting",
"quote the DELTA over stock (0.569), not the ratio: stock's mAP of "
"0.0575 is a near-zero denominator that makes any ratio look huge",
),
},
{
"key": "toxoplasma_from_hoechst_v1",
"name": "toxoplasma_from_hoechst_pv",
"kind": "cellpose",
"repo_id": "einarolafsson/toxoplasma-from-hoechst-cpsam",
"repo_type": "model",
"uri": None,
"sha256":
"8dc05ebced3550d1a418c13d24d319e0482c742988df29a525520026cb2f0d96",
"metrics": {'n_train': '2567', 'train_objects': 'not recorded', 'n_test': '463', 'test_objects': 'not recorded', 'cv': 'no', 'f1': '0.569', 'aji': '0.421', 'dice': '0.546', 'stock_f1': '0.002', 'stock_aji': '0.006', 'stock_dice': '0.016', 'train_loss': 'not recorded', 'val_loss': 'not recorded', 'best_epoch': 'not recorded'},
"display_name": "Toxoplasma from Hoechst (cross-channel)",
"architecture": "Cellpose-SAM (cpsam_v2)",
"dataset": "Toxoplasma PV masks predicted from the HOECHST channel alone; "
"2567 training and 463 held-out fields, split by well, "
"hosts HFF/HeLa/THP1",
"versus_stock": "F1 0.569 against 0.002 for stock cpsam_v2 on "
"463 well-grouped held-out fields, at IoU 0.5",
"trained_on": (
"cross-channel: given the Hoechst/nuclear image, predicts where the "
"Toxoplasma parasitophorous vacuoles are, with no parasite stain. "
"100 epochs, base cpsam_v2, AdamW lr 1e-05, targets are "
"PV-regenerated masks"
),
"trained_by": "einarolafsson",
"notes": (
"F1 0.569, AJI 0.421, Dice 0.546 at IoU 0.5 against "
"stock cpsam_v2's 0.002/0.006/0.016",
"the Hoechst route is harder than the cell-mask route -- compare "
"toxoplasma_from_cellmask_v1",
"the held-out split selects the checkpoint, so it is validation data "
"rather than an independent test set",
"accuracy falls above IoU 0.8 -- suited to counting, occupancy and "
"area rather than precise morphometry",
),
},
{
"key": "colony_yolo11n_makrai_v1",
"name": "colony_yolo11n_makrai_v1.pt",
"kind": "detector",
"repo_id": HF_MODELS_REPO,
"repo_type": "dataset",
"uri": "https://huggingface.co/datasets/einarolafsson/models/"
"resolve/main/colony_detector/v1/"
"colony_yolo11n_makrai_v1.pt?download=true",
"sha256":
"1bd99599c1d78fb0841216b798cfc61665be7c2aa8b356b970d0d06ae4b57557",
"size_bytes": 5641114,
"metrics": {'n_train': '369 photos (train/val/test split)', 'train_objects': '56,865 colony boxes in all', 'n_test': '39 held-out plates with 30-300 colonies', 'test_objects': 'not recorded', 'cv': 'no (one split)', 'f1': 'median count error 2.4%', 'aji': '69% of plates within 5%', 'dice': '87% of plates within 10%', 'stock_f1': 'thresholding: median count error 16.5%', 'stock_aji': 'not recorded', 'stock_dice': 'not recorded', 'train_loss': 'not recorded', 'val_loss': 'not recorded', 'best_epoch': 'not recorded'},
"display_name": "Colony Detector v1 (alpha)",
"architecture": "YOLO11n (ultralytics), one class 'colony'",
"dataset": "smartphone plate photos of 24 bacterial and fungal "
"species; 369 photos, 56,865 colony boxes, from Makrai "
"et al. 2023 (Scientific Data, CC BY 4.0)",
"versus_stock": "median count error 2.4% on 39 held-out plates with "
"30-300 colonies, against 16.5% for spaCR's "
"thresholding counter on the same plates",
"trained_on": (
"plate photos from Makrai et al. 2023, Scientific Data, figshare "
"doi:10.6084/m9.figshare.22022540 (CC BY 4.0): 369 smartphone "
"photos of 24 species with 56,865 colony boxes"
),
"trained_by": "einarolafsson",
"licence": "CC BY 4.0",
"notes": (
"counts colonies for colony_counting when colony_detector names "
"this key; ultralytics is needed",
"69% of held-out plates within 5% of the hand count and 87% "
"within 10%",
"one dataset and one split: photos from another lab, agar or "
"species may score worse, and confluent swarming colonies stay "
"hard",
"published in the einarolafsson/models dataset repo under "
"colony_detector/v1/",
),
},
{'key': 'toxoplasma_plaque_v3',
'name': 'cpsam_plaque_r5_geldoc',
'remote_name': 'cpsam_plaque_r5_geldoc',
'kind': 'cellpose',
'repo_id': 'einarolafsson/toxoplasma-plaque-segmentation-cpsam-r5-geldoc',
'repo_type': 'model',
'uri': 'https://huggingface.co/einarolafsson/toxoplasma-plaque-segmentation-cpsam-r5-geldoc/resolve/main/weights/cpsam_plaque_r5_geldoc',
'sha256': '8f59850dbe09c3a252728b0a7ca38ae5b52a28f2dbaa7a068174b4b7bfac768c',
'size_bytes': 1218645663,
'metrics': {'f1': {'finetuned': 0.6418703506907545, 'vanilla': None, 'delta': None},
'precision': {'finetuned': 0.6942528735632184, 'vanilla': None, 'delta': None},
'recall': {'finetuned': 0.5968379446640316, 'vanilla': None, 'delta': None},
'ap': {'finetuned': 0.4726134585289515, 'vanilla': None, 'delta': None},
'true_positives': {'finetuned': 302, 'vanilla': None, 'delta': None},
'false_positives': {'finetuned': 133, 'vanilla': None, 'delta': None},
'false_negatives': {'finetuned': 204, 'vanilla': None, 'delta': None},
'holdout': 'geldoc_heldout_min_size_0',
'holdout_version': '20261007T205603',
'n_fields': 18,
'n_objects': 506,
'n_train': '496 fields; 34 reviewed empty',
'train_objects': '25794',
'n_test': '18 Gel Doc wells, 3 plate groups',
'test_objects': '506',
'cv': 'no',
'aji': 'not recorded',
'dice': 'not recorded',
'stock_f1': 'not measured',
'stock_aji': 'not measured',
'stock_dice': 'not measured',
'train_loss': '0.18370248939254452',
'val_loss': 'see monitored NAS history; final epoch not measured',
'best_epoch': 'final checkpoint, 100 epochs; no best-epoch selection'},
'display_name': 'Toxoplasma Plaque v3 (Gel Doc r5, alternative)',
'architecture': 'Cellpose-SAM (cpsam, Cellpose 4.0.9)',
'dataset': '496 curated plaque fields, including 34 reviewed empty negatives and 71 Gel Doc '
'wells; 100 epochs; fixed physical-plate/source groups',
'versus_stock': 'Stock was not evaluated in this run; see the named incumbent comparison on the '
'model card.',
'trained_on': '496 curated plaque fields, including 34 reviewed empty negatives and 71 Gel Doc '
'wells; 100 epochs; fixed physical-plate/source groups',
'trained_by': 'einarolafsson',
'notes': ['Recommended alternative for Gel Doc wells: on the 153 test fields of the curated 2026-10-15 set it scores '
'object F1 0.954 on Gel Doc wells against 0.868 for the default toxoplasma_plaque_v5 (0.741 against 0.816 overall).',
'Alternative since 2026-10-11, when toxoplasma_plaque_v4 (base cpsam_v2) became the default by the maintainer\'s choice. '
'This model scores higher: on 263 held-out fields F1 0.844 against v4 0.815, and 0.913 against 0.804 on Gel Doc wells, '
'through spacr.plaque.segment_plaque_image. Kept so recorded runs reproduce.',
'Default 2026-10-10 to 2026-10-11: F1 0.844 against toxoplasma_plaque_v2 0.799, and 0.913 against 0.013 on Gel Doc wells '
'(features/data/model_compare_2026-10-10).',
'Its own release gate on independent literature crops did not pass at the time (literature F1 about 0.54 for every model).',
'Gel Doc score requires min_size=0; min_size=150 gives F1 0.1041.',
'Stock not measured; comparison baseline is incumbent r3.',
'Distinct from the cpsam_v2 round-5 model already registered as toxoplasma_plaque_v2.',
'Re-scored 2026-10-11 on curated ground truth (spaCR pipeline call, '
'defaults; receipt '
'features/data/model_rescore_curated_gt_2026-10-11/receipt.json): '
'strict hold-out, fields this model never saw, F1 0.732 at IoU 0.5 on n'
' = 232 of 562 fields (n differs between models); common subset unseen '
'by every plaque model, n = 114, F1 0.548 (95% bootstrap CI 0.449 to '
'0.645). Without the v3-seeded Gel Doc truth (n = 81): F1 0.444. '
'Earlier numbers in this row are superseded by this re-score where they'
' differ. 263-field set seen counts: v1 54, v2 178, v3 137, v4 0; GRA8 '
'Plate1 (hold-out of v3) is the same physical plate as ReimagedPlate1 '
'(training data of v4). Default unchanged.']},
{'key': 'toxoplasma_plaque_v4',
'name': 'cpsam_plaque_r6',
'remote_name': 'cpsam_plaque_r6',
'kind': 'cellpose',
'repo_id': 'einarolafsson/toxoplasma-plaque-segmentation-cpsam-r6',
'repo_type': 'model',
'uri': 'https://huggingface.co/einarolafsson/toxoplasma-plaque-segmentation-cpsam-r6/resolve/main/weights/cpsam_plaque_r6',
'sha256': '744aee3f71efa68f6a369a6f91fe5850dc61e35f3dc20d53dda2c04fce17b35f',
'size_bytes': 1218641683,
'metrics': {'f1': {'finetuned': 0.8154, 'vanilla': None, 'delta': None},
'precision': {'finetuned': 0.8252, 'vanilla': None, 'delta': None},
'recall': {'finetuned': 0.8058, 'vanilla': None, 'delta': None},
'holdout': 'retrain_20261009_plaque_test',
'holdout_version': '2026-10-09',
'n_fields': 263,
'n_objects': 12792,
'n_train': '335 fields; 5 reviewed skips excluded',
'train_objects': '16467',
'n_test': '263 fields: 3 held-out Gel Doc plates plus the r5 evaluation split',
'test_objects': '12792',
'cv': 'no',
'aji': 'not recorded',
'dice': 'not recorded',
'stock_f1': 'not measured with spacr.plaque.segment_plaque_image',
'stock_aji': 'not measured',
'stock_dice': 'not measured',
'train_loss': 'see training/epoch_history.json',
'val_loss': 'see training/epoch_history.json',
'best_epoch': '140 / 190'},
'display_name': 'Toxoplasma Plaque v4 (round 6, cpsam_v2)',
'architecture': 'Cellpose-SAM (cpsam_v2, Cellpose 4.0.9)',
'dataset': '663 curated plaque fields from four domains (published-figure wells, Gel Doc '
'wells, lab plates, staged literature crops), 32,717 plaques; 335 train / 65 '
'validation / 263 test',
'versus_stock': 'F1 0.815 (toxoplasma_plaque_v3 0.844) on the 263 held-out fields at IoU '
'0.5; stock cpsam_v2 0.373 there with the Cellpose evaluation scorer',
'trained_on': 'Toxoplasma plaques from published-figure wells, Bio-Rad Gel Doc wells, lab plates '
'and staged literature crops; round 6 splits, best of 190 epochs (epoch 140), '
'base cpsam_v2',
'trained_by': 'einarolafsson',
'notes': ['No longer the default: toxoplasma_plaque_v5 replaced it on 2026-10-17 (object F1 '
'0.816 against 0.725 on the 153 curated test fields); kept selectable so recorded runs reproduce.',
'The default plaque model 2026-10-11 to 2026-10-17, by the maintainer\'s rule that spaCR '
'models build on cpsam_v2. It scores below toxoplasma_plaque_v3 on the same 263 '
'held-out fields, through spacr.plaque.segment_plaque_image: F1 0.815 against 0.844, '
'lower on every source (Gel Doc 0.804 against 0.913); v3 stays selectable.',
'Fine-tuned from cpsam_v2 (2026-10-10). It replaces the first round-6 weights, '
'fine-tuned from first-generation cpsam (sha256 00f4a86b...a0c171), which remain in '
'the repository history.',
'Above toxoplasma_plaque_v2 overall (0.799) and on Gel Doc wells, where v2 scores 0.013.',
'Published-figure literature crops stay hard (F1 0.530).',
'Re-scored 2026-10-11 on curated ground truth (spaCR pipeline call, '
'defaults; receipt '
'features/data/model_rescore_curated_gt_2026-10-11/receipt.json): '
'strict hold-out, fields this model never saw, F1 0.765 at IoU 0.5 on n'
' = 329 of 562 fields (n differs between models); common subset unseen '
'by every plaque model, n = 114, F1 0.534 (95% bootstrap CI 0.440 to '
'0.625). Without the v3-seeded Gel Doc truth (n = 81): F1 0.459. '
'Earlier numbers in this row are superseded by this re-score where they'
' differ. 263-field set seen counts: v1 54, v2 178, v3 137, v4 0; GRA8 '
'Plate1 (hold-out of v3) is the same physical plate as ReimagedPlate1 '
'(training data of v4). Default unchanged.']},
{'key': 'toxoplasma_plaque_v5',
'name': 'cpsam_plaque_r7',
'remote_name': 'cpsam_plaque_r7',
'kind': 'cellpose',
'repo_id': 'einarolafsson/toxoplasma-plaque-segmentation-cpsam-r7',
'repo_type': 'model',
'uri': 'https://huggingface.co/einarolafsson/toxoplasma-plaque-segmentation-cpsam-r7/resolve/main/weights/cpsam_plaque_r7',
'sha256': '9243917beefbc3d1c8d997d56c1980b811ffef9ee075dac60b099aec715fa8a1',
'size_bytes': 1218648776,
'metrics': {'f1': {'finetuned': 0.8157, 'vanilla': 0.2268, 'delta': 0.5889},
'holdout': 'plaque_curated_20261015_test',
'holdout_version': '2026-10-15',
'n_fields': 153,
'n_objects': 6540,
'n_train': '1201 fields + 153 validation',
'train_objects': '49360',
'n_test': '153 fields',
'test_objects': '6540',
'cv': 'no',
'train_loss': 'see training/run_config.json',
'val_loss': 'see training/run_config.json',
'best_epoch': '100 / 100 (final epoch kept)'},
'display_name': 'Toxoplasma Plaque v5 (round 7, cpsam_v2, default)',
'architecture': 'Cellpose-SAM (cpsam_v2, Cellpose 4.2.1.1)',
'dataset': '1,507 curated plaque fields of 2026-10-15 (Gel Doc wells, lab plates, published '
'figures, earlier curation rounds); 1,201 train / 153 validation / 153 test',
'versus_stock': 'F1 0.816 (stock cpsam_v2 0.227) on the 153 held-out test fields at IoU 0.5; '
'toxoplasma_plaque_v4 0.725 and toxoplasma_plaque_v3 0.741 on the same fields',
'trained_on': 'Toxoplasma plaques from Gel Doc wells, lab plates, published figures and earlier '
'curation rounds (curated dataset of 2026-10-15); 100 epochs, final epoch kept, '
'base cpsam_v2',
'trained_by': 'einarolafsson',
'notes': ['The default plaque model since 2026-10-17, by the maintainer\'s choice; '
'toxoplasma_plaque_v3 is the recommended alternative for Gel Doc wells and '
'toxoplasma_plaque_v4 stays selectable.',
'Scored 2026-10-17 on the 153 test fields of the curated 2026-10-15 dataset '
'(spacr.plaque.segment_plaque_image, defaults): object F1 0.816 (95% bootstrap CI '
'0.778 to 0.848), v4 0.725, v3 0.741, stock cpsam_v2 0.227. On the 81 test fields '
'neither v3 nor v4 saw: v5 0.803, v4 0.661, v3 0.684. AJI is level with v4 '
'(0.737 against 0.744); Gel Doc wells 0.868 against v4 0.901 and v3 0.954.',
'The final epoch 100 ships by the maintainer\'s rule: an earlier checkpoint replaces '
'it only if better on every metric in every domain. Epoch 40 led on validation '
'object F1 (0.766 against 0.750) but trailed on boundary F1 (0.700 against 0.711), '
'AJI (0.661 against 0.667) and Dice (0.857 against 0.869); on test it scored F1 0.823.',
'Training ran 100 epochs; indices 91 to 99 were resumed from epoch 90 after a power '
'cut, with the same schedule and seeds (optimizer moments restarted at LR 2e-8).']},
{'key': 'toxoplasma_well_detector_v3',
'name': 'yolo_welldetect_v4_geldoc.pt',
'remote_name': 'yolo_welldetect_v4_geldoc.pt',
'kind': 'detector',
'repo_id': 'einarolafsson/toxoplasma-plaque-well-detector-yolo11-v4-geldoc',
'repo_type': 'model',
'uri': 'https://huggingface.co/einarolafsson/toxoplasma-plaque-well-detector-yolo11-v4-geldoc/resolve/main/weights/yolo_welldetect_v4_geldoc.pt',
'sha256': '046861021594256957292cfe9488a84e758077a12dac31fe66b8b4083f8dfa18',
'size_bytes': 5478042,
'metrics': {'mAP50': {'finetuned': 0.9940094439993655, 'vanilla': None, 'delta': None},
'mAP50_95': {'finetuned': 0.884212663807534, 'vanilla': None, 'delta': None},
'precision': {'finetuned': 0.96981851739557, 'vanilla': None, 'delta': None},
'recall': {'finetuned': 0.9979244614467876, 'vanilla': None, 'delta': None},
'holdout': 'detector_v3_validation_figure_clean',
'holdout_version': '20261007T205603',
'n_fields': 73,
'n_objects': 191,
'n_train': '452 images',
'train_objects': '1679',
'n_test': '73 clean incumbent-validation images; 3 Gel Doc validation plates '
'separately',
'test_objects': 'see training/fields.csv',
'cv': 'no',
'f1': 'not an object-F1 score; mAP50 0.9940',
'aji': 'not applicable to box detection',
'dice': 'not applicable to box detection',
'stock_f1': 'no stock YOLO plaque_well class; incumbent v3 baseline',
'stock_aji': 'not applicable',
'stock_dice': 'not applicable',
'train_loss': 'box/class/DFL losses in epoch_history.csv',
'val_loss': 'box/class/DFL losses in epoch_history.csv',
'best_epoch': 'validation-selected best.pt; see training checkpoint metadata'},
'display_name': 'Toxoplasma Well Detector v3 (YOLO11 Gel Doc, default)',
'architecture': 'YOLO11n (fine-tuned from detector v3)',
'dataset': '452 reviewed training images; 124 validation images; physical plate and figure '
'groups; 150 epochs; YOLO11n v3 initialization',
'versus_stock': 'Stock was not evaluated in this run; see the named incumbent comparison on the '
'model card.',
'trained_on': '452 reviewed training images; 124 validation images; physical plate and figure '
'groups; 150 epochs; YOLO11n v3 initialization',
'trained_by': 'einarolafsson',
'notes': ['The default well detector since 2026-10-12, by the maintainer\'s choice. On the 73 clean '
'incumbent-validation images its mAP50-95 is lower than the incumbent\'s; on the 3 Gel Doc '
'validation plates it is higher.',
'Gel Doc plates selected best.pt, so their scores are validation rather than '
'independent test.',
'Distinct from the YOLO26 checkpoint also historically called detector v4.',
'Locates wells, not plaques; one class plaque_well (0).']},
{
"key": "mitonet_mini_v1",
"name": "MitoNet_v1_mini.pth",
"kind": "mitonet",
"third_party": True,
"uri": "https://zenodo.org/records/6861565/files/"
"MitoNet_v1_mini.pth?download=1",
"sha256":
"6554a42fbdcb70a46700f20281b1a5f1ae2f4140a63f519193046c7c9a9f2f4c",
"size_bytes": 110074368,
"metrics": {'n_train': 'CEM-MitoLab (EMPIAR-11037), about 22,000 '
'2-D EM images, per the authors',
'train_objects': 'not recorded here',
'n_test': '10 TEM fields of spacr-example-tem '
'(EMPIAR-10982 TEM benchmark)',
'test_objects': '195 mitochondria',
'cv': 'no; the benchmark is its held-out test set',
'f1': '0.823 at IoU 0.5 (TP 142, FP 8, FN 53)',
'aji': 'not recorded',
'dice': '0.842 (foreground pixels)',
'stock_f1': 'cpsam 0.467 on the same fields',
'stock_aji': 'not recorded',
'stock_dice': 'cpsam 0.50',
'train_loss': 'not recorded',
'val_loss': 'not recorded',
'best_epoch': 'not recorded'},
"display_name": "MitoNet v1 mini (TEM mitochondria, third-party)",
"architecture": "Panoptic-DeepLab style network (empanada), "
"TorchScript; spaCR's own post-processing",
"dataset": "CEM-MitoLab, 2-D electron micrographs with every "
"mitochondrion outlined (Conrad and Narayan 2023)",
"versus_stock": "F1 0.823 against cpsam's 0.467 on the ten "
"spacr-example-tem fields at IoU 0.5",
"trained_on": (
"CEM-MitoLab: 2-D TEM and volume-EM images of mitochondria from "
"many tissues and labs, curated by the empanada authors "
"(Conrad and Narayan, Cell Systems 14:58-71, 2023)"
),
"trained_by": "Ryan Conrad and Kedar Narayan (empanada, NCI)",
"licence": "CC-BY-4.0",
"code_licence": "BSD-3-Clause",
"citation": (
"Conrad R, Narayan K. Instance segmentation of mitochondria in "
"electron microscopy images with a generalist deep learning "
"model trained on a diverse dataset. Cell Systems "
"14(1):58-71.e5, 2023. doi:10.1016/j.cels.2022.12.006"
),
"notes": (
"third-party weights from Zenodo record 6861565, licensed "
"CC-BY-4.0: cite Conrad and Narayan, Cell Systems 2023, "
"doi:10.1016/j.cels.2022.12.006",
"the empanada code is BSD-3-Clause; spaCR runs the TorchScript "
"weights with its own port of the post-processing and does not "
"need empanada-dl",
"Make Masks offers it as the MitoNet (TEM mitochondria) "
"detection method; the weights download on first use",
"trained by the lab that made the benchmark it is scored on, "
"though that benchmark is its held-out test set",
),
},
)
#: Models that are no longer OFFERED, by filename.
#:
#: ``toxo_plaque_cyto`` is retired: it
#: recalls 0.631 on the literature set -- it misses about a third of the
#: plaques -- against 0.811 for ``toxoplasma_plaque_v1``, and it published no
#: checksum, so its row in the picker had a Download button that could never
#: succeed.
#:
#: FILTERED FROM THE LISTING, NOT DELETED FROM DISK. The checkpoint still
#: ships, and ``plaque_model='bundled'`` still resolves to it, because a run
#: recorded against it has to stay reproducible: removing the weights would
#: silently change what re-running an old analysis produces, which is worse
#: than offering a model nobody should pick. It is simply no longer something
#: the zoo suggests.
RETIRED_MODEL_NAMES: frozenset = frozenset({
"toxo_plaque_cyto_e25000_X1120_Y1120.CP_model",
})
#: Keys :func:`rank` will sort on, with the direction and what the number is.
#:
#: There is deliberately no accuracy key. A benchmark here has no ground truth
#: (see the module docstring), so a column called "score" that sorts models
#: would be inventing one.
RANK_KEYS: Dict[str, str] = {
"qc": "fraction of fields spacr.seg_qc scored 'ok' — a quality-control "
"verdict on this model's own masks, not an accuracy (higher first)",
"seconds": "wall-clock segmentation time over the field set (lower first)",
}
#: :func:`rank`'s default key.
DEFAULT_RANK_KEY = "qc"
#: First bytes of a torch checkpoint. Everything torch has saved since 1.6 is a
#: zip; ``\\x80`` opens the legacy pickle protocol.
_TORCH_MAGICS = (b"PK\x03\x04", b"\x80")
#: ``name_v3`` -> ``('name', 3)``. See :func:`versioned_path`.
_VERSION_RE = re.compile(r"^(?P<base>.+)_v(?P<n>\d+)$")
_CURATED_SCORECARD_BINDINGS = {'toxoplasma_plaque_v1': {'checkpoint_sha256': 'eeecd2d6cd5cbb4dddee71564d5f460d26bb07ac125e0b494b7502fea4292d5d',
'field_ids_sha256': 'c0759011907c85b3a5fc0d0c1ebfa31ac4c8fd196a29427d8ac80c7e00f47ab5',
'full_card_sha256': '5467b6e446000c87cc61b34802aa4a212c8e8238f283f5fe13fa68afbf65df11',
'resource': 'toxoplasma-plaque-segmentation-cpsam.json',
'resource_sha256': '31d63e888366c0bf70048977efc191bba51bc38e06d0c2f846a9dd36c3335888',
'uri': 'https://huggingface.co/einarolafsson/toxoplasma-plaque-segmentation-cpsam/resolve/main/cpsam_plaque_r3?download=true'},
'toxoplasma_plaque_v2': {'checkpoint_sha256': '0927023a745ac6a19bae0ec72c89b7b864a4ff8d047a41f3f1e9767e1a4d0600',
'field_ids_sha256': 'cb872aac8934e014943e40e42c420f1335f1e815415f553210360c6b80b07fc3',
'full_card_sha256': '92e2362a345b8e27cf1595ac11e532126fe92d9dedd321445e6f06b05799684c',
'resource': 'toxoplasma-plaque-segmentation-cpsam-r5.json',
'resource_sha256': 'a541533956c3c570dfb85155839eb568951d76dfe87b63bf3727c487281a2d21',
'uri': 'https://huggingface.co/einarolafsson/toxoplasma-plaque-segmentation-cpsam-r5/resolve/main/weights/cpsam_plaque_r5?download=true'},
'toxoplasma_plaque_v3': {'checkpoint_sha256': '8f59850dbe09c3a252728b0a7ca38ae5b52a28f2dbaa7a068174b4b7bfac768c',
'field_ids_sha256': 'de6aa3f20249d5f8d7e3fd8cafae08f74f2a823983b36692b73de2fdd7c05424',
'full_card_sha256': '19af073e0a5c752df94dc4943ed66f2927d2351fffa745d070c3da7faf7e4806',
'resource': 'toxoplasma-plaque-segmentation-cpsam-r5-geldoc.json',
'resource_sha256': 'd9ecd22d0582c4b9668c802964e3e236254c6674afe95616375263c5ce2ef09a',
'uri': 'https://huggingface.co/einarolafsson/toxoplasma-plaque-segmentation-cpsam-r5-geldoc/resolve/main/weights/cpsam_plaque_r5_geldoc'},
'toxoplasma_plaque_v4': {'checkpoint_sha256': '744aee3f71efa68f6a369a6f91fe5850dc61e35f3dc20d53dda2c04fce17b35f',
'field_ids_sha256': '5e3dfefe23a2bc51877dc3cb72999b21b1d0b3e4fb393e5d9d2b2bcce686f86f',
'full_card_sha256': '6860674346b819da2045a8a887d39dfc6d343b0250cf748b74d86b4414e2c94a',
'resource': 'toxoplasma-plaque-segmentation-cpsam-r6.json',
'resource_sha256': '7dbefae4de57bf749fd9efe164b375288fcd65a0ca5c3198905d30600a03691f',
'uri': 'https://huggingface.co/einarolafsson/toxoplasma-plaque-segmentation-cpsam-r6/resolve/main/weights/cpsam_plaque_r6'},
'toxoplasma_plaque_v5': {'checkpoint_sha256': '9243917beefbc3d1c8d997d56c1980b811ffef9ee075dac60b099aec715fa8a1',
'field_ids_sha256': '981166ee8b999dad4b15deb0ea6a31c4dd5cd75713ae015650afc598d9cabc70',
'full_card_sha256': '33d7064088e3890a1a702b45a36c2d5aceb79667ab34318e39352d8ec4cc207d',
'resource': 'toxoplasma-plaque-segmentation-cpsam-r7.json',
'resource_sha256': 'e8943595cd5ba25b4a5d6aa31410d48cd028e4ed63acf3744a88a051e740b2d7',
'uri': 'https://huggingface.co/einarolafsson/toxoplasma-plaque-segmentation-cpsam-r7/resolve/main/weights/cpsam_plaque_r7'},
'toxoplasma_pv_v1': {'checkpoint_sha256': '182d8cf6b32c7b9ef2917c85870d188486e5e119f05e9c5c1f07652f6859f2d0',
'field_ids_sha256': '8192aa87281f8da7d1a4fd7884472e83f57c344b059953db0aa94e359b46a046',
'full_card_sha256': 'a74e00d9d0c0d4f7560a7650618243d393431a1bee78f43047f38c57b5e5edd5',
'resource': 'toxoplasma-pv-segmentation-cpsam.json',
'resource_sha256': '504429f98ee4d86824bb9aedd4291d799be2f5a272d542961c32b09f158f368c',
'uri': 'https://huggingface.co/einarolafsson/toxoplasma-pv-segmentation-cpsam/resolve/main/cpsam_v2_toxo_r2?download=true'},
'toxoplasma_pv_v2': {'checkpoint_sha256': '17c689e3b117745561e20a885c2a2a998ed360fa97cac8c0446316ae5905c10f',
'field_ids_sha256': 'cb581da5ded17152fb834d7b3c24af0c1ad0ce6f472cc292fa81ede122b4a9ff',
'full_card_sha256': '41015c1eea6fd3366205c121208272d931ccb9954f24b48b90c34fd01533fe69',
'resource': 'toxoplasma-pv-segmentation-cpsam-r5.json',
'resource_sha256': '4a763040ef118e1e940033723c93ec2750e7880b56eb47f5eb8da703b847235f',
'uri': 'https://huggingface.co/einarolafsson/toxoplasma-pv-segmentation-cpsam-r5/resolve/main/cpsam_v2_toxo_r5?download=true'},
'toxoplasma_pv_v3': {'checkpoint_sha256': '146ef269979b1d1ab45c11039b0ab164f68001adaa8f73f1f8f18be6fcfd060e',
'field_ids_sha256': 'b0196c1902fc479d76aa8df2afbc2c4046088413557ba09c96a1b3363618c1fb',
'full_card_sha256': '544e2321dd7af7b996af19e4ae16f5270e727279d0516ad36b932cd323ac3aad',
'resource': 'toxoplasma-pv-segmentation-cpsam-r6.json',
'resource_sha256': 'd8621a89fc25ff18801aebf32b8caf75a8af06045feba5b97d9833baa751f84a',
'uri': 'https://huggingface.co/einarolafsson/toxoplasma-pv-segmentation-cpsam-r6/resolve/main/weights/cpsam_v2_toxo_r6'},
'toxoplasma_pv_v4': {'checkpoint_sha256': '621a475c9bfb6865be7de12c5c4638f4509d734e0c4d39adb15ea541e48961d2',
'field_ids_sha256': 'd0a6542dd6647b7c1b46742151596d850ec5f2c5dfd0b30fe6582be7684988d8',
'full_card_sha256': '76ebd4a89abb9f98e65b668ce811d85e0cc5a6e5380de8cc168f7705c2b23c3a',
'resource': 'toxoplasma-pv-segmentation-cpsam-r7.json',
'resource_sha256': '0c92b1167245b2a6ae901d26623761269ed1d65ed0c109b5c0da91a797b50ff5',
'uri': 'https://huggingface.co/einarolafsson/toxoplasma-pv-segmentation-cpsam-r7/resolve/main/weights/cpsam_v2_toxo_r7'},
'toxoplasma_pv_v5': {'checkpoint_sha256': 'a1d7f1a978cf4925ec60db486a00eea33e87e6fc5a6b0d0f34c6db5f132d5c1c',
'field_ids_sha256': 'ff26b2e63f147aaa48da545addadab32505d8417654e7f5d9000ee54187b149d',
'full_card_sha256': '73b49ea573201d9d2d1e51513398dc9917ff10cbee792bae2a499d18c81ceed2',
'resource': 'toxoplasma-pv-segmentation-cpsam-r8.json',
'resource_sha256': '34c9df8121b591b5c5c8d0b8566dea5a8294fdb1c26197df7a42180db9d4d40f',
'uri': 'https://huggingface.co/einarolafsson/toxoplasma-pv-segmentation-cpsam-r8/resolve/main/weights/cpsam_v2_toxo_r8'}}
[docs]
class ModelZooError(Exception):
"""Base class for every refusal in this module."""
[docs]
class ChecksumMismatch(ModelZooError):
"""Downloaded bytes do not match the published digest; nothing was installed."""
[docs]
class ModelUnreadable(ModelZooError):
"""A model file is missing, empty, or not a checkpoint. Names the file."""
[docs]
class DownloadCancelled(ModelZooError):
"""Raised when the caller cancels a fetch; nothing is left at the destination."""
[docs]
class IncomparableBenchmarks(ModelZooError):
"""Benchmarks from different field sets cannot be ranked against each other."""
@dataclass(frozen=True)
[docs]
class ModelEntry:
"""Record of one model: its location, checksum and provenance.
Frozen because an entry is a *record of a file at a moment* — the hash, the
size and the provenance describe those bytes. Changing one in place would
silently invalidate the other two; :func:`dataclasses.replace` makes the
new record explicit.
:param key: stable id, unique within a listing. For a local file this is
derived from the filename; for a catalogue entry it is whatever the
catalogue declared.
:param name: the filename (or the published name) — what a human reads.
:param kind: ``'cellpose'`` or ``'classifier'``; see :data:`KINDS`.
:param source: ``'bundled'`` (ships with spaCR), ``'local'`` (found on this
machine) or ``'remote'`` (declared in a catalogue, not yet fetched).
:param path: absolute path on this machine, or ``''`` for a remote entry.
:param uri: where a remote entry is fetched from, or ``''``.
:param version: the zoo's own version number for a filename. ``'1'`` for a
plain name, ``'2'`` for ``foo_v2.CP_model`` (see
:func:`versioned_path`), or whatever a catalogue declared.
:param sha256: hex digest. For a downloaded model this is the digest of the
bytes that were actually written; for a catalogue entry it is the
published digest to check against; ``''`` means "no checksum known",
which :func:`fetch` treats as a refusal rather than a pass.
:param size_bytes: file size, ``0`` when unknown.
:param trained_on: what data produced this model, in prose, or
:data:`UNKNOWN`. Never ``''``.
:param trained_by: who produced it, or :data:`UNKNOWN`. Never ``''``.
:param metrics: whatever numbers came with it — for a classifier, the
best/last epoch metrics :func:`spacr.train_compare.load_run` recovered.
Excluded from equality: two records of the same bytes are the same
model whether or not somebody attached numbers to one of them.
:param notes: everything the reader needs to know that is not a field:
missing provenance, an unverified download, a file that does not look
like a checkpoint.
:param verified: True only when :attr:`sha256` was checked against a
published digest. A downloaded file whose hash was merely *recorded* is
not verified, and says so.
:param settings_path: where the provenance came from, for the reader who
wants to go and look at it.
:param licence: the licence the model or package is published under, as
its publisher states it (an SPDX identifier where there is one), or
``''`` when none is recorded.
:param _curated_card: Internal presentation marker set only after the
builtin declaration, checkpoint and packaged card match reviewed
identities. Serialized catalogue fields do not enable this marker.
"""
key: str
name: str
kind: str = "cellpose"
source: str = "local"
path: str = ""
uri: str = ""
version: str = "1"
sha256: str = ""
size_bytes: int = 0
trained_on: str = UNKNOWN
trained_by: str = UNKNOWN
metrics: Dict[str, Any] = _dc_field(default_factory=dict, compare=False)
notes: Tuple[str, ...] = ()
verified: bool = False
settings_path: str = ""
licence: str = ""
_curated_card: bool = _dc_field(default=False, compare=False, repr=False)
[docs]
def __post_init__(self):
"""Fill in the provenance fields and validate the kind.
A blank ``trained_on`` or ``trained_by`` reads as "no constraints", so
it is replaced with :data:`UNKNOWN`.
:raises ValueError: if ``kind`` is not one of the known model kinds.
"""
for attribute in ("trained_on", "trained_by"):
value = str(getattr(self, attribute) or "").strip()
object.__setattr__(self, attribute, value or UNKNOWN)
object.__setattr__(self, "notes", tuple(self.notes))
if self.kind not in KINDS:
raise ValueError(f"kind must be one of {KINDS}, got {self.kind!r}")
@property
[docs]
def exists(self) -> bool:
"""True when :attr:`path` names a file that is here now."""
return bool(self.path) and os.path.isfile(self.path)
@property
[docs]
def provenance_known(self) -> bool:
"""True when this model says what it was trained on."""
return self.trained_on != UNKNOWN
@property
[docs]
def scorecard_known(self) -> bool:
"""True when this model says how accurate it is.
The accuracy counterpart of :attr:`provenance_known`. A model with no
metrics did not score zero, but a table of empty cells reads as zero,
so the absence is reported as a state rather than shown as blank
cells.
"""
from .scorecard import scorecard_is_present
return scorecard_is_present(self.metrics or {})
@property
[docs]
def scorecard_holdout(self) -> str:
"""``name @ version`` of the hold-out set, or ``""``.
Scores from different hold-out sets are not comparable: two F1
values for the same model agree only if they were computed on the
same masks. The set name and version are therefore shown with the
numbers.
"""
metrics = self.metrics or {}
name = str(metrics.get("holdout") or "").strip()
version = str(metrics.get("holdout_version") or "").strip()
if not name:
return ""
return f"{name} @ {version}" if version else name
[docs]
def scorecard_lines(self) -> List[str]:
"""The scorecard as display lines, or the sentence saying there is none.
Trusted builtin cards show all six accepted matched-stock measurements
and their dataset and counts; their boundary F1 uses a two-pixel
tolerance. Other entries keep :func:`spacr.scorecard.headline` applied
to their fine-tuned and vanilla values, followed by the recorded
hold-out set. One-pixel measurements retain their original meaning.
Without a scorecard the list holds :data:`spacr.scorecard.NO_SCORECARD`.
"""
from .scorecard import NO_SCORECARD, headline
if not self.scorecard_known:
return [NO_SCORECARD]
metrics = self.metrics or {}
if self._curated_card:
return _curated_scorecard_lines(metrics)
finetuned = {k: v.get("finetuned") for k, v in metrics.items()
if isinstance(v, dict)}
baseline = {k: v.get("vanilla") for k, v in metrics.items()
if isinstance(v, dict)}
lines = list(headline(finetuned, baseline=baseline))
holdout = self.scorecard_holdout
if holdout:
lines.append(f"hold-out set {holdout}")
return lines
@property
[docs]
def checksum_state(self) -> str:
"""What the checksum column says, in one word.
``'none'``
no hash at all — nothing can be checked, and :func:`fetch` refuses
such an entry unless the caller overrides it.
``'published'``
a hash came with the entry but the bytes are not here yet; it is
the digest the download will be checked against.
``'recorded'``
the hash of the file on disk is known, but no published digest
exists to compare it against. It shows only that the file has not
changed since the hash was recorded.
``'verified'``
the bytes on disk were compared with a published digest and match.
"""
if not self.sha256:
return "none"
if self.verified:
return "verified"
return "recorded" if self.exists else "published"
[docs]
def summary_line(self) -> str:
"""One line for a list widget."""
bits = [self.name, self.kind, self.source, f"v{self.version}"]
bits.append(f"trained on: {_shorten(self.trained_on, 48)}")
bits.append(f"checksum {self.checksum_state}")
if self.notes:
bits.append(f"! {len(self.notes)} note"
f"{'s' if len(self.notes) > 1 else ''}")
return " · ".join(bits)
@property
[docs]
def model_card_url(self) -> str:
"""The Hugging Face page for this model, derived from its download uri.
A checksum and a metrics table are not enough on their own: the reader
wants the page that says what the model was trained on and shows its
training curves. Derived rather than declared, so every Hugging Face
entry has one without a per-entry field to forget.
"""
uri = str(self.uri or "")
marker = "huggingface.co/"
if marker not in uri:
return ""
rest = uri.split(marker, 1)[1]
parts = [p for p in rest.split("/") if p]
if len(parts) >= 2 and parts[0] == "datasets":
parts = parts[1:]
if len(parts) < 2:
return ""
return f"https://huggingface.co/{parts[0]}/{parts[1]}"
[docs]
def describe(self) -> str:
"""The multi-line provenance card shown next to a selected model.
Trusted builtin cards show the six accepted measurements and explicit
evaluation identity. Their historical evaluations and source digests
remain in the entry's data without cluttering this presentation.
Other entries retain their original complete metadata description.
"""
lines = [
f"{self.name} [{self.kind} · {self.source} · v{self.version}]",
f" path {self.path or '(not downloaded)'}",
]
if self.uri:
lines.append(f" uri {self.uri}")
lines.append(f" size {_human_bytes(self.size_bytes)}")
lines.append(f" sha256 {self.sha256 or '(none published)'} "
f"({self.checksum_state})")
lines.append(f" trained on {self.trained_on}")
lines.append(f" trained by {self.trained_by}")
if self.licence:
lines.append(f" licence {self.licence}")
card = self.model_card_url
if card:
lines.append(f" model card {card}")
if self.settings_path:
lines.append(f" provenance {self.settings_path}")
if self._curated_card:
lines.extend(" " + line for line in self.scorecard_lines())
for name, value in sorted(self.metrics.items()):
if self._curated_card and name not in (
"n_train", "train_objects", "cv", "train_loss", "val_loss", "best_epoch"):
continue
lines.append(f" {name:<10} {value}")
for note in self.notes:
lines.append(f" ! {note}")
if not self.provenance_known:
lines.append(
" ! this model does not say what it was trained on, so "
"nothing here tells you whether it suits your images.")
return "\n".join(lines)
def _shorten(text: Any, width: int) -> str:
"""Truncate text to a width, marking where it was cut.
:param text: the text.
:param width: the maximum length INCLUDING the ellipsis, so a column
laid out at this width never overflows.
:returns: the text, or its prefix with an ellipsis.
"""
text = str(text)
return text if len(text) <= width else text[:width - 1] + "…"
def _human_bytes(size: Any) -> str:
"""Bytes as something a person reads, ``'unknown'`` for 0/None."""
try:
n = float(size)
except (TypeError, ValueError):
return UNKNOWN
if n <= 0:
return UNKNOWN
for unit in ("B", "KB", "MB", "GB"):
if n < 1024:
return f"{n:.0f} {unit}" if unit == "B" else f"{n:.1f} {unit}"
if unit == "GB":
continue
n /= 1024.0
return f"{n:.1f} GB"
[docs]
def sha256_file(path: Any, chunk_size: int = 1 << 20) -> str:
"""Hex SHA-256 of a file, read in chunks so the whole file is never in memory.
:param path: the file.
:param chunk_size: bytes per read.
:returns: the lowercase hex digest.
:raises ModelUnreadable: when the file is missing or cannot be read, with
the path in the message.
"""
p = Path(path)
digest = hashlib.sha256()
try:
with p.open("rb") as handle:
while True:
block = handle.read(chunk_size)
if not block:
break
digest.update(block)
except FileNotFoundError:
raise ModelUnreadable(f"no such model file: {p}") from None
except OSError as e:
raise ModelUnreadable(f"could not read {p}: {e}") from None
return digest.hexdigest()
[docs]
def verify(entry: ModelEntry, expected: Optional[str] = None) -> bool:
"""Hash the file this entry points at and compare it to a known digest.
:param entry: the entry to check.
:param expected: the digest to compare against; defaults to
:attr:`ModelEntry.sha256`.
:returns: True when the file's digest matches.
:raises ModelUnreadable: when the entry has no local file, naming it.
:raises ModelZooError: when there is no digest to compare against — that is
a caller error, and returning False for it would read as "the file is
wrong" when what happened is "nobody said what right looks like".
"""
if not entry.path:
raise ModelUnreadable(
f"{entry.name} has not been downloaded, so there is nothing to "
f"verify (source={entry.source}, uri={entry.uri or 'none'})")
if not os.path.isfile(entry.path):
raise ModelUnreadable(f"no such model file: {entry.path}")
want = (expected if expected is not None else entry.sha256) or ""
want = want.strip().lower()
if not want:
raise ModelZooError(
f"no checksum recorded for {entry.name} ({entry.path}) — there is "
f"nothing to verify it against. Compute one with sha256_file() and "
f"put it in the catalogue.")
return sha256_file(entry.path) == want
def _looks_like_checkpoint(path: Path) -> bool:
"""True when the first bytes are a torch save (zip or legacy pickle)."""
try:
with path.open("rb") as handle:
head = handle.read(4)
except OSError:
return False
return any(head.startswith(magic) for magic in _TORCH_MAGICS)
[docs]
def classify_kind(path: Any) -> Optional[str]:
"""Say whether a file is a Cellpose model, a classifier, or not a model.
The rules, in order:
1. ``*.CP_model`` is a Cellpose checkpoint — that is what
:func:`spacr.submodules.train_cellpose` names its output.
2. ``*.pth`` / ``*.pt`` is a classifier checkpoint
(:func:`spacr.io._save_model` writes
``<model_type>_epoch_<n>_channels_<ch>.pth``) **unless** it sits in a
Cellpose folder or has ``cellpose``/``cp_model`` in its name.
3. An **extensionless** file inside a Cellpose folder is a Cellpose
checkpoint only if its first bytes are a torch save. ``cellpose.train``
writes ``<save_path>/models/<name>`` with no suffix, and that folder
also holds READMEs and logs — the magic-byte check is what keeps a
``README`` out of the zoo.
4. Anything else is not a model. CSVs, PNGs, ``.npy`` masks and settings
snapshots all land here and are ignored.
:param path: a file path.
:returns: ``'cellpose'``, ``'classifier'`` or None.
"""
p = Path(path)
low = p.name.lower()
near = {p.parent.name.lower(), p.parent.parent.name.lower()}
in_cellpose_dir = bool(near & set(CELLPOSE_DIR_NAMES))
if low.endswith(CELLPOSE_SUFFIXES):
return "cellpose"
if low.endswith(CLASSIFIER_SUFFIXES):
if "cellpose" in low or "cp_model" in low or in_cellpose_dir:
return "cellpose"
return "classifier"
if not p.suffix and in_cellpose_dir:
return "cellpose" if _looks_like_checkpoint(p) else None
return None
[docs]
def inspect_checkpoint(path: Any, loader: Optional[Callable[[str], Any]] = None,
deep: bool = False) -> Dict[str, Any]:
"""Check a file is a loadable checkpoint, failing with the filename in it.
The default failure for a wrong or corrupt checkpoint is a ``KeyError`` on
a state-dict key raised somewhere inside torch, which names nothing the
user chose and reads like a spaCR bug. This turns all of it — missing,
empty, truncated, a PNG somebody renamed, a Cellpose model handed to the
classifier path — into one :class:`ModelUnreadable` naming the file.
The shallow check needs no torch at all: it is a stat and four bytes.
:param path: the checkpoint.
:param loader: ``fn(path) -> object`` used for the deep check; defaults to
``torch.load(..., map_location='cpu')``, imported only if used.
:param deep: actually load the file. Off by default because loading a 2 GB
checkpoint to populate a list widget is not acceptable.
:returns: ``{'path', 'size_bytes', 'format', 'loaded'}``.
:raises ModelUnreadable: naming the file, always.
"""
p = Path(path)
if not p.exists():
raise ModelUnreadable(f"no such model file: {p}")
if p.is_dir():
raise ModelUnreadable(
f"{p} is a directory, not a model checkpoint — point at the file "
f"inside it")
size = p.stat().st_size
if size == 0:
raise ModelUnreadable(
f"{p} is empty (0 bytes) — an interrupted download or copy leaves "
f"exactly this")
try:
with p.open("rb") as handle:
head = handle.read(4)
except OSError as e:
raise ModelUnreadable(f"could not read {p}: {e}") from None
if not any(head.startswith(magic) for magic in _TORCH_MAGICS):
raise ModelUnreadable(
f"{p} is not a PyTorch checkpoint: it starts with {head!r}, and a "
f".pth / .CP_model file starts with a zip header or a pickle "
f"opcode. Check the path — this is usually a text file, an HTML "
f"error page saved by a failed download, or the wrong file "
f"entirely.")
fmt = "zip" if head.startswith(b"PK\x03\x04") else "pickle"
out: Dict[str, Any] = {"path": str(p), "size_bytes": int(size),
"format": fmt, "loaded": False}
if not deep:
return out
load = loader if loader is not None else _torch_loader
try:
load(str(p))
except ModelUnreadable:
raise
except Exception as e:
raise ModelUnreadable(
f"{p} could not be loaded as a model checkpoint "
f"({type(e).__name__}: {e}). The file is a torch save but not the "
f"architecture that was asked for — check that this is a "
f"{'Cellpose' if p.name.lower().endswith(CELLPOSE_SUFFIXES) else 'classifier'} "
f"model and not the other kind.") from None
out["loaded"] = True
return out
def _torch_loader(path: str) -> Any:
"""``torch.load`` on the CPU. Imported here so the module stays torch-free."""
import torch
return torch.load(path, map_location="cpu", weights_only=False)
def _read_key_value_csv(path: Path) -> Dict[str, Any]:
"""Read a ``Key,Value`` settings CSV the way the run diff reads one.
Reuses :func:`spacr.run_journal._read_settings_csv` rather than writing a
second parser: that one already handles the JSON/CSV/live-dict round trips
a spaCR settings dict takes, and the zoo has to agree with the run diff
about what a settings file says.
"""
from .run_journal import _read_settings_csv
return _read_settings_csv(path)
def _settings_beside(path: Path) -> Tuple[Dict[str, Any], str]:
"""Find the settings snapshot beside a Cellpose checkpoint.
:func:`spacr.submodules.train_cellpose` calls ``save_settings(settings,
name=model_name)``, which writes ``<src>/settings/<model_name>.csv``; the
bundled pack ships ``<model_file>_settings.csv`` next to the weights. Both
are looked for, nearest first.
:returns: ``(settings, where)``; ``({}, '')`` when there is none.
"""
candidates = [
path.with_name(path.name + "_settings.csv"),
path.with_name(path.stem + "_settings.csv"),
path.with_suffix(".csv"),
]
node = path.parent
for _ in range(DEFAULT_SCAN_DEPTH):
candidates.append(node / "settings" / f"{path.name}.csv")
candidates.append(node / "settings" / f"{path.stem}.csv")
node = node.parent
for candidate in candidates:
if not candidate.is_file():
continue
try:
settings = _read_key_value_csv(candidate)
except Exception:
continue
if settings:
return settings, str(candidate)
return {}, ""
def _describe_cellpose_training(settings: Mapping[str, Any]) -> str:
"""One prose line saying what a Cellpose model saw, or :data:`UNKNOWN`.
Deliberately concrete: magnification and confluence are not in the settings
file, but the source folder, the crop size, the diameter and the epoch
count are, and together they are enough for a reader to tell "this is not
my data".
"""
source = settings.get("img_src") or settings.get("src") or ""
bits: List[str] = []
if source:
bits.append(str(source))
shape = settings.get("width_height") or settings.get("target_size")
if shape:
bits.append(f"crops {shape}")
if settings.get("diameter"):
bits.append(f"diameter {settings['diameter']}")
if settings.get("n_epochs"):
bits.append(f"{settings['n_epochs']} epochs")
if settings.get("grayscale") is True or str(settings.get("grayscale")).lower() == "true":
bits.append("greyscale")
return ", ".join(bits) if bits else UNKNOWN
def _describe_classifier_training(settings: Mapping[str, Any]) -> str:
"""One prose line saying what a classifier saw, or :data:`UNKNOWN`."""
bits: List[str] = []
source = settings.get("src") or ""
if source:
bits.append(str(source))
if settings.get("model_type"):
bits.append(str(settings["model_type"]))
if settings.get("classes"):
bits.append(f"classes {settings['classes']}")
if settings.get("image_size"):
bits.append(f"{settings['image_size']}px crops")
if settings.get("epochs"):
bits.append(f"{settings['epochs']} epochs")
return ", ".join(bits) if bits else UNKNOWN
def _who(settings: Mapping[str, Any]) -> str:
"""Who trained it, from whatever the settings recorded, or :data:`UNKNOWN`."""
for key in ("user", "author", "trained_by", "operator", "hostname", "host"):
value = settings.get(key)
if value not in (None, "", "nan"):
return str(value)
return UNKNOWN
def _version_of(name: str) -> str:
"""The zoo's version number for a filename (see :func:`versioned_path`)."""
stem = Path(name).stem
match = _VERSION_RE.match(stem)
return match.group("n") if match else "1"
def _key_for(path: Path) -> str:
"""A short, stable key for a local file: its name, made filesystem-safe."""
return re.sub(r"[^A-Za-z0-9_.-]+", "_", path.name)
[docs]
def entry_from_file(path: Any, kind: Optional[str] = None,
source: str = "local", key: Optional[str] = None,
runs: Optional[Mapping[str, Any]] = None,
compute_hash: bool = False,
sha256: str = "", verified: bool = False,
extra_notes: Sequence[str] = ()) -> ModelEntry:
"""Build a :class:`ModelEntry` for a checkpoint on this machine.
Provenance is recovered from whatever spaCR already wrote beside the model:
a ``*_settings.csv`` or ``<src>/settings/<name>.csv`` for a Cellpose model,
and — for a classifier — the training run the checkpoint sits in, loaded
through :func:`spacr.train_compare.load_run` so the zoo and the training-run
comparison agree about where settings live and what they say.
:param path: the checkpoint.
:param kind: override :func:`classify_kind`.
:param source: ``'local'`` or ``'bundled'``.
:param key: override the generated key.
:param runs: ``{folder: TrainingRun}`` from :func:`_runs_under`, so a scan
of 40 checkpoints in one run folder reads that folder once.
:param compute_hash: hash the file now. Off by default: hashing every
checkpoint on a machine to populate a list widget is minutes.
:param sha256: a digest already known for these bytes.
:param verified: whether ``sha256`` was checked against a published digest.
:param extra_notes: notes to carry onto the entry.
:returns: the entry.
:raises ModelUnreadable: when the path is not a file.
"""
p = Path(path)
if not p.is_file():
raise ModelUnreadable(f"no such model file: {p}")
kind = kind or classify_kind(p) or "classifier"
notes: List[str] = list(extra_notes)
settings: Dict[str, Any] = {}
settings_path = ""
metrics: Dict[str, Any] = {}
if kind == "classifier":
run = _run_for(p, runs)
if run is not None:
settings = dict(run.settings)
settings_path = run.settings_path
metrics = _metrics_from_run(run)
if not settings:
settings, settings_path = _settings_beside(p)
trained_on = _describe_classifier_training(settings)
else:
settings, settings_path = _settings_beside(p)
trained_on = _describe_cellpose_training(settings)
if trained_on == UNKNOWN:
notes.append(
"no settings snapshot found beside this model, so what it was "
"trained on is unknown — treat it as untested on your images")
if not _looks_like_checkpoint(p):
notes.append(
f"{p.name} does not start with a torch header; it may be a Git LFS "
f"pointer, a failed download, or not a model at all")
try:
size = p.stat().st_size
except OSError:
size = 0
digest = sha256
if compute_hash and not digest:
digest = sha256_file(p)
return ModelEntry(
key=key or _key_for(p),
name=p.name,
kind=kind,
source=source,
path=str(p.resolve()),
version=_version_of(p.name),
sha256=digest,
size_bytes=int(size),
trained_on=trained_on,
trained_by=_who(settings),
metrics=metrics,
notes=tuple(notes),
verified=bool(verified and digest),
settings_path=settings_path,
)
def _metrics_from_run(run: Any) -> Dict[str, Any]:
"""Best/last accuracy off a :class:`spacr.train_compare.TrainingRun`.
Both, never one: the best epoch of a validation curve was chosen using that
curve, so it is optimistically biased; the last epoch is unbiased but may
be well past the optimum. :mod:`spacr.train_compare` makes that argument at
length and reports both for it — the zoo shows the same pair for the same
reason.
"""
out: Dict[str, Any] = {}
final = getattr(run, "final_metrics", None) or {}
for label, entry in final.items():
if str(entry.get("split")) != "val":
continue
best = (entry.get("best") or {}).get("accuracy")
last = (entry.get("last") or {}).get("accuracy")
if best:
out[f"{label} best accuracy"] = f"{best['value']:.4f} @ epoch {best['epoch']}"
if last:
out[f"{label} last accuracy"] = f"{last['value']:.4f} @ epoch {last['epoch']}"
if not out:
for label, entry in final.items():
best = (entry.get("best") or {}).get("accuracy")
if best:
out[f"{label} best accuracy"] = (
f"{best['value']:.4f} @ epoch {best['epoch']} "
f"(train split — not held out)")
return out
def _run_for(path: Path, runs: Optional[Mapping[str, Any]]) -> Any:
"""The discovered training run a checkpoint belongs to, or one loaded now."""
if runs:
for folder in (path.parent, path.parent.parent):
run = runs.get(str(folder.resolve()))
if run is not None:
return run
from .train_compare import load_run
for folder in (path.parent, path.parent.parent):
try:
return load_run(folder)
except Exception:
continue
return None
def _runs_under(root: Path) -> Dict[str, Any]:
"""``{folder: TrainingRun}`` for every training run below ``root``.
Straight reuse of :func:`spacr.train_compare.find_runs` — the classifier
half of the zoo *is* the training-run scan, and a second discovery pass
would drift out of step with it the first time the on-disk layout changed.
Failures are swallowed: a zoo that cannot list local models because one
folder was unreadable is worse than one with thinner provenance.
"""
try:
from .train_compare import find_runs
return {str(Path(run.path).resolve()): run for run in find_runs(root)}
except Exception:
return {}
def _model_cache_root() -> Path:
"""Return spaCR's shared checkpoint cache, including an explicit relocation."""
configured = os.environ.get("SPACR_MODEL_CACHE", "").strip()
return Path(configured).expanduser() if configured else _spacr_home() / "models"
[docs]
def package_model_root() -> Path:
"""``<spacr>/resources/models`` — where the bundled pack lives.
The same folder :func:`spacr.utils.download_models` fills and
:func:`spacr.submodules.analyze_plaques` reads from.
"""
return Path(__file__).resolve().parent / "resources" / "models"
[docs]
def default_local_roots() -> List[Path]:
"""Default folders to scan for models when the caller names none.
The bundled pack, the Cellpose user folder, and spaCR's own model cache.
Only the ones that exist come back.
"""
roots = [
package_model_root(),
Path.home() / ".cellpose" / "models",
_model_cache_root(),
]
return [r for r in roots if r.is_dir()]
def _walk(base: Path, max_depth: int, limit: int) -> Iterator[Path]:
"""Every file at most ``max_depth`` levels below ``base``, hidden dirs skipped.
``limit`` caps the files *examined*, not the models found, so a folder with
a million PNGs in it is slow rather than fatal and does not silently drop
the models below them.
"""
stack: List[Tuple[Path, int]] = [(base, 0)]
seen = 0
while stack and seen < limit:
node, depth = stack.pop(0)
try:
children = sorted(node.iterdir())
except OSError:
continue
for child in children:
if child.name.startswith("."):
continue
if child.is_dir():
if depth < max_depth:
stack.append((child, depth + 1))
elif child.is_file():
seen += 1
yield child
if seen >= limit:
return
[docs]
def discover_local(roots: Any = None, max_depth: int = DEFAULT_SCAN_DEPTH,
compute_hashes: bool = False,
limit: int = DEFAULT_SCAN_LIMIT) -> List[ModelEntry]:
"""Find the model checkpoints already on this machine.
Cellpose models and classifier checkpoints are told apart by
:func:`classify_kind`; everything else in the folders — settings CSVs, mask
``.npy`` files, montage PNGs, logs — is ignored.
Nothing is downloaded and nothing is hashed unless ``compute_hashes`` is
set: this is the function behind a list widget, and it has to be fast
enough to run on a folder the user just typed.
:param roots: a folder, a file, or an iterable of them; None uses
:func:`default_local_roots`.
:param max_depth: how deep below each root to look.
:param compute_hashes: hash every file found (minutes on a big folder).
:param limit: stop after examining this many files per root.
:returns: entries, Cellpose first, then by name.
"""
entries: List[ModelEntry] = []
seen: set = set()
for root in _as_paths(roots if roots is not None else default_local_roots()):
if root.is_file():
candidates: List[Path] = [root]
runs: Dict[str, Any] = {}
elif root.is_dir():
candidates = list(_walk(root, max_depth, limit))
runs = _runs_under(root)
else:
continue
for path in candidates:
kind = classify_kind(path)
if kind is None:
continue
try:
resolved = str(path.resolve())
except OSError:
continue
if resolved in seen:
continue
seen.add(resolved)
source = ("bundled" if _under(path, package_model_root())
else "local")
try:
entries.append(entry_from_file(path, kind=kind, source=source,
runs=runs,
compute_hash=compute_hashes))
except ModelUnreadable:
continue
entries.sort(key=lambda e: (0 if e.kind == "cellpose" else 1, e.name))
return entries
def _under(path: Path, root: Path) -> bool:
"""Whether a path is inside a root, after resolving both.
Resolved first, so a symlink or a ``..`` cannot escape the root while
appearing to be under it -- this gates where a downloaded checkpoint may
be written.
:param path: the path to test.
:param root: the root it must be under.
:returns: ``True`` when it is; ``False`` when it is not, and also when
either path cannot be resolved.
"""
try:
path.resolve().relative_to(root.resolve())
return True
except (ValueError, OSError):
return False
def _as_paths(roots: Any) -> List[Path]:
"""Normalise a folder / path / iterable-of-those into a list of Paths."""
if roots is None:
return []
if isinstance(roots, (str, os.PathLike)):
return [Path(roots)]
return [Path(r) for r in roots]
[docs]
def hf_uri(repo_id: str, filename: str, repo_type: str = "dataset") -> str:
"""The download URL for a file in a Hugging Face repo.
:param repo_id: Hugging Face repository identifier.
:param filename: repository-relative name of the file to download.
:param repo_type: ``"dataset"`` (the default, and what spaCR shipped
first) or ``"model"``.
Exactly the URL :func:`spacr.utils.download_models` and
:func:`spacr.qt.hf_download._download_one` build, kept in one place so the
zoo cannot drift away from the downloader spaCR already ships.
The two repository kinds have different URLs: a dataset
file lives under ``/datasets/<repo>/resolve/...`` and a model file under
``/<repo>/resolve/...``. Asking for one at the other's URL returns a 404
page, and a downloader that does not check the content type writes that
HTML into the destination and leaves a "checkpoint" that fails to load
with a torch error naming neither the URL nor the repo.
``dataset`` remains the default because :data:`HF_MODELS_REPO` is a
DATASET repo -- ``einarolafsson/models`` -- and every entry written before
this parameter existed assumes it. New model repos pass ``"model"``.
"""
kind = str(repo_type or "dataset").lower()
if kind not in ("dataset", "model"):
raise ValueError(
f"repo_type must be 'dataset' or 'model', not {repo_type!r}")
prefix = "datasets/" if kind == "dataset" else ""
return (f"https://huggingface.co/{prefix}{repo_id}/resolve/main/"
f"{filename}?download=true")
def _entry_from_mapping(data: Mapping[str, Any],
source: str = "remote") -> ModelEntry:
"""One catalogue record -> a :class:`ModelEntry`."""
name = str(data.get("name") or data.get("key") or "")
if not name:
raise ValueError("a catalogue entry needs at least a name")
uri = data.get("uri")
if not uri:
uri = hf_uri(str(data.get("repo_id") or HF_MODELS_REPO), name,
str(data.get("repo_type") or "dataset"))
notes = tuple(str(n) for n in (data.get("notes") or ()))
sha = str(data.get("sha256") or "").strip().lower()
if not sha:
notes = notes + (
"no published checksum — this entry cannot be verified, and fetch "
"refuses it unless you explicitly accept that",)
return ModelEntry(
key=str(data.get("key") or name),
name=name,
kind=str(data.get("kind") or "cellpose"),
source=str(data.get("source") or source),
path=str(data.get("path") or ""),
uri=str(uri),
version=str(data.get("version") or "1"),
sha256=sha,
size_bytes=int(data.get("size_bytes") or 0),
trained_on=data.get("trained_on") or UNKNOWN,
trained_by=data.get("trained_by") or UNKNOWN,
metrics=dict(data.get("metrics") or {}),
notes=notes,
licence=str(data.get("licence") or data.get("license") or ""),
)
def _builtin_entry_from_mapping(record: Mapping[str, Any]) -> ModelEntry:
"""Attach an accepted card only to its exact trusted builtin checkpoint.
The declaration object, download URI, checkpoint digest and packaged card
digest must all match the reviewed bindings. Custom declarations and changed
files retain their original measurements. The earlier evaluation remains
available in ``historical_scorecards``; no score is recalculated.
:param record: One declaration from ``BUNDLED_REMOTE_MODELS``.
:returns: Its model entry, with an approved matched-stock evaluation when bound.
"""
from copy import deepcopy
entry = _entry_from_mapping(record)
binding = _CURATED_SCORECARD_BINDINGS.get(entry.key)
if (binding is None or not any(record is raw for raw in BUNDLED_REMOTE_MODELS)
or entry.source != "remote" or entry.path
or entry.uri != binding["uri"]
or entry.sha256 != binding["checkpoint_sha256"]):
return entry
from . import model_cards
try:
payload = (model_cards._CARD_DIR / binding["resource"]).read_bytes()
except OSError:
return entry
payload = payload.replace(b"\r\n", b"\n")
if hashlib.sha256(payload).hexdigest() != binding["resource_sha256"]:
return entry
card = json.loads(payload)
scores = card["scores"]
rows = {row["key"]: row for row in scores["rows"]}
holdout = scores["holdout"]
metrics = {key: deepcopy(value) for key, value in entry.metrics.items()
if key in ("n_train", "train_objects", "cv", "train_loss", "val_loss", "best_epoch")}
for key, source in (("f1", "object_f1"), ("boundary_f1_2px", "boundary_f1"),
("precision", "precision"), ("recall", "recall"),
("aji", "aji"), ("dice", "dice")):
row = rows[source]
mine, stock = row["this"], row["stock"]
metrics[key] = {"finetuned": mine, "vanilla": stock,
"delta": None if stock is None else mine - stock}
counts = scores.get("counts", {}).get("this", {})
for key, source in (("true_positives", "tp"), ("false_positives", "fp"),
("false_negatives", "fn")):
if source in counts:
metrics[key] = counts[source]
metrics.update({"holdout": holdout.get("set", "gt_rescore_20261011"),
"holdout_version": "strict hold-out, pooled",
"n_fields": holdout["n_fields"], "n_objects": holdout["n_objects"],
"n_test": str(holdout["n_fields"]) + " fields",
"test_objects": str(holdout["n_objects"]),
"historical_scorecards": [{"metrics": deepcopy(entry.metrics)}],
"scorecard_provenance": {"schema": "spacr-approved-curated-card-v1",
"stock_model": scores["stock_model"], "boundary_tolerance_px": 2,
"holdout": deepcopy(holdout), "checkpoint_uri": entry.uri,
"checkpoint_sha256": entry.sha256,
"resource_sha256": binding["resource_sha256"],
"full_card_sha256": binding["full_card_sha256"],
"field_ids_sha256": binding["field_ids_sha256"],
"field_ids_included": False}})
return replace(entry, metrics=metrics, _curated_card=True)
_SHARED_CATALOGUE_CACHE: Dict[str, Any] = {"fetched_at": 0.0, "entries": ()}
#: Set while a background refresh is in flight, so a screen that opens twice
#: in a second starts one fetch rather than two.
_SHARED_CATALOGUE_FETCHING = threading.Event()
def _on_the_qt_gui_thread() -> bool:
"""Whether this call is running on Qt's GUI thread.
Answers False for a process with no Qt, for a worker thread, and for
anything that goes wrong while asking -- so the only thing this can do
is turn a blocking fetch into a background one, never the reverse.
"""
try:
from PySide6.QtCore import QCoreApplication, QThread
except Exception: # noqa: BLE001
return False
try:
app = QCoreApplication.instance()
return app is not None and QThread.currentThread() is app.thread()
except Exception: # noqa: BLE001
return False
def _refresh_shared_catalogue_in_background(uri: Optional[str],
timeout: float) -> None:
"""Fetch the catalogue on a daemon thread, for the cache to serve later.
Nothing waits on the thread and nothing is redrawn when it lands: the
point is only that the NEXT caller answers from a warm cache instead of
from the network.
"""
if _SHARED_CATALOGUE_FETCHING.is_set():
return
_SHARED_CATALOGUE_FETCHING.set()
def run() -> None:
"""Fetch the catalogue off the GUI thread, and always release the flag.
The flag is cleared in a `finally` block. `_SHARED_CATALOGUE_FETCHING`
stops a second refresh being started while this one is in flight, so a
fetch that raises must still clear it or no later refresh can ever
begin.
"""
try:
shared_catalogue(uri, timeout=timeout, force=True, block=True)
finally:
_SHARED_CATALOGUE_FETCHING.clear()
threading.Thread(target=run, daemon=True,
name="spacr-model-catalogue").start()
[docs]
def shared_catalogue(uri: Optional[str] = None, *,
timeout: float = DEFAULT_TIMEOUT,
force: bool = False,
block: Optional[bool] = None) -> Tuple["ModelEntry", ...]:
"""The community catalogue, fetched from :data:`REMOTE_CATALOGUE_URI`.
:param uri: override the catalogue location.
:param timeout: seconds to wait for the request.
:param force: ignore the cache and re-fetch.
:param block: whether to wait for the network. ``None`` (the default)
waits everywhere except Qt's GUI thread, where it answers from the
cache and refreshes on a daemon thread. ``True`` waits wherever it is called, which only a caller that
knows it is on a worker or in a CLI may ask for. ``False`` never
waits.
:returns: the entries, or ``()`` when the catalogue cannot be read.
It does not fetch on Qt's GUI thread. This function is reached from ``spacr.settings.downloaded_zoo_models`` while a settings
panel is being built, so it ran inside ``MainWindow._on_nav_selected``
with nothing able to paint or answer the compositor. Measured
with the catalogue host non-routable (10.255.255.1, the shape of a down
VPN or a captive portal -- the connect neither completes nor is refused):
opening the Mask module took **32.2 s**, all of it a GUI thread stuck in
``urlopen``. GNOME asks a window whether it is alive after five seconds,
so the user saw GNOME's "force quit" dialog. With the fetch moved off the
thread the same open took 2.4 s.
The cost of not waiting is a first module open whose Cellpose dropdown
lists the bundled and local models but not the community ones; the
background refresh means the second one has them.
It does not raise. This runs when a user opens a module that offers a
model list, and the list is useful without the community rows: the bundled
entries and any local models are still there. A laptop on a train, a lab
behind a proxy and a Hugging Face outage all produce the same thing -- a
shorter list and a log line -- rather than a module that will not open.
A corrupt or hostile catalogue is the failure that would go unnoticed,
so entries that do not parse are dropped individually, and an
entry without a ``sha256`` still cannot be installed by :func:`fetch`
without an explicit override. A catalogue row is a claim about where a file
lives; the checksum is what makes it a claim about which file.
A row's own ``source`` field is dropped: every row here is ``shared``,
and which heading it lands under is :func:`source_of`'s decision, made
from its checksum and the repository it is fetched from.
"""
import time
if os.environ.get("SPACR_MODEL_DOWNLOAD", "").strip().casefold() in (
"0", "off", "false", "no"):
return tuple(_SHARED_CATALOGUE_CACHE["entries"])
target = uri or REMOTE_CATALOGUE_URI
now = time.time()
if (not force and float(_SHARED_CATALOGUE_CACHE["fetched_at"]) > 0
and now - float(_SHARED_CATALOGUE_CACHE["fetched_at"])
< CATALOGUE_CACHE_SECONDS):
return tuple(_SHARED_CATALOGUE_CACHE["entries"])
wait = (not _on_the_qt_gui_thread()) if block is None else bool(block)
if not wait:
_refresh_shared_catalogue_in_background(uri, timeout)
return tuple(_SHARED_CATALOGUE_CACHE["entries"])
try:
import urllib.request
with urllib.request.urlopen(target, timeout=timeout) as response:
payload = json.loads(response.read().decode("utf-8"))
except Exception as exc: # noqa: BLE001
_SHARED_CATALOGUE_CACHE["fetched_at"] = now
first = not _SHARED_CATALOGUE_CACHE.get("warned")
_SHARED_CATALOGUE_CACHE["warned"] = True
(LOG.info if first else LOG.debug)(
"shared model catalogue unavailable (%s): %s",
type(exc).__name__, exc)
return tuple(_SHARED_CATALOGUE_CACHE["entries"])
records = payload.get("models") if isinstance(payload, Mapping) else payload
entries: List[ModelEntry] = []
for record in (records or ()):
if isinstance(record, Mapping):
record = {k: v for k, v in record.items() if k != "source"}
try:
entries.append(_entry_from_mapping(record, source="shared"))
except Exception as exc: # noqa: BLE001
LOG.warning("skipping a shared catalogue entry: %s", exc)
_SHARED_CATALOGUE_CACHE.update(fetched_at=now, entries=tuple(entries))
return tuple(entries)
[docs]
def shared_catalogue_is_stale() -> bool:
"""Whether :func:`shared_catalogue` would go to the network to answer.
For a caller that wants to do the waiting somewhere it is allowed to --
a worker thread -- rather than get the cached answer and not know it was
one.
"""
stamp = float(_SHARED_CATALOGUE_CACHE["fetched_at"])
return stamp <= 0 or (time.time() - stamp) >= CATALOGUE_CACHE_SECONDS
[docs]
def publish_model(local_path: Any, repo_id: str, *,
key: str,
kind: str = "cellpose",
trained_on: str = UNKNOWN,
trained_by: str = UNKNOWN,
private: bool = False,
notes: Sequence[str] = ()) -> Dict[str, Any]:
"""Upload a model to Hugging Face and return its catalogue row.
:param local_path: the checkpoint to upload.
:param repo_id: ``<user>/<repo>`` -- a repository under your own account.
:param key: the short name spaCR will offer the model under.
:param kind: one of :data:`KINDS`.
:param trained_on: what the model was trained on. Other labs read this
field to decide whether the model applies to their images.
:param trained_by: who trained it, and roughly when.
:param private: keep the repo private. A private model cannot be fetched by
other spaCR users, so it is off by default.
:param notes: caveats worth carrying next to the model.
:returns: the catalogue row, with the sha256 filled in.
:raises ImportError: when ``huggingface_hub`` is not installed.
The checksum is computed here, from the file that was actually uploaded.
:func:`fetch` refuses an entry it cannot verify, so a row written by hand
without a hash produces a model nobody can install without disabling the
check.
Publishing does not distribute the model on its own: add the returned row
to the shared catalogue (:data:`REMOTE_CATALOGUE_URI`) and every spaCR user
sees it within :data:`CATALOGUE_CACHE_SECONDS`.
"""
try:
from huggingface_hub import HfApi
except ImportError as exc:
raise ImportError(
"Publishing a model needs the 'huggingface_hub' package:\n"
" pip install huggingface_hub\n"
"then log in with `huggingface-cli login`.") from exc
path = Path(str(local_path))
if not path.is_file():
raise ModelZooError(f"{path} is not a file")
if kind not in KINDS:
raise ValueError(f"kind must be one of {KINDS}, got {kind!r}")
api = HfApi()
api.create_repo(repo_id, repo_type="model", private=bool(private),
exist_ok=True)
api.upload_file(path_or_fileobj=str(path), path_in_repo=path.name,
repo_id=repo_id, repo_type="model")
row = {
"key": key,
"name": path.name,
"kind": kind,
"repo_id": repo_id,
"repo_type": "model",
"sha256": sha256_file(path),
"size_bytes": path.stat().st_size,
"trained_on": trained_on,
"trained_by": trained_by,
"notes": tuple(notes),
}
LOG.info("published %s to %s; add this row to the shared catalogue",
path.name, repo_id)
return row
[docs]
def load_catalogue_file(path: Any) -> List[ModelEntry]:
"""Read a JSON catalogue of remote models.
Format — a list, or an object with a ``models`` list::
{"models": [
{"key": "hela_60x",
"name": "hela_60x_confluent.CP_model",
"kind": "cellpose",
"uri": "https://…/hela_60x_confluent.CP_model",
"sha256": "9f86d0…",
"size_bytes": 26566572,
"trained_on": "HeLa, 60x, confluent monolayer, 512px crops",
"trained_by": "A. Researcher, 2026-02",
"metrics": {"note": "benchmarked on plate3 fields 1-3"}}
]}
``sha256`` is the field that decides whether the entry is usable without an
explicit override: :func:`fetch` refuses an entry without one unless
``require_checksum=False`` is passed.
:param path: the JSON file.
:returns: the entries.
:raises ModelZooError: when the file cannot be read or is not a catalogue,
naming the file.
"""
p = Path(path)
try:
data = json.loads(p.read_text(encoding="utf-8"))
except FileNotFoundError:
raise ModelZooError(f"no such catalogue file: {p}") from None
except (OSError, ValueError) as e:
raise ModelZooError(f"could not read the catalogue {p}: {e}") from None
records = data.get("models") if isinstance(data, dict) else data
if not isinstance(records, list):
raise ModelZooError(
f"{p} is not a model catalogue — expected a list of entries, or an "
f"object with a 'models' list, got {type(records).__name__}")
out: List[ModelEntry] = []
for i, record in enumerate(records):
if not isinstance(record, Mapping):
raise ModelZooError(
f"{p}: entry {i} is a {type(record).__name__}, not an object")
try:
out.append(_entry_from_mapping(record))
except ValueError as e:
raise ModelZooError(f"{p}: entry {i} is unusable — {e}") from None
return out
#: What each Cellpose stock model is, for the zoo listing. Cellpose publishes
#: names, not descriptions, and a row reading only "cpdino" tells nobody
#: whether it applies to their images.
STOCK_CELLPOSE_NOTES = {
"cpsam": ("Cellpose-SAM v1, the original SAM-based generalist. Superseded "
"by cpsam_v2 but kept for reproducing older runs."),
"cpsam_v2": ("Cellpose-SAM v2, the current Cellpose generalist and the "
"base every spaCR fine-tune here starts from."),
"cpdino": "Cellpose-DINO, a DINO-backbone generalist.",
"cpdino-vitb": "Cellpose-DINO with the larger ViT-B backbone.",
}
[docs]
def stock_cellpose_entries() -> List["ModelEntry"]:
"""Every model the installed Cellpose can fetch for itself.
These are not spaCR's files and carry no checksum of ours: Cellpose
downloads and verifies them, and the name is the path -- passing "cpsam"
to Cellpose resolves it. They are listed so that the zoo shows every model
the user can segment with, not only the models trained for spaCR.
"""
import sys
loaded = sys.modules.get("cellpose.models")
names = list(getattr(loaded, "MODEL_NAMES", ()) or ()) if loaded else []
names = names or list(STOCK_CELLPOSE_NOTES)
home = Path.home() / ".cellpose" / "models"
out = []
for name in names:
local = home / name
out.append(ModelEntry(
key=_key_for(Path(name)), name=name,
path=str(local) if local.is_file() else name,
kind="cellpose", source="stock", uri="", sha256="",
size_bytes=local.stat().st_size if local.is_file() else 0,
trained_on=STOCK_CELLPOSE_NOTES.get(
name, "Cellpose stock model; see the Cellpose documentation."),
trained_by="Cellpose"))
return out
#: bioimage.io's published collection.
BIOIMAGEIO_COLLECTION = ("https://hypha.aicell.io/bioimage-io/artifacts/"
"bioimage.io/children?limit=1000&pagination=false")
#: What spaCR's Cellpose can actually load. The test is COMPATIBILITY, not the
#: word "cellpose": spaCR runs Cellpose 4, so the SAM and DINO backbones load
#: and cyto3 and earlier do not -- they carry the same "cellpose" tag, load,
#: and then produce nonsense. The older collection.json does not list these
#: models at all, which is why reading it found nothing.
BIOIMAGEIO_COMPATIBLE = ("cellpose sam", "cpsam", "cellposedino",
"cellpose dino", "cpdino")
#: The words that make a Cellpose 4 model a Cellpose-DINO one, which runs
#: through the Cellpose-DINO backend: the Cellpose spaCR installs builds a
#: DINO network only when the ``dinov3`` package is there, and it is not.
_BIOIMAGEIO_DINO = ("cellposedino", "cellpose dino", "cpdino")
#: How long a fetched collection is trusted before it is fetched again.
BIOIMAGEIO_CACHE_HOURS = 24
#: The smallest weights file that can hold a Cellpose network. Cellpose 3's
#: is 26 MB and Cellpose-SAM's 1.2 GB. At least one bioimage.io Cellpose
#: package (OC1 Project 11, happy-elephant) publishes a 1596-byte stand-in
#: instead, and its own code downloads the real checkpoint from
#: GitHub when it runs; a row offering that file would download it, load it
#: and fail, or worse.
_BIOIMAGEIO_MIN_WEIGHTS = 1 << 20
#: How the note saying spaCR cannot run a bioimage.io row begins. The picker
#: finds it by this, shows it in place of Download and Use, and says it.
_CANNOT_RUN = "spaCR cannot run this model: "
#: The run settings a bioimage.io Cellpose package names, which its row
#: reports so a user can compare them with the ones spaCR will use.
_BIOIMAGEIO_RUN_KWARGS = ("diameter", "diam_mean", "flow_threshold",
"cellprob_threshold", "channels")
#: Where bioimage.io serves an artifact's files.
_BIOIMAGEIO_FILES = ("https://hypha.aicell.io/bioimage-io/artifacts/{alias}/"
"files/{name}")
#: What a Cellpose 3-format row says about how to use it.
_CELLPOSE3_USE = (
"runs through the Cellpose 3 backend exactly as a stock Cellpose 3 model "
"does: install that backend from this list, download this model, and "
"press Use this model, which puts cellpose3:<its path> in the object's "
"model setting")
#: What a Cellpose-SAM row says about how to use it.
_CELLPOSE4_USE = (
"runs through spaCR's own Cellpose 4: download it and press Use this "
"model, which puts its path in the object's model setting")
#: What a Cellpose-DINO row says about how to use it.
_CELLPOSE_DINO_USE = (
"runs through the Cellpose-DINO backend, Cellpose 4 with DINOv3 in an "
"environment of its own: install that backend from this list, download "
"this model, and press Use this model, which puts cellpose_dino:<its "
"path> in the object's model setting")
def _bioimageio_text(manifest: Mapping[str, Any]) -> str:
"""A manifest's name, description and tags, lower case, words spaced."""
parts = [str(manifest.get("name") or ""),
str(manifest.get("description") or "")]
parts += [str(t) for t in (manifest.get("tags") or ())]
return " ".join(parts).lower().replace("-", " ").replace("_", " ")
def _looks_like_cellpose_sam(manifest: Mapping[str, Any]) -> bool:
"""Whether a bioimage.io manifest is loadable by the Cellpose spaCR runs.
STRICTLY. spaCR loads these as drop-in replacements, and only the Cellpose
4 backbones work that way: a cyto3 or a ResNet model carries the same
"cellpose" tag, loads, and then produces nonsense.
"""
if str(manifest.get("type") or "") != "model":
return False
text = _bioimageio_text(manifest)
return any(key in text for key in BIOIMAGEIO_COMPATIBLE)
def _looks_like_cellpose_dino(manifest: Mapping[str, Any]) -> bool:
"""Whether a Cellpose 4 manifest is a Cellpose-DINO model."""
text = _bioimageio_text(manifest)
return any(key in text for key in _BIOIMAGEIO_DINO)
def _bioimageio_cache(name: str) -> Path:
"""Where a bioimage.io cache file lives: ``~/.spacr/<name>``."""
return _spacr_home() / name
def _bioimageio_payload(timeout: float, url: Optional[str],
allow_network: bool) -> Tuple[Any, bool]:
"""bioimage.io's collection: the cache while it is fresh, bioimage.io
when allowed, and the stale cache when bioimage.io cannot be reached.
A stale collection is still bioimage.io's own list, a day or more old,
and an empty category is the worse answer. :func:`catalogue` must not
touch the network -- it is called on offline paths and from the GUI
thread -- so the fetch happens in the picker's background warm-up and
this reads what that left behind.
:returns: ``(payload or None, whether it was fetched just now)``.
"""
import urllib.request
cache = _bioimageio_cache("bioimageio_children.json")
stale = None
try:
if cache.is_file():
stale = json.loads(cache.read_text(encoding="utf-8"))
if (time.time() - cache.stat().st_mtime
< BIOIMAGEIO_CACHE_HOURS * 3600):
return stale, False
except Exception:
stale = None
if not allow_network:
return stale, False
try:
with urllib.request.urlopen(url or BIOIMAGEIO_COLLECTION,
timeout=timeout) as response:
payload = json.loads(response.read().decode("utf-8"))
cache.parent.mkdir(parents=True, exist_ok=True)
cache.write_text(json.dumps(payload), encoding="utf-8")
return payload, True
except Exception:
return stale, False
def _bioimageio_sizes() -> Dict[str, int]:
"""The weights sizes measured so far, by URL; empty when none were."""
try:
sizes = json.loads(
_bioimageio_cache("bioimageio_weight_sizes.json").read_text(encoding="utf-8"))
except Exception:
return {}
if not isinstance(sizes, Mapping):
return {}
return {str(k): int(v) for k, v in sizes.items()
if isinstance(v, int) and not isinstance(v, bool)}
def _weights_size(uri: str, timeout: float) -> int:
"""The size of the file at ``uri``, asking for its first byte only.
bioimage.io answers HEAD with 405, and a ranged GET with the total in
``Content-Range``. A 1.2 GB checkpoint is measured without being read.
:returns: bytes, or 0 when the server does not say.
"""
import urllib.request
request = urllib.request.Request(uri, headers={"Range": "bytes=0-0"})
with urllib.request.urlopen(request, timeout=timeout) as response:
total = str(response.headers.get("Content-Range") or "").rpartition("/")[2]
if total.isdigit():
return int(total)
length = str(response.headers.get("Content-Length") or "")
return int(length) if length.isdigit() and response.status == 200 else 0
def _measure_bioimageio_weights(uris: Iterable[str], sizes: Dict[str, int],
timeout: float) -> Dict[str, int]:
"""Measure the weights files not measured yet, and remember them.
A failed measurement is kept as 0, so a server that will not say is asked
again only when the collection itself is fetched again.
"""
sizes = dict(sizes)
for uri in uris:
try:
sizes[uri] = _weights_size(uri, timeout)
except Exception:
sizes[uri] = 0
try:
cache = _bioimageio_cache("bioimageio_weight_sizes.json")
cache.parent.mkdir(parents=True, exist_ok=True)
cache.write_text(json.dumps(sizes), encoding="utf-8")
except Exception:
pass
return sizes
[docs]
def bioimageio_entries(timeout: float = 5.0,
url: Optional[str] = None,
allow_network: bool = False) -> List["ModelEntry"]:
"""Every Cellpose model bioimage.io publishes, as zoo rows.
Read from bioimage.io's own collection, so the category cannot go stale
the way a typed list would. Three kinds of row:
* a Cellpose-SAM or Cellpose-DINO model, which spaCR's own Cellpose 4
loads: a ``cellpose`` row that downloads its weights;
* a Cellpose 3-format checkpoint: a ``cellpose3`` row that downloads its
weights for the Cellpose 3 backend, which runs it as it runs cyto3;
* one of either that spaCR cannot run, which says why in its first note
(:func:`_bioimageio_cannot_run`) and offers nothing to download.
Each weights file is checked against the SHA-256 its manifest publishes,
and carries the licence the uploader chose and what it was trained on.
Best effort and never raises: no network, a slow mirror or a changed
schema all mean fewer rows, never a zoo that fails to open. With
``allow_network`` the collection is refreshed once it is a day old, and
the weights files are measured, so a stand-in file is refused before
anyone downloads it.
"""
payload, fetched = _bioimageio_payload(timeout, url, allow_network)
if payload is None:
return []
sizes = _bioimageio_sizes()
rows = _bioimageio_rows(payload, sizes)
if allow_network:
unmeasured = [r.uri for r in rows if r.uri and (
r.uri not in sizes or (fetched and not sizes[r.uri]))]
if unmeasured:
sizes = _measure_bioimageio_weights(unmeasured, sizes, timeout)
rows = _bioimageio_rows(payload, sizes)
return rows
def _bioimageio_rows(payload: Any, sizes: Mapping[str, int]) -> List["ModelEntry"]:
"""The rows a bioimage.io collection makes, in its order."""
if isinstance(payload, list):
items = payload
elif isinstance(payload, Mapping):
items = payload.get("items", [])
else:
items = []
out = []
for item in items if isinstance(items, list) else []:
row = _bioimageio_row(item, sizes)
if row is not None:
out.append(row)
return out
def _bioimageio_row(item: Any, sizes: Mapping[str, int]) -> Optional["ModelEntry"]:
"""One collection item as a row, or None when it is no Cellpose model."""
if not isinstance(item, Mapping):
return None
manifest = item.get("manifest") or {}
if not isinstance(manifest, Mapping):
return None
if str(manifest.get("type") or "") != "model":
return None
alias = str(item.get("alias") or "")
if not alias:
return None
cellpose3 = _cellpose3_weights(manifest)
cellpose4 = cellpose3 is None and _looks_like_cellpose_sam(manifest)
if (cellpose3 is None and not cellpose4
and "cellpose" not in _bioimageio_text(manifest).replace(" ", "")):
return None
title = str(manifest.get("name") or alias)
slug = re.sub(r"[^a-z0-9]+", "_", title.lower()).strip("_") or alias
authors = ", ".join(
str(a.get("name")) for a in (manifest.get("authors") or ())
if isinstance(a, Mapping))
if cellpose3 is not None:
row = _cellpose3_download_entry(
alias, slug, title, manifest, authors, *cellpose3)
else:
row = _cellpose4_download_entry(alias, slug, title, manifest, authors)
size = int(sizes.get(row.uri) or 0) if row.uri else 0
if size:
row = replace(row, size_bytes=size)
reason = _bioimageio_refusal(manifest, row, size,
known=cellpose3 is not None or cellpose4)
if reason:
row = replace(row, uri="", notes=(
_CANNOT_RUN + reason, f"bioimage.io model {alias}"))
return row
def _bioimageio_refusal(manifest: Mapping[str, Any], row: "ModelEntry",
size: int, *, known: bool) -> str:
"""Why spaCR cannot run a bioimage.io row, or ``''`` when it can.
:param known: whether the manifest names a Cellpose 3 or a Cellpose 4
network at all.
:param size: the measured size of its weights file, 0 when unmeasured.
"""
if not known:
called = _bioimageio_architecture(manifest) or "no network named"
return (f"its package runs a network of its own ({called}), neither "
f"Cellpose 3's nor Cellpose-SAM's, and those are the two "
f"Cellpose networks spaCR runs.")
if not row.uri:
return "its package names no weights file to download."
if 0 < size < _BIOIMAGEIO_MIN_WEIGHTS:
return (f"the weights file bioimage.io publishes for it is a "
f"{size}-byte stand-in, not a Cellpose network. The package's "
f"own code downloads the real checkpoint from elsewhere when "
f"it runs, and spaCR runs only the weights bioimage.io "
f"publishes and checksums.")
return ""
def _bioimageio_cannot_run(entry: Any) -> str:
"""The sentence saying spaCR cannot run this row, or ``''``.
:param entry: any zoo row; only a bioimage.io row can carry one.
"""
for note in getattr(entry, "notes", ()) or ():
if str(note).startswith(_CANNOT_RUN):
return str(note)
return ""
def _bioimageio_architecture(manifest: Mapping[str, Any]) -> str:
"""The network a manifest's PyTorch weights name, or ``''``."""
weights = manifest.get("weights")
state = weights.get("pytorch_state_dict") if isinstance(weights, Mapping) else None
architecture = state.get("architecture") if isinstance(state, Mapping) else None
return (str(architecture.get("callable") or "")
if isinstance(architecture, Mapping) else "")
def _bioimageio_trained_on(title: str, manifest: Mapping[str, Any]) -> str:
"""What a bioimage.io model was trained on, as its package says.
``training_data`` comes first when the package names it, so the cap on
the cell cuts the description rather than the data.
"""
data = manifest.get("training_data")
ident = str((data.get("id") if isinstance(data, Mapping) else data) or "")
if ident.startswith("10."):
ident = f"https://doi.org/{ident}"
described = str(manifest.get("description") or "").strip()
text = title
if ident:
text += f" — trained on {ident}."
if described:
text += f" — {described}" if not ident else f" {described}"
return text[:300]
def _bioimageio_run_settings(manifest: Mapping[str, Any]) -> Tuple[str, ...]:
"""What the package itself runs the model with, as a note, or nothing.
spaCR runs a bioimage.io model with the object's Mask generation settings,
as it runs every other model, so the two can differ; the note says so.
"""
weights = manifest.get("weights")
state = weights.get("pytorch_state_dict") if isinstance(weights, Mapping) else None
architecture = state.get("architecture") if isinstance(state, Mapping) else None
kwargs = architecture.get("kwargs") if isinstance(architecture, Mapping) else None
if not isinstance(kwargs, Mapping):
return ()
named = [f"{key}={kwargs[key]}" for key in _BIOIMAGEIO_RUN_KWARGS
if kwargs.get(key) is not None]
if not named:
return ()
return (f"Its bioimage.io package runs it with {', '.join(named)}; spaCR "
f"runs it with the object's Mask generation settings instead, as "
f"it runs every other model.",)
#: The architectures a Cellpose 3-format bioimage.io model names.
#: ``CellPoseWrapper`` wraps a Cellpose 3 checkpoint and ``CPnetBioImageIO``
#: is Cellpose's own export of one; either way the ``pytorch_state_dict``
#: weights ARE the checkpoint, which Cellpose 3's ``CellposeModel`` loads.
_BIOIMAGEIO_CELLPOSE3_ARCHITECTURES = ("CellPoseWrapper", "CPnetBioImageIO")
def _cellpose3_weights(manifest: Mapping[str, Any]) -> Optional[Tuple[str, str]]:
"""The Cellpose 3 checkpoint a bioimage.io manifest publishes.
Only the two architectures in :data:`_BIOIMAGEIO_CELLPOSE3_ARCHITECTURES`
count -- the word "cellpose" in a tag does not, and a Cellpose-SAM model
is the Cellpose 4 kind, listed as such.
:returns: ``(source, sha256)`` of the weights file, or None.
"""
if str(manifest.get("type") or "") != "model":
return None
if _looks_like_cellpose_sam(manifest):
return None
weights = manifest.get("weights")
state = weights.get("pytorch_state_dict") if isinstance(weights, Mapping) else None
if not isinstance(state, Mapping):
return None
called = _bioimageio_architecture(manifest)
source = str(state.get("source") or "").strip()
if called not in _BIOIMAGEIO_CELLPOSE3_ARCHITECTURES or not source:
return None
return source, str(state.get("sha256") or "").strip().lower()
def _bioimageio_uri(alias: str, source: str) -> str:
"""Where a manifest's weights file is downloaded from."""
if not source:
return ""
if source.startswith(("http://", "https://")):
return source
return _BIOIMAGEIO_FILES.format(alias=alias, name=source)
def _cellpose3_download_entry(alias: str, slug: str, title: str,
manifest: Mapping[str, Any], authors: str,
source: str, sha256: str) -> "ModelEntry":
"""A bioimage.io Cellpose 3 checkpoint as a row that downloads it.
The weights file is fetched from bioimage.io itself and checked against
the SHA-256 its manifest publishes, like any other zoo download, and the
licence the uploader chose travels with the row.
"""
suffix = Path(source).suffix if Path(source).suffix in (".pth", ".pt") else ".pth"
return ModelEntry(
key=f"bioimageio_{slug}", name=f"{slug}{suffix}", path="",
kind="cellpose3", source="bioimage.io",
uri=_bioimageio_uri(alias, source), sha256=sha256,
trained_on=_bioimageio_trained_on(title, manifest),
trained_by=authors or "bioimage.io",
licence=str(manifest.get("license") or ""),
notes=(f"bioimage.io model {alias}; {_CELLPOSE3_USE}",)
+ _bioimageio_run_settings(manifest))
def _cellpose4_download_entry(alias: str, slug: str, title: str,
manifest: Mapping[str, Any],
authors: str) -> "ModelEntry":
"""A bioimage.io Cellpose-SAM or Cellpose-DINO model as a row that
downloads it: a Cellpose-SAM one for spaCR's own Cellpose 4 to load from
its path, a Cellpose-DINO one (kind ``cellpose_dino``) for the
Cellpose-DINO backend."""
weights = manifest.get("weights")
state = weights.get("pytorch_state_dict") if isinstance(weights, Mapping) else None
state = state if isinstance(state, Mapping) else {}
source = str(state.get("source") or "").strip()
suffix = Path(source).suffix if Path(source).suffix in (".pth", ".pt") else ""
dino = _looks_like_cellpose_dino(manifest)
return ModelEntry(
key=f"bioimageio_{slug}", name=f"{slug}{suffix}", path="",
kind="cellpose_dino" if dino else "cellpose", source="bioimage.io",
uri=_bioimageio_uri(alias, source),
sha256=str(state.get("sha256") or "").strip().lower(),
trained_on=_bioimageio_trained_on(title, manifest),
trained_by=authors or "bioimage.io",
licence=str(manifest.get("license") or ""),
notes=(f"bioimage.io model {alias}; "
f"{_CELLPOSE_DINO_USE if dino else _CELLPOSE4_USE}",)
+ _bioimageio_run_settings(manifest))
#: What each Cellpose 3 model is, for its zoo row.
_CELLPOSE3_NOTES = {
"cyto3": "Cellpose 3's generalist whole-cell model.",
"cyto2": "Cellpose 2's whole-cell model.",
"cyto": "The original Cellpose whole-cell model.",
"nuclei": "Cellpose's nucleus model.",
}
def _cellpose3_model_entries() -> List["ModelEntry"]:
"""The Cellpose 3 backend's own models, listed whether or not it is here.
The name IS the path, as for the Cellpose 4 stock models: ``cyto3`` in
an object's model setting, with segmentation_backend set to cellpose3,
is what runs it. Until the backend is installed the row has no path and
says what it needs.
"""
from ._segmentation_backends import _CELLPOSE3, _SPECS, _backend_state
spec = _SPECS[_CELLPOSE3]
ready = _backend_state(_CELLPOSE3).ready
out = []
for model in spec.models:
out.append(ModelEntry(
key=f"cellpose3_{model}", name=model, kind="cellpose3",
source="stock", path=model if ready else "",
uri=f"backend:{_CELLPOSE3}", sha256="", size_bytes=0,
trained_on=(f"{_CELLPOSE3_NOTES.get(model, '')} Runs through the "
f"Cellpose 3 backend.").strip(),
trained_by="Cellpose", licence=spec.licence,
notes=() if ready else (
"needs the Cellpose 3 backend, which installs from this "
"list into an environment of its own",)))
return out
def _mask_model_kinds() -> Tuple[str, ...]:
"""The kinds an object's model setting in Mask generation can run:
Cellpose-SAM, Cellpose 3, Cellpose-DINO and each prefixed backend's
models. What every Model zoo button beside such a setting asks for."""
return ("cellpose", "cellpose3", "cellpose_dino") + PREFIXED_KINDS
def _prefixed_model_entries() -> List["ModelEntry"]:
"""The models of each prefixed backend (StarDist, InstanSeg, Omnipose),
listed whether or not the backend is here.
As for Cellpose 3's own models the name is the path: Use this model
writes ``<prefix><name>`` into the object's model setting. Until the
backend is installed the row has no path, says what it needs, and
choosing it offers the install.
"""
from ._segmentation_backends import _backend_state
out = []
for name in PREFIXED_KINDS:
spec = _BACKEND_SPECS[name]
ready = _backend_state(name).ready
for model in spec.models:
out.append(ModelEntry(
key=f"{name}_{model}", name=model, kind=name,
source="stock", path=model if ready else "",
uri=f"backend:{name}", sha256="", size_bytes=0,
trained_on=(f"{spec.label}'s own {model} model. Runs "
f"through the {spec.label} backend."),
trained_by=spec.label, licence=spec.licence,
notes=() if ready else (
f"needs the {spec.label} backend, which installs from "
f"this list into an environment of its own",)))
return out
def _backend_for(entry: Any) -> str:
"""The optional segmentation backend a zoo row needs, or ``''``.
A backend row names itself in its ``backend:<name>`` uri; every
``cellpose3`` model needs the Cellpose 3 backend, and every
``cellpose_dino`` model the Cellpose-DINO one.
"""
uri = str(getattr(entry, "uri", "") or "")
if uri.startswith("backend:"):
return uri.split(":", 1)[1]
kind = getattr(entry, "kind", "")
if kind in ("cellpose3", "cellpose_dino") + PREFIXED_KINDS:
return kind
return ""
#: ``name -> (label, install uri, import name, what it is)`` for every
#: optional ``segmentation_backend``. These are PACKAGES, not checkpoints: the
#: zoo lists them so a user learns they exist, and each installs into an
#: environment of its own, never into spaCR's. Cellpose-DINO segments but is
#: no ``segmentation_backend`` value -- a model setting chooses it -- so the
#: backend box that reads this does not offer it.
INSTALLABLE_BACKENDS = {
_name: (_spec.label, f"backend:{_name}", _spec.module, _spec.blurb)
for _name, _spec in _BACKEND_SPECS.items()
if _spec.segments and _name in _SEGMENTATION_BACKEND_NAMES
}
[docs]
def installable_backend_entries() -> List["ModelEntry"]:
"""Every optional segmentation backend, in whatever state it is here.
Each backend is listed whether or not it is installed, and its
``source`` says where it stands -- ``installed``,
``installable``, ``installing`` or ``not installable here`` -- with the
reason as its first note and its licence on the row. Installing one
builds it an environment of its own under ``~/.spacr/backends`` and
leaves spaCR's own environment alone.
"""
from ._segmentation_backends import _backend_state, _credential_note
out = []
for name, spec in _BACKEND_SPECS.items():
state = _backend_state(name)
out.append(ModelEntry(
key=f"{name}_v1", name=spec.label,
path=state.env if state.ready and not state.in_process else "",
kind="backend", source=state.state, uri=f"backend:{name}",
sha256="", size_bytes=0, trained_on=spec.blurb,
trained_by=spec.label, licence=spec.licence,
notes=tuple(note for note in (f"{state.state}: {state.reason}",
spec.licence_note, spec.published,
_credential_note(name))
if note)))
return out
#: Where spaCR's Add button puts community submissions.
COMMUNITY_REPO = "einarolafsson/user-models"
#: Submissions live under this prefix until somebody promotes them.
COMMUNITY_PREFIX = "staging/"
#: The one-line warning that must travel with every community row.
COMMUNITY_WARNING = (
"Community upload — NOT vetted. Anyone can submit through spaCR's Add "
"button; nobody has checked what this file is, what it was trained on, or "
"whether its reported scores are real.")
#: The five places a zoo row can come from, in the order they are offered.
#:
#: WHY FIVE HEADINGS RATHER THAN ONE LIST AND A BOOLEAN. The zoo used to be a
#: single list with "show unvetted community uploads" beside it. Ten models
#: trained here, four stock Cellpose-SAM backbones, bioimage.io's collection,
#: the Cellpose 3 backend's own four and a shared catalogue is more than a
#: list can carry, and a boolean that means "also show these" cannot fold away
#: the four groups a reader is not looking at.
#:
#: These names are IDENTIFIERS as well
#: as captions: :func:`source_of` returns one of them, and the picker
#: remembers which are on by this spelling. Renaming one is a migration.
ZOO_SOURCES: Tuple[str, ...] = (
"cellposeSAM", "spaCR", "spaCR community", "bioimage.io", "cellpose3",
"Third-party")
#: The headings a user who has never touched them sees turned on.
#:
#: The stock weights and the models trained here: the two a new user can act
#: on immediately. The other three are opt-in -- community uploads because
#: nobody has vetted them, bioimage.io and cellpose3 because they are long
#: lists of other people's models that would bury the ten this project ships.
DEFAULT_ZOO_SOURCES: Tuple[str, ...] = ("cellposeSAM", "spaCR")
#: ``ModelEntry.source`` values that mean "spaCR's own catalogue".
#:
#: ``remote`` is :data:`BUNDLED_REMOTE_MODELS` and any JSON catalogue named by
#: :data:`CATALOGUE_ENV_VAR`; ``bundled`` is what shipped in the package;
#: ``local`` is a checkpoint found on this machine or listed through the Add
#: button. All three are models this project offers, so all three are spaCR.
_SPACR_OWN_SOURCES = ("remote", "bundled", "local")
#: ``ModelEntry.source`` values that mean an unvetted upload.
#:
#: ``shared`` comes from :func:`shared_catalogue`, ``community`` from
#: :func:`community_entries`; the two are different transports for the same
#: thing, which is a file somebody uploaded that nobody has checked.
_COMMUNITY_SOURCES = ("shared", "community")
#: Hugging Face accounts whose repositories hold spaCR's own models.
#:
#: A shared or community row whose download URL lives in one of these
#: accounts is a model the maintainer trained and uploaded, and is listed
#: under spaCR whatever its key. :data:`COMMUNITY_REPO` is the one exception:
#: it belongs to the same account but is the inbox the Add button fills, so
#: what is in it was uploaded by somebody else.
SPACR_HF_OWNERS: Tuple[str, ...] = (HF_MODELS_REPO.split("/", 1)[0],)
_HF_REPO_PART = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$")
[docs]
def hf_repo_of(uri: Any) -> str:
"""The ``owner/repo`` a Hugging Face download URL points into, or ``''``.
Read from the URL a row is FETCHED from, never from a ``trained_by`` or
``owner`` field a catalogue row can claim for itself: the bytes come from
this repository, so it is the only owner that cannot be borrowed.
Strict on purpose. Only ``https://huggingface.co/<owner>/<repo>/...`` and
``https://huggingface.co/datasets/<owner>/<repo>/...`` qualify; another
host, a port, user info, ``..`` or percent-escapes in the path give
``''``, so ``huggingface.co.example.org/einarolafsson/...`` is nobody's.
:param uri: a download URL.
:returns: ``owner/repo`` or ``''``.
"""
from urllib.parse import urlsplit
try:
parts = urlsplit(str(uri or "").strip())
except ValueError:
return ""
if parts.scheme != "https" or parts.netloc.lower() != "huggingface.co":
return ""
path = parts.path
if ".." in path or "%" in path or "\\" in path:
return ""
segments = path.split("/")[1:]
if segments[:1] == ["datasets"]:
segments = segments[1:]
if len(segments) < 2:
return ""
owner, repo = segments[0], segments[1]
if not (_HF_REPO_PART.match(owner) and _HF_REPO_PART.match(repo)):
return ""
return f"{owner}/{repo}"
[docs]
def is_spacr_owned(entry: Any) -> bool:
"""Whether a row's file lives in a repository owned by the einarolafsson account.
:param entry: any zoo row.
:returns: True when :func:`hf_repo_of` its ``uri`` is in an account named
by :data:`SPACR_HF_OWNERS` and is not :data:`COMMUNITY_REPO`.
"""
repo = hf_repo_of(getattr(entry, "uri", "")).lower()
if not repo or repo == COMMUNITY_REPO.lower():
return False
owners = {owner.lower() for owner in SPACR_HF_OWNERS}
return repo.split("/", 1)[0] in owners
def _is_spacr_published(entry: Any) -> bool:
"""Whether a shared or community row is a checkpoint spaCR publishes.
Item 685: the Toxoplasma plaque model also travels in the shared
catalogue, where it was filed under spaCR community. A row is spaCR's
own when its key AND its checksum match a
:data:`BUNDLED_REMOTE_MODELS` record, so an upload that merely borrows
a key stays a community row. :func:`is_spacr_owned` is the other way
in: a file in a repository owned by the einarolafsson account.
:param entry: any zoo row.
:returns: True for the very file a bundled record names.
"""
key = str(getattr(entry, "key", "") or "")
sha = str(getattr(entry, "sha256", "") or "").strip().lower()
if not key or not sha:
return False
return any(str(record.get("key") or "") == key
and str(record.get("sha256") or "").strip().lower() == sha
for record in BUNDLED_REMOTE_MODELS)
def _record_is_alternative(record: Mapping[str, Any]) -> bool:
"""Whether a bundled record marks itself as an alternative.
An explicit ``alternative`` field wins; otherwise the record's
``display_name`` saying "alternative" marks it, so moving a model's
default status means editing its own record and nothing else.
"""
if "alternative" in record:
return bool(record.get("alternative"))
return "alternative" in str(record.get("display_name") or "").lower()
#: Keys of published models that are offered beside a default rather than
#: instead of it. They are listed under spaCR like any other, but a model
#: family's row opens on its default version, not on the alternative.
ALTERNATIVE_MODEL_KEYS: frozenset = frozenset(
str(record.get("key") or "") for record in BUNDLED_REMOTE_MODELS
if _record_is_alternative(record))
THIRD_PARTY_MODEL_KEYS: frozenset = frozenset(
str(record.get("key") or "") for record in BUNDLED_REMOTE_MODELS
if record.get("third_party"))
[docs]
def is_third_party(entry: Any) -> bool:
"""Whether a row is a published model that spaCR did not train.
Item 691. Such a row is offered by spaCR, with its download and checksum
in :data:`BUNDLED_REMOTE_MODELS`, but its weights belong to someone else
and carry their own licence and citation, so it is listed under
``Third-party`` rather than ``spaCR``.
:param entry: any zoo row.
:returns: True for a key in :data:`THIRD_PARTY_MODEL_KEYS` or a row of
kind ``mitonet``.
"""
return (str(getattr(entry, "key", "") or "") in THIRD_PARTY_MODEL_KEYS
or str(getattr(entry, "kind", "") or "") == "mitonet")
[docs]
def is_alternative(entry: Any) -> bool:
"""Whether a row is a published alternative rather than a default.
:param entry: any zoo row.
:returns: True for a key in :data:`ALTERNATIVE_MODEL_KEYS`.
"""
return str(getattr(entry, "key", "") or "") in ALTERNATIVE_MODEL_KEYS
def _record_is_candidate(record: Mapping[str, Any]) -> bool:
"""Whether a bundled record marks itself as an unpromoted candidate."""
if "candidate" in record:
return bool(record.get("candidate"))
return bool(re.search(r"\bcandidate\b",
str(record.get("display_name") or "").lower()))
_CANDIDATE_MODEL_KEYS: frozenset = frozenset(
str(record.get("key") or "") for record in BUNDLED_REMOTE_MODELS
if _record_is_candidate(record))
[docs]
def default_version_index(pairs: Sequence[Tuple[str, Any]]) -> int:
"""Which version a model family's row opens on.
The newest version that is neither an alternative nor a candidate, so
publishing either never changes what a family's row hands out by
default. The README and the website's model cards use the same rule.
:param pairs: ``(label, entry)`` pairs, newest first.
:returns: an index into ``pairs``; 0 when every version is an
alternative or a candidate, or there are none.
"""
for index, (_label, entry) in enumerate(pairs):
if not is_alternative(entry) and str(
getattr(entry, "key", "") or "") not in _CANDIDATE_MODEL_KEYS:
return index
return 0
[docs]
def source_of(entry: Any) -> str:
"""Which of :data:`ZOO_SOURCES` this row belongs under.
Decided from the row's own ``source``, ``kind`` and ``uri``, never from a
separate list of names: such a list goes stale the first time a model is
added, and a model missing from it appears in the catalogue under no
heading, where nobody can see it.
The order of the tests is the rule, and it matters in one place: a
Cellpose 3 checkpoint published on bioimage.io is a bioimage.io row, not
a Cellpose 3 one. ``cellpose3`` means the backend's own models -- cyto,
cyto2, cyto3, nuclei -- and the backend package that runs them.
A shared or community row is spaCR's when it is the file a bundled
record names (key and checksum) or when it is fetched from a repository
owned by the einarolafsson account (:func:`is_spacr_owned`). Anything else
uploaded by somebody else stays under ``spaCR community``.
A row that matches nothing is filed under ``spaCR`` and a warning is
logged. It is the fallback rather than a sixth heading because a model
under the wrong heading is still visible, while a model under no heading
appears to the user as a missing model.
:param entry: any zoo row -- a :class:`ModelEntry`, or anything with
``kind``, ``source`` and ``uri``.
:returns: one of :data:`ZOO_SOURCES`.
"""
source = str(getattr(entry, "source", "") or "")
kind = str(getattr(entry, "kind", "") or "")
uri = str(getattr(entry, "uri", "") or "")
if str(getattr(entry, "key", "") or "") in THIRD_PARTY_MODEL_KEYS and (
source not in _COMMUNITY_SOURCES or _is_spacr_published(entry)):
return "Third-party"
if source in _COMMUNITY_SOURCES:
if _is_spacr_published(entry) or is_spacr_owned(entry):
return "spaCR"
return "spaCR community"
if source == "bioimage.io" or "bioimage" in uri.lower():
return "bioimage.io"
if kind == "cellpose3" or _backend_for(entry) == "cellpose3":
return "cellpose3"
if kind == "cellpose" and source == "stock":
return "cellposeSAM"
if is_third_party(entry):
return "Third-party"
if (kind == "backend" or kind in PREFIXED_KINDS
or source in _SPACR_OWN_SOURCES):
return "spaCR"
LOG.warning(
"model zoo: %r (kind=%r, source=%r) matches no source heading; "
"listing it under spaCR so it stays visible",
getattr(entry, "name", "") or getattr(entry, "key", ""), kind, source)
return "spaCR"
[docs]
def group_by_source(entries: Iterable[Any]) -> Dict[str, List[Any]]:
"""Split a listing into :data:`ZOO_SOURCES`, keeping each source's order.
Every heading is present even when it has no rows, so a caller drawing
the strip does not have to know which of the five happened to be empty
this time.
:param entries: the rows to split.
:returns: heading -> rows, in :data:`ZOO_SOURCES` order.
"""
out: Dict[str, List[Any]] = {name: [] for name in ZOO_SOURCES}
for entry in entries:
out[source_of(entry)].append(entry)
return out
[docs]
def entries_from_sources(entries: Iterable[Any],
sources: Iterable[str]) -> List[Any]:
"""The rows belonging to the headings that are on, in the given order.
:param entries: the rows to filter.
:param sources: the headings currently on.
"""
wanted = set(sources)
return [entry for entry in entries if source_of(entry) in wanted]
[docs]
def catalogue(include_bundled: bool = True, remote: bool = True,
catalogue_path: Any = None,
include_plugins: bool = True,
block: Optional[bool] = None) -> List[ModelEntry]:
"""Everything the zoo knows about without scanning the user's disks.
That is: the models bundled with the installed package (whatever
:func:`spacr.utils.download_models` has put in ``resources/models``), plus
the declared remote entries — :data:`BUNDLED_REMOTE_MODELS` and, if one is
configured, the JSON catalogue named by ``catalogue_path`` or the
:data:`CATALOGUE_ENV_VAR` environment variable.
Local apart from one thing:
with ``remote`` on, this also asks :func:`shared_catalogue` for the
community rows, which is a network call. It works offline either way --
that fetch never raises -- and it never blocks Qt's GUI thread, which
:func:`shared_catalogue` enforces for itself.
:param include_bundled: list the models in the package resources folder.
:param remote: list declared remote entries.
:param catalogue_path: a JSON catalogue to add; defaults to
``$SPACR_MODEL_CATALOGUE`` when that names a file.
:param include_plugins: include entries returned by installed spaCR model
providers. Provider failures are recorded in plugin diagnostics and do
not hide built-in entries.
:param block: passed to :func:`shared_catalogue`. ``False`` takes the
community rows from its cache rather than waiting for the network;
``None`` lets that function decide from the thread it is on.
:returns: bundled entries first, then remote ones already present locally
are dropped (a downloaded model is listed once, as the local file).
"""
entries: List[ModelEntry] = []
if include_bundled:
root = package_model_root()
if root.is_dir():
entries.extend(discover_local(root, max_depth=2))
entries.extend(stock_cellpose_entries())
entries.extend(installable_backend_entries())
entries.extend(_cellpose3_model_entries())
entries.extend(_prefixed_model_entries())
if remote:
entries.extend(bioimageio_entries())
if remote:
have = {(e.key, e.name) for e in entries}
for record in BUNDLED_REMOTE_MODELS:
entry = _builtin_entry_from_mapping(record)
if (entry.key, entry.name) not in have:
entries.append(entry)
have.add((entry.key, entry.name))
path = catalogue_path or os.environ.get(CATALOGUE_ENV_VAR, "")
if path and os.path.isfile(str(path)):
for entry in load_catalogue_file(path):
if (entry.key, entry.name) not in have:
entries.append(entry)
have.add((entry.key, entry.name))
for entry in shared_catalogue(block=block):
if (entry.key, entry.name) not in have:
entries.append(entry)
have.add((entry.key, entry.name))
entries = [e for e in entries if e.name not in RETIRED_MODEL_NAMES]
if include_plugins:
try:
from .plugins import (
load_object,
model_providers,
record_diagnostic,
)
have = {(entry.key, entry.name) for entry in entries}
for plugin_name, contribution in model_providers():
try:
provider = load_object(contribution.provider)
if not callable(provider):
raise TypeError(
f"{contribution.provider!r} is not callable"
)
produced = provider()
if isinstance(produced, (ModelEntry, Mapping)):
produced = (produced,)
for item in produced or ():
entry = (
item if isinstance(item, ModelEntry)
else _entry_from_mapping(item)
)
identity = (entry.key, entry.name)
if identity not in have:
entries.append(entry)
have.add(identity)
except Exception as exc:
record_diagnostic(
plugin_name,
f"Model provider {contribution.key!r} failed",
exc,
)
except Exception:
LOG.exception("Could not initialise plugin model providers")
return entries
[docs]
def resolve(key_or_path: Any,
entries: Optional[Sequence[ModelEntry]] = None) -> ModelEntry:
"""Turn a key, a name or a path into a :class:`ModelEntry`.
A path that exists wins over a key: pointing the zoo at a file you just
trained has to work without registering it anywhere first.
:param key_or_path: an entry key, a model filename, or a path to a file.
:param entries: the listing to search; defaults to :func:`catalogue`.
:returns: the entry.
:raises ModelUnreadable: when it looks like a path and no file is there.
:raises ModelZooError: when no entry matches, listing the near misses.
"""
text = str(key_or_path or "").strip()
if not text:
raise ModelZooError("no model given")
if os.path.isfile(text):
return entry_from_file(text)
if os.sep in text or text.startswith("~"):
raise ModelUnreadable(
f"no such model file: {text} — and nothing in the zoo is called "
f"that either")
pool = list(entries) if entries is not None else catalogue()
for entry in pool:
if entry.key == text or entry.name == text:
return entry
lowered = text.lower()
near = [e.key for e in pool if lowered in e.key.lower()
or lowered in e.name.lower()]
if len(near) == 1:
return next(e for e in pool if e.key == near[0])
raise ModelZooError(
f"no model called {text!r} in the zoo"
+ (f" — did you mean one of: {', '.join(near[:5])}?" if near
else f" ({len(pool)} entries known; pass a path to use a file "
f"that is not registered)"))
def _model_input_kinds(module, settings):
"""Describe file-model consumers without loading a model or inspecting paths."""
key = getattr(module, "key", str(module))
result = {}
if key in {"mask", "timelapse", "make_masks"}:
from .object_roles import SEGMENTED_ROLES
from ._segmentation_backends import (_backend_name, _cellpose3_choice,
_cellpose_dino_choice, _prefixed_backend)
for name in [*(f"{role}_model_name" for role in SEGMENTED_ROLES), "pathogen_model"]:
if name not in settings:
continue
value = settings.get(name)
backend = (_prefixed_backend(value)
or ("cellpose3" if _cellpose3_choice(value) is not None else None)
or ("cellpose_dino" if _cellpose_dino_choice(value) is not None else None)
or _backend_name(settings.get("segmentation_backend")))
result[name] = (backend,) if backend in KINDS else ()
result['real_object_classifier'] = ('classifier',)
elif key == "train_cellpose":
result = {"base_model": ("cellpose",)}
elif key == "cellpose_masks":
result = {"model_name": ("cellpose",), "custom_model": ("cellpose",)}
elif key == "analyze_plaques":
result = {"plaque_model": ("cellpose",), "well_detection": ("detector",),
"figure_detector": ("detector",)}
elif key in {"classify", "classify_merged", "train_only", "activation"}:
if key == "classify_merged":
from .classify import resolve_family
if resolve_family(settings) != "cv":
return {}
result = {name: ("classifier",) for name in
("model_path", "custom_model_path", "resume_checkpoint")}
elif key == "ops":
result = {"cellpose_model": ("cellpose",)}
elif key == "spatial_masks":
result = {"spatial_cellpose_model": ("cellpose",)}
return {name: kinds for name, kinds in result.items() if name in settings}
def _capture_model_metadata(module, settings):
"""Pin explicit trusted catalogue selections without file or network I/O."""
from copy import deepcopy
kinds = _model_input_kinds(module, settings)
metadata = _validate_model_metadata(settings, model_kinds=kinds)
result = deepcopy(settings)
records = metadata["settings"] if metadata is not None else {}
for name, accepted in kinds.items():
requested = settings[name]
if name in records or not isinstance(requested, str) or not requested.strip():
continue
if "/" in requested or "\\" in requested or requested.startswith("~"):
continue
matches = []
for raw in BUNDLED_REMOTE_MODELS:
if requested != raw.get("key") or raw.get("kind", "cellpose") not in accepted:
continue
pin = _model_reference_from_record(requested, raw)
if pin not in matches:
matches.append(pin)
if len(matches) > 1:
_bad_model_metadata(name, "one exact catalogue identity")
if matches:
records[name] = matches[0]
if records or metadata is not None:
result["model_metadata"] = {"schema_version": 1, "settings": records}
return result
_MODEL_REFERENCE_VERIFIED = object()
_MODEL_REFERENCE_STAT_FIELDS = ("st_dev", "st_ino", "st_size", "st_mtime_ns", "st_ctime_ns")
@dataclass(frozen=True)
class _VerifiedModelReference:
"""Immutable background verification result, retaining no open descriptor.
:param record: immutable existing-schema key/value pairs.
:param path: absolute local file whose bytes were verified.
:param identity: device, inode, size and nanosecond modification/change times.
:param _verification: private verifier marker preventing accidental raw-pin
attachment; this is not a same-process adversary or authentication boundary.
"""
record: Tuple[Tuple[str, Any], ...]
path: str
identity: Tuple[int, ...]
_verification: Any = _dc_field(default=None, repr=False, compare=False)
def _verified_model_reference(entry, path, *, requested_value=None, cancelled=None):
"""Verify a selected trusted file in a worker and preserve its exact request.
:param entry: selected ModelEntry; custom, stock or untrusted entries remain
unpinned rather than being identified from their filenames.
:param path: actual local file, read in bounded chunks on one held descriptor.
:param requested_value: exact setting value, normally ``str(path)``; an
existing matching backend prefix may wrap that same path.
:param cancelled: optional cooperative callback checked between reads.
:returns: immutable verified result, or None for an untrusted selection.
:raises ModelZooError: for conflicting identity, unreadable/changed bytes or
a request which does not name the selected file.
:raises DownloadCancelled: when cancellation is requested; no result is issued.
"""
import stat
if not isinstance(entry, ModelEntry) or entry.source == "stock":
return None
matches = [raw for raw in BUNDLED_REMOTE_MODELS
if raw.get("key", raw.get("name")) == entry.key
and raw.get("name") == entry.name
and raw.get("kind", "cellpose") == entry.kind
and raw.get("sha256") == entry.sha256
and re.fullmatch(r"[0-9a-f]{64}", str(raw.get("sha256") or ""))]
if not matches:
return None
local_value = os.fspath(path)
requested = local_value if requested_value is None else requested_value
if requested != local_value:
from ._segmentation_backends import (
_cellpose3_choice, _cellpose_dino_choice, _prefixed_choice,
)
selected = (_cellpose3_choice(requested) if entry.kind == "cellpose3"
else _cellpose_dino_choice(requested) if entry.kind == "cellpose_dino"
else _prefixed_choice(entry.kind, requested)
if entry.kind in PREFIXED_KINDS else None)
if selected != local_value:
_bad_model_metadata("requested_value", str(local_value))
record = _model_reference_from_record(requested, matches[0])
_match_model_reference(record)
actual_path = Path(path).expanduser().absolute()
expected = record["sha256"]
try:
with os.fdopen(os.open(actual_path, os.O_RDONLY | getattr(os, "O_NONBLOCK", 0)), "rb") as handle:
before = os.fstat(handle.fileno())
if not stat.S_ISREG(before.st_mode) or before.st_size <= 0:
_bad_model_metadata(str(path), f"sha256={expected}")
digest = hashlib.sha256()
total = 0
while True:
if cancelled is not None and cancelled():
raise DownloadCancelled()
chunk = handle.read(min(1024 * 1024, before.st_size - total + 1))
if not chunk:
break
total += len(chunk)
if total > before.st_size:
_bad_model_metadata(str(path), f"sha256={expected}")
digest.update(chunk)
after = os.fstat(handle.fileno())
current = actual_path.stat()
identity = tuple(getattr(before, key) for key in _MODEL_REFERENCE_STAT_FIELDS)
if (total != before.st_size or digest.hexdigest() != expected
or tuple(getattr(after, key) for key in _MODEL_REFERENCE_STAT_FIELDS) != identity
or tuple(getattr(current, key) for key in _MODEL_REFERENCE_STAT_FIELDS) != identity):
_bad_model_metadata(str(path), f"sha256={expected}")
if cancelled is not None and cancelled():
raise DownloadCancelled()
except OSError:
_bad_model_metadata(str(path), f"sha256={expected}")
return _VerifiedModelReference(tuple(record.items()), str(actual_path), identity,
_MODEL_REFERENCE_VERIFIED)
def _attach_model_reference(module, settings, name, verified):
"""Attach one verified picker result to a copied matching settings payload.
:param module: actual consuming module key.
:param settings: current settings; existing pins are validated before merging.
:param name: selected setting, whose request and consumer kind must still match.
:param verified: verifier-produced immutable result, never a raw supplied dict.
:returns: independent settings with one ordinary JSON-schema pin attached.
:raises ModelZooError: for stale input/file identity, forged result, invalid
existing pins or incompatible current consumer. No hashing occurs here.
"""
from copy import deepcopy
kinds = _model_input_kinds(module, settings)
metadata = _validate_model_metadata(settings, model_kinds=kinds)
if (not isinstance(verified, _VerifiedModelReference)
or verified._verification is not _MODEL_REFERENCE_VERIFIED):
_bad_model_metadata(name, "model_metadata")
record = dict(verified.record)
candidate = {**settings, "model_metadata": {"schema_version": 1, "settings": {name: record}}}
_validate_model_metadata_structure(candidate, model_kinds=kinds)
_match_model_reference(record)
try:
current = Path(verified.path).stat()
except OSError:
_bad_model_metadata(name, f"sha256={record['sha256']}")
if tuple(getattr(current, key) for key in _MODEL_REFERENCE_STAT_FIELDS) != verified.identity:
_bad_model_metadata(name, f"sha256={record['sha256']}")
result = deepcopy(settings)
pins = metadata["settings"] if metadata is not None else {}
pins[name] = record
result["model_metadata"] = {"schema_version": 1, "settings": pins}
return result
def _bad_model_metadata(field, expected):
"""Raise a translated diagnostic naming the incompatible metadata field."""
from .qt.i18n import tr
raise ModelZooError(tr("Bad settings") + f": {field} ({expected})")
def _validate_model_metadata_structure(settings, *, model_kinds=None):
"""Validate recorded identities against their exact original setting values."""
from copy import deepcopy
if not isinstance(settings, Mapping):
_bad_model_metadata("settings", "mapping")
if "model_metadata" not in settings:
return None
metadata = settings["model_metadata"]
if not isinstance(metadata, dict) or set(metadata) != {"schema_version", "settings"}:
_bad_model_metadata("model_metadata", "schema_version, settings")
if type(metadata["schema_version"]) is not int or metadata["schema_version"] != 1:
_bad_model_metadata("model_metadata.schema_version", "1")
records = metadata["settings"]
if not isinstance(records, dict):
_bad_model_metadata("model_metadata.settings", "mapping")
required = {"requested_value", "key", "version", "version_status", "sha256", "kind", "filename"}
for name, record in records.items():
field = f"model_metadata.settings[{name!r}]"
if not isinstance(name, str) or not name.strip():
_bad_model_metadata(field, "setting key")
if not isinstance(record, dict) or set(record) != required:
_bad_model_metadata(field, ", ".join(sorted(required)))
for key in required - {"version"}:
if not isinstance(record[key], str) or not record[key].strip():
_bad_model_metadata(f"{field}.{key}", "string")
if not isinstance(settings.get(name), str) or settings[name] != record["requested_value"]:
_bad_model_metadata(f"{field}.requested_value", f"settings[{name!r}]")
if not re.fullmatch(r"[0-9a-f]{64}", record["sha256"]):
_bad_model_metadata(f"{field}.sha256", "[0-9a-f]{64}")
if record["version_status"] == "declared":
if not isinstance(record["version"], str) or not record["version"].strip():
_bad_model_metadata(f"{field}.version", "string")
elif record["version_status"] != "unspecified" or record["version"] is not None:
_bad_model_metadata(f"{field}.version_status/version", "declared/string | unspecified/null")
filename = record["filename"]
if ("\x00" in filename or re.split(r"[/\\]", filename)[-1] in {"", "."}
or filename.startswith(("/", "\\")) or re.match(r"^[A-Za-z]:", filename)
or "://" in filename or ".." in re.split(r"[/\\]", filename)):
_bad_model_metadata(f"{field}.filename", "relative catalogue filename")
if record["kind"] not in KINDS:
_bad_model_metadata(f"{field}.kind", " | ".join(KINDS))
if model_kinds is not None:
accepted = model_kinds.get(name, ())
if isinstance(accepted, str):
accepted = (accepted,)
if record["kind"] not in accepted:
_bad_model_metadata(f"{field}.kind", f"{name}: {tuple(accepted)}")
return deepcopy(metadata)
def _validate_model_metadata(settings, *, model_kinds=None):
"""Check bindings and reject contradictions with the bundled trusted index."""
metadata = _validate_model_metadata_structure(settings, model_kinds=model_kinds)
if metadata is not None:
known = {row.get("key") or row.get("name") for row in BUNDLED_REMOTE_MODELS}
for record in metadata["settings"].values():
if record["key"] in known:
_match_model_reference(record)
return metadata
def _invalidate_model_metadata(settings, keys):
"""Remove valid pins for explicit edits without changing the input payload."""
from copy import deepcopy
metadata = _validate_model_metadata(settings)
result = deepcopy(settings)
if metadata is not None:
for key in keys:
metadata["settings"].pop(key, None)
result["model_metadata"] = metadata
return result
def _model_reference_from_record(requested, record):
"""Build a pin from a trusted raw catalogue row without inferring a version."""
if not isinstance(record, Mapping):
_bad_model_metadata("catalogue", "mapping")
declared = record.get("version")
if declared is not None and not isinstance(declared, str):
_bad_model_metadata("catalogue.version", "string | null")
version = declared if declared and declared.strip() else None
result = {"requested_value": requested, "key": record.get("key") or record.get("name"),
"version": version, "version_status": "declared" if version is not None else "unspecified",
"sha256": record.get("sha256"), "kind": record.get("kind", "cellpose"),
"filename": record.get("name")}
_validate_model_metadata_structure({"model": requested, "model_metadata": {
"schema_version": 1, "settings": {"model": result}}})
return result
def _match_model_reference(record, catalogue_records=None):
"""Select a conjunction of raw trusted identity fields without first-match fallback."""
if not isinstance(record, Mapping):
_bad_model_metadata("model_metadata", "mapping")
_validate_model_metadata_structure({"model": record.get("requested_value"), "model_metadata": {
"schema_version": 1, "settings": {"model": record}}})
matches = []
rows = BUNDLED_REMOTE_MODELS if catalogue_records is None else catalogue_records
for candidate in rows:
if candidate.get("key", candidate.get("name")) != record["key"]:
continue
pin = _model_reference_from_record(record["requested_value"], candidate)
if pin == record and pin not in matches:
matches.append(pin)
elif pin != record:
_bad_model_metadata(f"model_metadata[{record['key']!r}]", "catalogue identity")
if len(matches) != 1:
_bad_model_metadata(f"model_metadata[{record['key']!r}]", "one exact catalogue identity")
return matches[0]
def _resolve_pinned_model_file(settings, name, *, kinds=(), download=True):
"""Verify recorded local bytes or acquire the exact trusted recorded identity."""
metadata = _validate_model_metadata(settings)
if metadata is None or name not in metadata["settings"]:
return None
record = metadata["settings"][name]
if kinds and record["kind"] not in kinds:
_bad_model_metadata(f"model_metadata.settings[{name!r}].kind", ", ".join(kinds))
try:
requested = Path(record["requested_value"]).expanduser()
except RuntimeError:
requested = None
if requested is not None and requested.is_file():
actual = sha256_file(requested)
if actual != record["sha256"]:
_bad_model_metadata(
f"{name}={record['requested_value']!r}; {record['key']} {record['version']!r}; {requested}",
f"sha256={record['sha256']}; actual={actual}")
return requested
_match_model_reference(record)
raw = next(row for row in BUNDLED_REMOTE_MODELS
if row.get("key", row.get("name")) == record["key"]
and _model_reference_from_record(record["requested_value"], row) == record)
resolved = _ensure_catalogue_model_file(
_entry_from_mapping(raw), download=download, setting=name,
requested=record['requested_value'], publisher_version=record['version'])
if resolved is None:
_bad_model_metadata(name, f"{record['key']} sha256={record['sha256']}")
resolved = Path(resolved)
if resolved.is_file():
actual = sha256_file(resolved)
if actual != record["sha256"]:
if not download:
return versioned_path(resolved.parent, resolved.name)
_bad_model_metadata(
f"{name}={record['requested_value']!r}; {record['key']} {record['version']!r}; {resolved}",
f"sha256={record['sha256']}; actual={actual}")
elif download:
_bad_model_metadata(str(resolved), f"sha256={record['sha256']}")
return resolved
def _planning_model_entries():
"""Bounded trusted declarations for planning without providers or network."""
return [_entry_from_mapping(record) for record in BUNDLED_REMOTE_MODELS]
def _resolution_model_entries(download):
"""Select declared metadata without executing providers on an offline node.
:param download: false retains the existing bundled-only pure planning index.
:returns: online catalogue entries, or offline bundled/local/cached declarations
in first-key/name order without provider calls or metadata transport.
"""
if not download:
return _planning_model_entries()
if os.environ.get("SPACR_MODEL_DOWNLOAD", "").strip().casefold() not in (
"0", "off", "false", "no"):
return catalogue(remote=True)
entries = _planning_model_entries()
path = os.environ.get(CATALOGUE_ENV_VAR, "")
if path and os.path.isfile(path):
entries.extend(load_catalogue_file(path))
entries.extend(_SHARED_CATALOGUE_CACHE["entries"])
selected = []
have = set()
for entry in entries:
identity = (entry.key, entry.name)
if identity not in have and entry.name not in RETIRED_MODEL_NAMES:
selected.append(entry)
have.add(identity)
return selected
def _ensure_model_file(requested, *, kinds=(), download=True, setting=None, original_request=None):
"""Resolve a catalogue checkpoint and fetch missing verified weights.
Local paths win. Unknown identifiers return None for the caller's
existing stock-model or missing-path handling. With download=False,
return the intended cache path without fetching checkpoint bytes.
A Qt GUI-thread caller must use the model picker for a missing model.
Optional ``setting`` names the actual input field in a successful-download log.
``original_request`` preserves a caller's unmodified prefix/target string in
that log; None uses the lookup request. It never changes identity selection.
"""
text = os.fspath(requested).strip() if requested is not None else ""
if not text:
return None
try:
local = Path(text).expanduser()
except RuntimeError:
local = None
if local is not None and local.is_file():
return local
if "://" in text:
return None
path_like = "/" in text or "\\" in text or text.startswith("~")
identifier = Path(text.replace("\\", "/")).name if path_like else text
entries = _resolution_model_entries(download)
matches = [item for item in entries
if item.key == identifier or item.name == identifier
or (path_like and Path(item.name).name == identifier)]
if len(matches) > 1:
from .qt.i18n import tr
choices = "\n".join(
f"{item.key}@{item.version} [sha256={item.sha256}]" for item in matches)
raise ModelZooError(
tr("{name} · {n} matches", name=repr(text), n=len(matches))
+ ":\n" + choices)
if not matches:
return None
entry = matches[0]
if kinds and entry.kind not in kinds:
raise ModelZooError(
f"{text!r} is a {entry.kind} model; this operation requires "
f"one of {', '.join(kinds)}")
return _ensure_catalogue_model_file(
entry, download=download, setting=setting,
requested=text if original_request is None else original_request,
publisher_version=next((row.get("version") for row in BUNDLED_REMOTE_MODELS
if row.get("key", row.get("name")) == entry.key
and row.get("name") == entry.name
and row.get("kind", "cellpose") == entry.kind
and row.get("sha256") == entry.sha256
and _entry_from_mapping(row).uri == entry.uri), Ellipsis))
def _model_resolution_context(requested, *, setting, kinds=()):
"""Describe an unresolved input using only trusted declared candidate identities.
:param requested: original missing request, retained verbatim in the diagnostic.
:param setting: actual setting or direct API parameter name supplied by its caller.
:param kinds: accepted model kinds; other declarations are excluded.
:returns: localized setting/request context and exact declared keys, versions and hashes.
"""
from .qt.i18n import tr
choices = [
f"{row.get('key', row.get('name'))}@{row.get('version')!r} "
f"[sha256={row.get('sha256')}]"
for row in BUNDLED_REMOTE_MODELS
if not kinds or row.get("kind", "cellpose") in kinds]
return tr(
"Model setting {setting}: no existing checkpoint or exact model-zoo identity matches {requested}. Available compatible zoo identities: {candidates}.",
setting=setting, requested=repr(requested), candidates="\n".join(choices) or "[]")
_MODEL_LOCK_STALE_ENV = "SPACR_MODEL_LOCK_STALE_SECONDS"
_MODEL_LOCK_STALE_SECONDS = 600.0
def _model_lock_stale_seconds() -> float:
"""Seconds an unrefreshed download lock lasts before another node reclaims it.
:returns: ``SPACR_MODEL_LOCK_STALE_SECONDS`` when it is a number of at
least one second, otherwise 600.
"""
try:
value = float(os.environ.get(_MODEL_LOCK_STALE_ENV, "").strip()
or _MODEL_LOCK_STALE_SECONDS)
except ValueError:
return _MODEL_LOCK_STALE_SECONDS
return value if value >= 1.0 else _MODEL_LOCK_STALE_SECONDS
def _refresh_model_download_lock(path, stop, interval) -> None:
"""Touch a held download lock until ``stop`` is set.
Waiting jobs on other nodes read the changing modification time as proof
that the download is still running.
:param path: the lock file.
:param stop: ``threading.Event`` set when the download ends.
:param interval: seconds between touches.
"""
while not stop.wait(interval):
try:
os.utime(path)
except OSError:
return
def _model_lock_signature(path):
"""Identity and modification time of a lock file, or None when it is absent."""
try:
info = os.stat(path)
except OSError:
return None
return info.st_dev, info.st_ino, info.st_mtime_ns, info.st_size
def _acquire_model_download_lock(lock, lock_path, label) -> None:
"""Hold the download lock for one model, waiting for another job or node.
Qt reclaims a lock left by a dead process on this host. A lock written on
another host cannot be checked that way, so its holder refreshes the lock
file's modification time while it downloads. A lock that stays unchanged
for ``SPACR_MODEL_LOCK_STALE_SECONDS`` (600 by default) of this node's own
clock is taken as left by a crashed node and removed. Clocks of different
nodes are never compared.
:param lock: an unlocked ``QLockFile`` for ``lock_path``.
:param lock_path: the lock file in the shared model cache.
:param label: model key named in the waiting log line.
:raises OSError: EACCES when the lock cannot be written, EBUSY for any
other lock failure.
"""
import errno
from PySide6.QtCore import QLockFile
stale_after = _model_lock_stale_seconds()
poll = max(50, min(1000, int(stale_after * 250)))
seen = None
since = time.monotonic()
announced = False
while not lock.tryLock(poll):
if lock.error() != QLockFile.LockFailedError:
code = errno.EACCES if lock.error() == QLockFile.PermissionError else errno.EBUSY
raise OSError(code, os.strerror(code), str(lock_path))
if not announced:
LOG.info("Waiting for another job to finish downloading model %s (lock %s)",
label, lock_path)
announced = True
signature = _model_lock_signature(lock_path)
if signature != seen:
seen, since = signature, time.monotonic()
elif signature is not None and time.monotonic() - since >= stale_after:
LOG.warning("Removing download lock %s: unchanged for %.0f s, its holder stopped",
lock_path, stale_after)
lock.removeStaleLockFile()
seen, since = None, time.monotonic()
def _manual_model_copy_command(entry, cached) -> str:
"""Shell command that places one verified checkpoint where an offline node looks.
Run it on a machine that reaches the model's source and shares the cache
folder with the node, or run it elsewhere and copy the file to the same
path.
:param entry: the selected catalogue entry.
:param cached: the cache path the node checks.
:returns: a one-line POSIX shell command that downloads and checks SHA-256.
"""
import shlex
cached = Path(cached).as_posix()
target = shlex.quote(cached)
uri = str(entry.uri or "")
if uri.startswith(("http://", "https://")):
get = f"curl -fL -o {target} {shlex.quote(uri)}"
else:
source = uri[len("file://"):] if uri.startswith("file://") else uri
get = f"cp {shlex.quote(source)} {target}"
check = shlex.quote(f"{str(entry.sha256 or '').lower()} {cached}")
return (f"mkdir -p {shlex.quote(Path(cached).parent.as_posix())} && {get} "
f"&& echo {check} | sha256sum -c -")
def _ensure_catalogue_model_file(entry, *, download=True, setting=None, requested=None,
publisher_version=Ellipsis):
"""Acquire one identity under the cache lock and report successful contextual downloads.
:param entry: exact selected catalogue entry.
:param download: permit acquisition; false retains pure planning behavior.
:param setting: actual input field, or None when the caller has no field context.
:param requested: original path/key, preserved independently of the effective path.
:param publisher_version: raw recorded version, including None when unspecified;
Ellipsis uses the selected entry version without inferring one from storage suffixes.
:returns: verified existing or acquired path, or None for a non-acquirable entry.
"""
if entry.path and Path(entry.path).is_file():
return Path(entry.path)
if not entry.uri:
return None
destination = _model_cache_root()
cached = destination / Path(entry.name).name
if cached.is_file():
if entry.sha256 and sha256_file(cached) == entry.sha256.lower():
return cached
if entry.sha256 and destination.is_dir():
for version in sorted(destination.glob(
f"{cached.stem}_v*{cached.suffix}")):
if version.is_file() and sha256_file(version) == entry.sha256.lower():
return version
if not download:
return cached
if os.environ.get("SPACR_MODEL_DOWNLOAD", "").strip().casefold() in (
"0", "off", "false", "no"):
from .qt.i18n import tr
raise ModelUnreadable(tr(
"Automatic model downloads are disabled by SPACR_MODEL_DOWNLOAD. "
"Copy verified model {model} (sha256 {sha256}) to {path}.",
model=entry.key, sha256=entry.sha256, path=str(cached))
+ "\n" + _manual_model_copy_command(entry, cached))
if _on_the_qt_gui_thread():
raise ModelUnreadable(
f"{entry.key!r} is not cached. Download it in Model Zoo, or run "
f"the operation through spacr-run or Python without a GUI.")
from PySide6.QtCore import QLockFile
destination.mkdir(parents=True, exist_ok=True)
identity = hashlib.sha256(
f"{cached.name}\0{str(entry.sha256 or '').lower()}".encode("utf-8")).hexdigest()
lock_path = destination / f".model-{identity}.lock"
lock = QLockFile(str(lock_path))
lock.setStaleLockTime(0)
_acquire_model_download_lock(lock, lock_path, entry.key)
stop = threading.Event()
heartbeat = threading.Thread(
target=_refresh_model_download_lock, name="spacr-model-lock-heartbeat",
args=(lock_path, stop, min(30.0, _model_lock_stale_seconds() / 4)), daemon=True)
heartbeat.start()
try:
if cached.is_file() and entry.sha256 and sha256_file(cached) == entry.sha256.lower():
return cached
if entry.sha256:
for version in sorted(destination.glob(
f"{cached.stem}_v*{cached.suffix}")):
if version.is_file() and sha256_file(version) == entry.sha256.lower():
return version
LOG.info("Downloading requested model %s into %s", entry.key, destination)
try:
acquired = fetch(entry, destination)
except ChecksumMismatch as error:
if setting is None:
raise
version = (f"ModelEntry.version={entry.version!r}"
if publisher_version is Ellipsis else repr(publisher_version))
context = f"{setting}={requested!r}; {entry.key} {version} sha256={entry.sha256}"
raise ChecksumMismatch(context + "\n" + str(error)) from error
if setting is not None:
from .qt.i18n import tr
version = (f"ModelEntry.version={entry.version!r}"
if publisher_version is Ellipsis else repr(publisher_version))
LOG.info(tr(
"{setting}: {requested} was not found here; downloaded zoo model {key} {version} (sha256 {sha256}) to {path}.",
setting=setting, requested=requested, key=entry.key,
version=version, sha256=entry.sha256, path=str(acquired)))
return acquired
finally:
stop.set()
heartbeat.join()
lock.unlock()
[docs]
def versioned_path(dest: Any, filename: str) -> Path:
"""The first free destination for ``filename`` in ``dest``.
``foo.CP_model`` -> ``foo.CP_model``, then ``foo_v2.CP_model``,
``foo_v3.CP_model``… An existing checkpoint is never overwritten: two
models with the same filename are a normal thing to have (the same author
retrained, or two people picked the same name), and the failure mode of
overwriting — a run that used the old weights becoming unreproducible with
no trace — is silent.
An input that already carries ``_vN`` counts from there rather than
becoming ``foo_v2_v2``.
:param dest: destination directory.
:param filename: the name to place there.
:returns: a path that does not exist yet.
"""
folder = Path(dest)
p = Path(filename)
suffix = p.suffix
match = _VERSION_RE.match(p.stem)
base = match.group("base") if match else p.stem
n = int(match.group("n")) if match else 1
candidate = folder / p.name
while candidate.exists() or candidate.is_symlink():
n += 1
candidate = folder / f"{base}_v{n}{suffix}"
return candidate
[docs]
def open_uri(uri: str, timeout: int = DEFAULT_TIMEOUT,
chunk_size: int = DEFAULT_CHUNK
) -> Tuple[Iterable[bytes], int]:
"""Open a model URI for streaming. ``(chunks, total_bytes)``.
``http://`` and ``https://`` stream over ``requests`` — the same call
:func:`spacr.utils.download_models` and
:func:`spacr.qt.hf_download._download_one` make, imported here so this
module has no hard dependency on it. ``file://`` and a plain existing path
are read from disk, which is what a lab mirror on a NAS looks like and what
the tests use, so the whole fetch path is exercised without a network.
:param uri: where the model lives.
:param timeout: seconds, HTTP only.
:param chunk_size: bytes per chunk.
:returns: ``(iterable of byte chunks, total size or 0 when unknown)``.
:raises ModelZooError: for a scheme this does not speak.
"""
text = str(uri or "")
if text.startswith(("http://", "https://")):
import requests
response = requests.get(text, stream=True, timeout=timeout)
response.raise_for_status()
total = int(response.headers.get("content-length") or 0)
return _resumed_chunks(text, response, total, timeout,
chunk_size), total
if text.startswith("file://"):
from urllib.parse import urlparse
from urllib.request import url2pathname
parsed = urlparse(text)
host = parsed.netloc if parsed.netloc not in ("", "localhost") else ""
local = Path(url2pathname(("//" + host if host else "") + parsed.path))
elif os.path.exists(text):
local = Path(text)
else:
raise ModelZooError(
f"do not know how to fetch {text!r} — expected an http(s):// URL, "
f"a file:// URL, or a path that exists")
if not local.is_file():
raise ModelUnreadable(f"no such model file: {local}")
return _read_chunks(local, chunk_size), local.stat().st_size
#: How many times a dropped download picks up where it stopped.
_RESUMES = 5
def _resumed_chunks(url: str, response: Any, total: int, timeout: int,
chunk_size: int) -> Iterator[bytes]:
"""A download's chunks, resumed with a ``Range`` request when the
server drops the connection part-way.
bioimage.io's file server dropped two of three 1.2 GB Cellpose-SAM
downloads part-way on 2026-09-25 ("Response ended prematurely"), each
after minutes of transfer. A server that answers a range with 206 and
the offset asked for is resumed from the last byte received; any other
answer ends the download as before. The checksum :func:`fetch` checks
covers the joined file, so a bad join is refused like any bad file.
"""
import requests
done = 0
resumes = _RESUMES
while True:
try:
for block in response.iter_content(chunk_size=chunk_size):
done += len(block)
yield block
return
except (requests.exceptions.ChunkedEncodingError,
requests.exceptions.ConnectionError) as dropped:
resumes -= 1
if resumes < 0 or (total and done >= total):
raise
response = requests.get(url, stream=True, timeout=timeout,
headers={"Range": f"bytes={done}-"})
span = str(response.headers.get("content-range") or "")
if response.status_code != 206 or not span.startswith(
f"bytes {done}-"):
raise dropped
def _read_chunks(path: Path, chunk_size: int) -> Iterator[bytes]:
"""Read a file in fixed-size blocks.
:param path: the file.
:param chunk_size: the block size.
:returns: an iterator over the blocks -- streamed rather than read
whole, because these are multi-gigabyte checkpoints being hashed.
"""
with path.open("rb") as handle:
while True:
block = handle.read(chunk_size)
if not block:
return
yield block
[docs]
def fetch(entry: ModelEntry, dest: Any,
expected_sha256: Optional[str] = None,
require_checksum: bool = True,
opener: Optional[Callable[[str], Any]] = None,
progress: Optional[Callable[[int, int], None]] = None,
cancel: Optional[Callable[[], bool]] = None,
chunk_size: int = DEFAULT_CHUNK,
timeout: int = DEFAULT_TIMEOUT) -> Path:
"""Download a model, verify it, and only then put it where it belongs.
The steps run in this order:
1. bytes stream into a temporary file **inside** ``dest``, so the rename in
step 4 is a same-filesystem ``os.replace`` and therefore atomic;
2. the checksum of what actually arrived is computed;
3. if it does not match the published digest the temporary file is deleted
and :class:`ChecksumMismatch` is raised — nothing is installed, and the
destination still holds whatever it held before;
4. only now is the temporary file renamed, to a
:func:`versioned_path` that does not exist yet.
Every failure — a dead server, a cancel, a bad hash, a full disk — leaves
the destination directory exactly as it was. There is no window in which a
half-written file sits at a name that looks like a model.
:param entry: what to fetch. :attr:`ModelEntry.uri` is the source.
:param dest: destination directory; created if missing.
:param expected_sha256: digest to require, overriding
:attr:`ModelEntry.sha256`.
:param require_checksum: refuse to install when no digest is known. True by
default, so an unverifiable download needs an explicit opt-in. Pass
False to accept one knowingly; the
entry :func:`install` returns then reports ``verified=False``.
:param opener: ``fn(uri) -> chunks`` or ``fn(uri) -> (chunks, total)``;
defaults to :func:`open_uri`.
:param progress: ``fn(done_bytes, total_bytes)``; ``total`` is 0 when the
server did not say.
:param cancel: ``fn() -> bool``, polled between chunks. Returning True
deletes the partial file and raises :class:`DownloadCancelled`.
:param chunk_size: bytes per read.
:param timeout: seconds, HTTP only.
:returns: the path the model was written to.
:raises ChecksumMismatch: the bytes are not the published bytes.
:raises DownloadCancelled: ``cancel()`` returned True.
:raises ModelZooError: no URI, or no checksum with ``require_checksum``.
"""
if entry.kind == "timepose":
record = _timepose_record(entry.key)
folder = _timepose_files(entry.key, dest=Path(dest) / "timepose")
return folder / record["weights"]
if not entry.uri:
raise ModelZooError(
f"{entry.name} has no uri to fetch it from (source={entry.source})")
want = (expected_sha256 if expected_sha256 is not None
else entry.sha256) or ""
want = want.strip().lower()
if require_checksum and not want:
raise ModelZooError(
f"refusing to install {entry.name}: no sha256 was published for "
f"it, so a truncated or substituted checkpoint could not be told "
f"from the real one. Supply expected_sha256=…, put a hash in the "
f"catalogue, or pass require_checksum=False to accept it "
f"unverified.")
folder = Path(dest)
folder.mkdir(parents=True, exist_ok=True)
import tempfile
handle = tempfile.NamedTemporaryFile(
dir=str(folder), prefix=f".{Path(entry.name).stem}.", suffix=".part",
delete=False)
temp = Path(handle.name)
done = 0
try:
stream = (opener(entry.uri) if opener is not None
else open_uri(entry.uri, timeout=timeout,
chunk_size=chunk_size))
chunks, total = stream if isinstance(stream, tuple) else (stream, 0)
total = int(total or entry.size_bytes or 0)
if progress is not None:
progress(0, total)
for block in chunks:
if cancel is not None and cancel():
raise DownloadCancelled(
f"download of {entry.name} cancelled after {done} byte(s) "
f"— nothing was written to {folder}")
if not block:
continue
handle.write(block)
done += len(block)
if progress is not None:
progress(done, total)
handle.flush()
os.fsync(handle.fileno())
handle.close()
if done == 0:
raise ModelZooError(
f"{entry.uri} returned no data for {entry.name} — nothing was "
f"written to {folder}")
if (entry.source == "bioimage.io"
and entry.kind in ("cellpose", "cellpose3", "cellpose_dino")
and done < _BIOIMAGEIO_MIN_WEIGHTS):
raise ModelZooError(
f"{entry.name} from {entry.uri} is {done} bytes, too small to "
f"be a Cellpose network, and was NOT installed in {folder}. "
f"Its bioimage.io package publishes a stand-in and fetches "
f"the real weights from elsewhere when it runs, which spaCR "
f"does not do.")
got = sha256_file(temp)
if want and got != want:
raise ChecksumMismatch(
f"{entry.name} does not match its published checksum and was "
f"NOT installed in {folder}.\n"
f" expected sha256 {want}\n"
f" got sha256 {got}\n"
f" ({done} bytes from {entry.uri})\n"
f"A checkpoint that fails this still loads and still produces "
f"masks — they are just not the masks the author's model "
f"produces. Re-download, or get the right hash from whoever "
f"published it.")
return _claim(temp, folder, entry.name)
except BaseException:
try:
handle.close()
except Exception:
pass
try:
if temp.exists():
temp.unlink()
except OSError:
pass
raise
def _claim(temp: Path, folder: Path, name: str) -> Path:
"""Move ``temp`` onto the first free version of ``name``, race-free.
:func:`versioned_path` alone is check-then-act: two downloads of the same
model finishing together both see ``foo.CP_model`` free, and the second
``os.replace`` silently destroys the first. So the name is *reserved* with
``O_CREAT | O_EXCL`` — which the kernel makes atomic — before the rename,
and a lost race just tries the next version.
The reservation is a zero-byte file that exists for microseconds and is
then replaced. If the process is killed inside that window what remains is
an empty file, which :func:`inspect_checkpoint` rejects by name — unlike a
truncated checkpoint, which loads.
"""
while True:
target = versioned_path(folder, name)
try:
descriptor = os.open(str(target),
os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o644)
except FileExistsError:
continue
try:
created = os.fstat(descriptor)
finally:
os.close(descriptor)
try:
os.replace(str(temp), str(target))
except BaseException:
try:
if os.path.samestat(target.lstat(), created):
target.unlink()
except OSError:
pass
raise
return target
[docs]
def install(entry: ModelEntry, dest: Any, **kwargs: Any) -> ModelEntry:
""":func:`fetch` the model and return the registered local entry.
The returned entry carries the digest of the bytes that were actually
written — not the one the catalogue claimed — and
:attr:`ModelEntry.verified` is True only when the two were compared and
matched. Provenance (``trained_on``, ``trained_by``), metrics and licence
are carried over from the catalogue entry.
:param entry: the remote entry.
:param dest: destination directory.
:param kwargs: passed to :func:`fetch`.
:returns: a ``source='local'`` entry pointing at the new file.
"""
expected = kwargs.get("expected_sha256")
want = ((expected if expected is not None else entry.sha256) or "").strip()
path = fetch(entry, dest, **kwargs)
digest = sha256_file(path)
notes = tuple(n for n in entry.notes if "no published checksum" not in n)
if not want:
notes = notes + (
f"installed without a published checksum to check against; the "
f"recorded sha256 {digest[:12]}… is of the bytes that arrived, "
f"which proves nothing about where they came from",)
return replace(
entry,
key=_key_for(path),
name=path.name,
source="local",
path=str(path.resolve()),
version=_version_of(path.name),
sha256=digest,
size_bytes=path.stat().st_size,
verified=bool(want),
notes=notes,
)
[docs]
def download_bundled_models(**kwargs: Any) -> str:
"""Pull the bundled Hugging Face model pack via the existing downloader.
Thin, deliberate wrapper over :func:`spacr.utils.download_models` — the
downloader spaCR already ships and the one
:func:`spacr.submodules.analyze_plaques` depends on. It is *not*
reimplemented here, so there is one code path that fills
``resources/models`` and one place to fix when the repo moves.
The downloader stages the whole pack in a temporary sibling directory
and atomically renames the completed directory into place after every
download succeeds. It does not verify published checksums and trusts an
existing non-empty cache without checking its completeness or origin.
Prefer a catalogue entry with a hash and :func:`install` when checksum
verification is required.
``spacr.utils`` imports torch, so it is imported here and not at module
level. Nothing else in this module reaches for it.
:param kwargs: forwarded to :func:`spacr.utils.download_models`.
:returns: the local directory the pack landed in.
"""
return _bulk_downloader()(**kwargs)
def _bulk_downloader() -> Callable[..., str]:
"""The legacy bulk downloader, imported late (it pulls torch in with it)."""
from .utils import download_models
return download_models
@dataclass
[docs]
class FieldBenchmark:
"""One field's result for one model.
:param field: the field name.
:param n_objects: labels in the mask this model produced.
:param severity: :mod:`spacr.seg_qc`'s verdict — ``'ok'``, ``'warn'`` or
``'fail'``, or ``'-'`` when QC was off.
:param flags: the named defects seg_qc raised.
:param note: seg_qc's verdict in prose, with its numbers in it.
"""
field: str
n_objects: int = 0
severity: str = "-"
flags: Tuple[str, ...] = ()
note: str = ""
@dataclass
[docs]
class BenchmarkResult:
"""One model over one field set. Only comparable to results on the *same* set.
:param entry: the model that ran.
:param fieldset: :func:`fieldset_id` of the images — a hash of the pixels,
so two runs over the same three fields share it and two runs over
different fields never do, whatever the folders were called.
:param fieldset_label: a human-readable description of the field set:
the field count, the source folder and the first field names.
:param rows: one :class:`FieldBenchmark` per field, in field order.
:param seconds: wall-clock seconds the model spent segmenting.
:param honoured: the parameters that reached the model.
:param ignored: what was set and Cellpose 4 dropped (``diam_mean`` and
friends) — carried so a benchmark cannot silently be a benchmark of
settings nothing read.
:param notes: warnings the reader must see before the numbers.
:param masks: the label images, when kept, so a GUI can draw them.
:param images: the source fields, likewise.
:param object_type: what was segmented.
"""
entry: ModelEntry
fieldset: str = ""
fieldset_label: str = ""
rows: List[FieldBenchmark] = _dc_field(default_factory=list)
seconds: float = 0.0
honoured: Dict[str, Any] = _dc_field(default_factory=dict)
ignored: Dict[str, Any] = _dc_field(default_factory=dict)
notes: List[str] = _dc_field(default_factory=list)
masks: List[Any] = _dc_field(default_factory=list)
images: List[Any] = _dc_field(default_factory=list)
object_type: str = "cell"
@property
[docs]
def fields(self) -> List[str]:
"""Every field this model was benchmarked on.
:returns: the field names.
"""
return [r.field for r in self.rows]
@property
[docs]
def n_fields(self) -> int:
"""How many fields were benchmarked.
:returns: the field count.
"""
return len(self.rows)
@property
[docs]
def total_objects(self) -> int:
"""Every object the model found, across all fields.
:returns: the object count.
"""
return sum(r.n_objects for r in self.rows)
@property
[docs]
def mean_objects(self) -> float:
"""Objects per field.
NaN rather than zero when nothing was benchmarked: no fields is a
different statement from a model that found nothing.
:returns: the mean, or NaN.
"""
return self.total_objects / self.n_fields if self.rows else float("nan")
@property
[docs]
def n_failed(self) -> int:
"""Fields :mod:`spacr.seg_qc` scored ``'fail'``."""
return sum(1 for r in self.rows if r.severity == "fail")
@property
[docs]
def n_ok(self) -> int:
"""How many fields came back without a quality complaint.
:returns: the count of fields at severity ``ok``.
"""
return sum(1 for r in self.rows if r.severity == "ok")
@property
[docs]
def qc_score(self) -> float:
"""Fraction of fields seg_qc scored ``'ok'``; ``nan`` without QC.
A quality-control verdict on this model's own masks — it says the masks
are not obviously broken, not that they are right. There is no ground
truth in a benchmark (see the module docstring), so the zoo reports no
accuracy score.
"""
scored = [r for r in self.rows if r.severity != "-"]
if not scored:
return float("nan")
return sum(1 for r in scored if r.severity == "ok") / len(scored)
@property
[docs]
def summary(self) -> str:
"""The model, what it found, over how many fields, and its QC score.
:returns: a one-line summary.
"""
score = self.qc_score
return (f"{self.entry.name}: {self.total_objects} "
f"{self.object_type}(s) over {self.n_fields} field(s) "
f"({self.mean_objects:.1f}/field) in {self.seconds:.1f}s; "
f"seg_qc scored {self.n_ok}/{self.n_fields} field(s) ok"
+ ("" if score != score else f" ({score * 100:.0f}%)")
+ f", {self.n_failed} fail.")
[docs]
def fieldset_id(names: Sequence[str], images: Sequence[Any]) -> str:
"""A stable id for a set of fields, taken from the **pixels**.
Folder names are not identity: ``plate1/1`` on two machines is two
different sets of images, and the same three images copied to a new folder
are the same benchmark input. So the id hashes each array's bytes, shape and
dtype together with its name.
:func:`rank` uses this id to refuse results from different field sets;
without it, benchmarks run on different data would be sorted together.
:param names: field names, in order.
:param images: the arrays, in the same order.
:returns: a 16-character hex id.
"""
digest = hashlib.sha256()
for name, image in zip(names, images):
array = np.ascontiguousarray(np.asarray(image))
digest.update(str(name).encode("utf-8", "replace"))
digest.update(b"\x00")
digest.update(str(array.shape).encode())
digest.update(str(array.dtype).encode())
digest.update(hashlib.sha256(array.tobytes()).digest())
return digest.hexdigest()[:16]
def _model_compare():
""":mod:`spacr.model_compare`, imported late so browsing stays cheap."""
from . import model_compare
return model_compare
[docs]
def config_for(entry: ModelEntry, overrides: Optional[Mapping[str, Any]] = None):
"""The :class:`spacr.model_compare.ModelConfig` that runs this entry.
A local checkpoint is passed by path, which
:func:`spacr.utils._choose_model` and
:func:`spacr.model_compare.segment_with_cellpose` both load as
``pretrained_model``; anything else goes through by name and is subject to
Cellpose 4's legacy-name remapping, which the config reports.
:param entry: the model.
:param overrides: eval settings (``diameter``, ``flow_threshold``, …).
:returns: the config.
"""
mc = _model_compare()
settings: Dict[str, Any] = dict(overrides or {})
settings["model"] = entry.path or entry.name
settings.setdefault("name", entry.name)
return mc.ModelConfig.from_mapping(settings)
[docs]
def benchmark(entry: ModelEntry, images: Optional[Sequence[Any]] = None,
source: Any = None, n_fields: int = DEFAULT_N_FIELDS,
field_names: Optional[Sequence[str]] = None,
segment_fn: Optional[Callable] = None,
settings: Optional[Mapping[str, Any]] = None,
object_type: str = "cell", qc: bool = True,
keep_images: bool = True, channel: Optional[int] = None,
progress: Optional[Callable[[str, int, int], None]] = None,
) -> BenchmarkResult:
"""Run one model over N fields and report what came out. "Test on 3 fields".
This is :mod:`spacr.model_compare`'s harness with one model instead of two:
the same :func:`~spacr.model_compare.load_fields` reader, the same
:class:`~spacr.model_compare.ModelConfig` (so the same arguments are
honoured and the same ones reported as ignored), the same
:func:`~spacr.model_compare.segment_with_cellpose` backend, and the same
:mod:`spacr.seg_qc` scorecards. To put two models side by side use
:func:`compare_entries`, which calls
:func:`~spacr.model_compare.compare_models` proper.
The checkpoint is checked before it is loaded, so a missing or corrupt file
fails with its own name in the message rather than a torch ``KeyError``.
:param entry: the model to run.
:param images: fields already in memory; None loads them from ``source``.
:param source: a folder of fields (``.tif`` / ``.png`` / ``.npy`` /
``.npz``), read by :func:`spacr.model_compare.load_fields`.
:param n_fields: how many fields to take from ``source``.
:param field_names: names for the rows.
:param segment_fn: ``fn(images, config) -> masks``; defaults to
:func:`spacr.model_compare.segment_with_cellpose`. This is the seam the
GUI and the tests use, and the reason no test here loads Cellpose.
:param settings: eval overrides (``diameter``, ``flow_threshold``, …).
:param object_type: what is being segmented, for the seg_qc scorecards.
:param qc: score the masks with :mod:`spacr.seg_qc`.
:param keep_images: keep images and masks on the result for a GUI to draw.
:param channel: index into the last axis for multi-channel fields.
:param progress: ``fn(message, done, total)``.
:returns: a :class:`BenchmarkResult`.
:raises ModelUnreadable: when the checkpoint is missing or not a checkpoint.
:raises ValueError: when there is no field, or the model returned the wrong
number of masks.
"""
mc = _model_compare()
if images is None:
if source is None:
raise ValueError(
"benchmark needs either images= or source= (a folder of fields)")
field_names, images = mc.load_fields(source, n_fields=n_fields,
channel=channel)
fields = [np.asarray(image) for image in images]
if not fields:
raise ValueError("no field to benchmark: pass at least one image")
names = ([str(n) for n in field_names] if field_names is not None
else [f"field_{i:04d}" for i in range(len(fields))])
if len(names) != len(fields):
raise ValueError(
f"got {len(names)} field name(s) for {len(fields)} field(s)")
config = config_for(entry, settings)
notes = list(entry.notes) + list(config.notes())
if not entry.provenance_known:
notes.append(
f"{entry.name} does not record what it was trained on, so a good "
f"score here says it works on these fields and nothing more.")
if entry.path:
inspect_checkpoint(entry.path)
total_steps = 2
def _tick(message: str, done: int) -> None:
"""Report one benchmark milestone through the captured callback.
:param message: stage description for the progress display.
:param done: completed-step index from zero through two.
:returns: None. When a callback was supplied it receives the message,
completed index, and captured total of two; otherwise this is a
no-op.
"""
if progress is not None:
progress(message, done, total_steps)
_tick(f"Segmenting {len(fields)} field(s) with {entry.name}…", 0)
run = segment_fn if segment_fn is not None else _default_segmenter(
[entry], mc)
started = time.perf_counter()
produced = list(run(fields, config))
seconds = time.perf_counter() - started
if len(produced) != len(fields):
raise ValueError(
f"{entry.name} returned {len(produced)} mask(s) for "
f"{len(fields)} field(s)")
masks = [mc._as_labels(m) for m in produced]
_tick("Scoring masks…", 1)
scores = mc._score(masks, names, object_type) if qc else [None] * len(fields)
rows = [
FieldBenchmark(
field=names[i],
n_objects=int(np.unique(masks[i]).size - (1 if (masks[i] == 0).any()
else 0)),
severity=scores[i].severity if scores[i] else "-",
flags=tuple(scores[i].flags) if scores[i] else (),
note=scores[i].note if scores[i] else "",
)
for i in range(len(fields))
]
_tick("Done", 2)
return BenchmarkResult(
entry=entry,
fieldset=fieldset_id(names, fields),
fieldset_label=_fieldset_label(names, source),
rows=rows,
seconds=seconds,
honoured=config.honoured_parameters(),
ignored=config.ignored_parameters(),
notes=notes,
masks=masks if keep_images else [],
images=fields if keep_images else [],
object_type=object_type,
)
def _cellpose3_segment(model: str, images: Sequence[Any],
config: Any) -> List[np.ndarray]:
"""Segment ``images`` with a Cellpose 3 model, in the Cellpose 3 backend.
spaCR's Cellpose 4 loads a Cellpose 3 checkpoint without complaint and
then segments nonsense with it, so a ``cellpose3`` row is never handed
to it: the benchmark and the A/B comparison run it where Mask generation
does, in the Cellpose 3 environment, with the config's eval settings.
:param model: a Cellpose 3 model name or checkpoint path.
:param images: the fields.
:param config: the :class:`spacr.model_compare.ModelConfig` for the side.
:returns: one integer label image per field.
:raises ImportError: when the Cellpose 3 backend is not installed.
"""
from ._segmentation_backends import _load_backend
backend = _load_backend("cellpose3", model_name=model)
masks, _flows, _styles = backend.eval(
[np.asarray(image, dtype=np.float32) for image in images],
**config.eval_kwargs())
return [np.asarray(mask).astype(np.int32) for mask in masks]
def _default_segmenter(entries: Sequence[ModelEntry], mc: Any) -> Callable:
"""``fn(images, config) -> masks`` that runs each entry where it belongs.
A ``cellpose3`` entry goes to the Cellpose 3 backend, told apart by the
model its config names; everything else to spaCR's own Cellpose,
exactly as before.
"""
cellpose3 = {str(e.path or e.name) for e in entries
if e.kind == "cellpose3"}
def _segment(images, config):
"""Segment ``images`` with the backend the config's model belongs to."""
model = str(getattr(config, "model", "") or "")
if model in cellpose3:
return _cellpose3_segment(model, images, config)
return mc.segment_with_cellpose(images, config)
return _segment
def _fieldset_label(names: Sequence[str], source: Any) -> str:
""""3 field(s) from …: a, b, c" — what a group header says."""
where = f" from {os.fspath(source)}" if isinstance(
source, (str, os.PathLike)) else ""
listed = ", ".join(str(n) for n in list(names)[:4])
if len(names) > 4:
listed += f", … (+{len(names) - 4})"
return f"{len(names)} field(s){where}: {listed}"
[docs]
def compare_entries(entry_a: ModelEntry, entry_b: ModelEntry,
images: Optional[Sequence[Any]] = None,
source: Any = None, n_fields: int = DEFAULT_N_FIELDS,
field_names: Optional[Sequence[str]] = None,
settings_a: Optional[Mapping[str, Any]] = None,
settings_b: Optional[Mapping[str, Any]] = None,
**kwargs: Any):
"""Compare two zoo entries on the same fields.
Straight delegation to :func:`spacr.model_compare.compare_models` — the
metrics, the split/merge attribution and the "neither model is ground
truth" wording all come from there, unchanged. This function's only job is
turning two :class:`ModelEntry` objects into two
:class:`~spacr.model_compare.ModelConfig` objects.
:param entry_a: the A side.
:param entry_b: the B side.
:param images: fields already in memory; None loads them from ``source``.
:param source: a folder of fields.
:param n_fields: how many fields to take from ``source``.
:param field_names: names for the rows.
:param settings_a: eval overrides for A.
:param settings_b: eval overrides for B.
:param kwargs: forwarded to :func:`spacr.model_compare.compare_models`.
:returns: a :class:`spacr.model_compare.ComparisonReport`.
"""
mc = _model_compare()
if images is None:
if source is None:
raise ValueError(
"compare_entries needs either images= or source=")
field_names, images = mc.load_fields(source, n_fields=n_fields)
for entry in (entry_a, entry_b):
if entry.path:
inspect_checkpoint(entry.path)
kwargs.setdefault("segment_fn", _default_segmenter([entry_a, entry_b], mc))
return mc.compare_models(
images,
config_for(entry_a, settings_a),
config_for(entry_b, settings_b),
field_names=field_names,
**kwargs)
[docs]
def group_by_fieldset(results: Sequence[BenchmarkResult]
) -> Dict[str, List[BenchmarkResult]]:
"""Bucket benchmarks by the field set they ran on, first-seen order.
:param results: benchmarks.
:returns: ``{fieldset_id: [results]}``.
"""
groups: Dict[str, List[BenchmarkResult]] = {}
for result in results:
groups.setdefault(result.fieldset, []).append(result)
return groups
def _rank_value(result: BenchmarkResult, key: str) -> Tuple:
"""Build the sort key for one benchmark result.
``nan`` sorts LAST rather than first: a model nobody scored is not the
best model.
:param result: the benchmark result.
:param key: what to rank by.
:returns: the sort key.
:raises ValueError: for an unknown key. There is deliberately no
accuracy key -- a benchmark here has no ground truth, so a column
sorting models by "score" would be inventing one.
"""
if key == "qc":
score = result.qc_score
return (-(score if score == score else -1.0), result.seconds,
result.entry.name)
if key == "seconds":
return (result.seconds, result.entry.name)
raise ValueError(
f"rank key must be one of {tuple(RANK_KEYS)}, got {key!r}. There is "
f"deliberately no accuracy key: a benchmark here has no ground truth, "
f"so a column that sorted models by 'score' would be inventing one.")
[docs]
def rank(results: Sequence[BenchmarkResult],
key: str = DEFAULT_RANK_KEY) -> List[BenchmarkResult]:
"""Order benchmarks best-first — **within one field set only**.
A model's numbers on your three fields say nothing about its numbers on
somebody else's: different cell density, different exposure, different
magnification. Sorting results from two field sets into one list would
produce a ranking that looks valid but compares unrelated numbers, so this
function raises instead.
:param results: benchmarks, all from the same field set.
:param key: one of :data:`RANK_KEYS`.
:returns: the results, best first.
:raises IncomparableBenchmarks: when the results span more than one field
set. Use :func:`rank_groups` or :func:`format_benchmarks`, which group
and label instead.
:raises ValueError: on an unknown ``key``.
"""
groups = group_by_fieldset(results)
if len(groups) > 1:
detail = "; ".join(
f"{members[0].fieldset_label} "
f"[{', '.join(m.entry.name for m in members)}]"
for members in groups.values())
raise IncomparableBenchmarks(
f"refusing to rank {len(results)} benchmark(s) that ran on "
f"{len(groups)} different field sets — a score on one set of "
f"fields says nothing about a score on another, so sorting them "
f"together would invent a ranking. The sets were: {detail}. Use "
f"rank_groups() or format_benchmarks() to rank within each set.")
return sorted(results, key=lambda r: _rank_value(r, key))
[docs]
def rank_groups(results: Sequence[BenchmarkResult],
key: str = DEFAULT_RANK_KEY
) -> Dict[str, List[BenchmarkResult]]:
"""Rank inside each field set, keeping the sets apart.
:param results: benchmarks from any number of field sets.
:param key: one of :data:`RANK_KEYS`.
:returns: ``{fieldset_id: [results, best first]}``.
"""
return {fieldset: sorted(members, key=lambda r: _rank_value(r, key))
for fieldset, members in group_by_fieldset(results).items()}
def _render_table(rows: Sequence[Sequence[str]],
header: Sequence[str]) -> List[str]:
"""A fixed-width text table.
:func:`spacr.model_compare._render_table`, reused so the zoo's console
output is the same shape as the comparison's rather than a second table
style two lines apart in the same terminal.
"""
return _model_compare()._render_table(rows, header)
_ZOO_COLUMNS = (
("model", lambda e: e.name),
("kind", lambda e: e.kind),
("source", lambda e: e.source),
("v", lambda e: e.version),
("size", lambda e: _human_bytes(e.size_bytes)),
("checksum", lambda e: e.checksum_state),
("trained on", lambda e: _shorten(e.trained_on, 46)),
("trained by", lambda e: _shorten(e.trained_by, 22)),
)
SCORECARD_ROWS = (
("train", "n_train"), ("train obj.", "train_objects"),
("test", "n_test"), ("test obj.", "test_objects"), ("CV", "cv"),
("F1 @ IoU 0.5", "f1"), ("AJI", "aji"), ("Dice", "dice"),
("final train loss", "train_loss"), ("final val loss", "val_loss"),
("best epoch", "best_epoch"),
)
def _curated_scorecard_rows():
"""Name the six approved measurements without changing one-pixel semantics.
:returns: Translated labels paired with exact approved metric keys.
"""
from .qt.i18n import tr
return ((tr("F1 @ IoU 0.5"), "f1"), (tr("Boundary F1 (2 px)"), "boundary_f1_2px"),
(tr("Precision"), "precision"), (tr("Recall"), "recall"),
(tr("AJI"), "aji"), (tr("Dice"), "dice"))
def _curated_scorecard_lines(metrics):
"""Render the accepted six-row evaluation and its explicit holdout identity.
:param metrics: Measurements attached by the trusted builtin constructor.
:returns: Six matched-stock lines followed by dataset, counts and tolerance.
"""
from .qt.i18n import tr
lines = []
for label, key in _curated_scorecard_rows():
row = metrics[key]
if row["vanilla"] is None:
lines.append(tr("{metric}: {value:.4f}", metric=label, value=row["finetuned"]))
else:
lines.append(tr("{metric}: {value:.4f} (stock {stock:.4f}, difference {delta:+.4f})",
metric=label, value=row["finetuned"], stock=row["vanilla"],
delta=row["delta"]))
provenance = metrics["scorecard_provenance"]
lines.append(tr("Evaluation: {dataset}; {fields} fields, {objects} objects; boundary tolerance {pixels} px; stock {stock}.",
dataset=metrics["holdout"], fields=metrics["n_fields"], objects=metrics["n_objects"],
pixels=provenance["boundary_tolerance_px"], stock=provenance["stock_model"]))
return lines
[docs]
def scorecard_html(entry) -> str:
"""The model's scorecard as an HTML table, for a tooltip.
The table is the same one the model card prints, so two models' numbers
can be compared without reading prose. When an entry publishes no
metrics this returns ``""``, and the Zoo screen's tooltip shows the
entry's prose description instead.
Trusted builtin cards include all six accepted matched-stock rows and
their evaluation identity; boundary F1 uses a two-pixel tolerance.
Other entries keep their published fields and existing interpretation.
Metrics that hold none of the scorecard's keys -- a free-form note, a
training loss under a name of its own -- are not a scorecard either, and
return nothing too: a table of eleven "not recorded" rows would replace a
two-line note that said something.
:param entry: a catalogue entry, or anything with ``metrics`` and a name.
:returns: the HTML table, or ``""`` when there is no scorecard to show.
"""
metrics = dict(getattr(entry, "metrics", None) or {})
detector_rows = tuple((label, key) for label, key in
(("mAP50", "mAP50"), ("mAP50–95", "mAP50_95"))
if key in metrics)
curated = bool(getattr(entry, "_curated_card", False))
display_rows = (_curated_scorecard_rows() if curated else SCORECARD_ROWS + detector_rows)
if not any(key in metrics for _label, key in display_rows):
return ""
name = (getattr(entry, "display_name", "") or getattr(entry, "name", "")
or getattr(entry, "key", ""))
def cell(value):
"""A scorecard value as shown, or "not recorded" when blank."""
if value is None:
return "not recorded"
return value if str(value).strip() else "not recorded"
rows = []
for label, key in display_rows:
stock = metrics.get(f"stock_{key}", "") if key in ("f1", "aji", "dice") else ""
value = metrics.get(key, "")
if isinstance(value, Mapping):
stock = value.get("vanilla")
value = value.get("finetuned")
if isinstance(value, float):
value = f"{value:.4f}"
if isinstance(stock, float):
stock = f"{stock:.4f}"
rows.append(
f"<tr><td>{label}</td>"
f"<td align='right'><b>{cell(value)}</b></td>"
f"<td align='right'>{cell(stock) if stock or key in ('f1','aji','dice') else ''}</td></tr>")
trained = getattr(entry, "trained_on", "") or ""
card = getattr(entry, "model_card_url", "")
return (f"<p><b>{name}</b></p>"
"<table cellspacing='0' cellpadding='3'>"
"<tr><th align='left'></th><th align='right'>this model</th>"
"<th align='right'>stock</th></tr>"
+ "".join(rows) + "</table>"
+ ("<p>" + html.escape(_curated_scorecard_lines(metrics)[-1]) + "</p>"
if curated else "")
+ (f"<p>{trained[:200]}</p>" if trained else "")
+ (f"<p>{card}</p>" if card else ""))
_UNRECORDED_PREFIXES = ("not recorded", "not applicable", "not measured")
def _recorded_text(value: Any) -> str:
"""A recorded metric string, or ``""`` when it says nothing was measured."""
if isinstance(value, Mapping) or value is None or isinstance(value, bool):
return ""
text = str(value).strip()
if not text or text.lower().startswith(_UNRECORDED_PREFIXES):
return ""
return text
[docs]
def scorecard_state(entry: Any) -> str:
"""How far a model's scorecard has come, as one word.
``"reviewed"``
a trusted builtin card: six matched-stock measurements on the
curated hold-out set, bound to the exact published checkpoint.
``"held_out"``
structured measurements on a named hold-out set, not yet a reviewed
card.
``"recorded"``
numbers recorded with the model when it was published, kept as
written; no reviewed held-out comparison yet.
``"not_evaluated"``
no measurement at all.
:param entry: a :class:`ModelEntry`, or anything with ``metrics``.
:returns: one of the four words above.
"""
if getattr(entry, "_curated_card", False):
return "reviewed"
metrics = dict(getattr(entry, "metrics", None) or {})
from .scorecard import scorecard_is_present
if scorecard_is_present(metrics):
return "held_out"
if any(_recorded_text(metrics.get(key)) for key in ("f1", "aji", "dice", "mAP50", "mAP50_95")):
return "recorded"
return "not_evaluated"
[docs]
def scorecard_state_text(entry: Any) -> str:
"""The sentence a spaCR model's tooltip and card give for an unfinished scorecard.
Only models spaCR publishes get one; a stock, third-party or local model
keeps its own description.
:param entry: a zoo row.
:returns: the translated sentence for a ``"recorded"`` or
``"not_evaluated"`` :func:`scorecard_state`, else ``""``.
"""
if not (_is_spacr_published(entry) or is_spacr_owned(entry)):
return ""
state = scorecard_state(entry)
from .qt.i18n import tr
if state == "not_evaluated":
return tr("Scorecard: not yet evaluated. No held-out comparison against the "
"stock model has been measured for this model.")
if state == "recorded":
return tr("Scorecard: these numbers were recorded with the model when it was "
"published; a reviewed held-out comparison against the stock model "
"is not yet available.")
return ""
def _scorecard_state_html(entry: Any) -> str:
""":func:`scorecard_state_text` as an HTML paragraph, or ``""``."""
text = scorecard_state_text(entry)
return f"<p><i>{html.escape(text)}</i></p>" if text else ""
def _number_or_none(text: str) -> Optional[float]:
"""``float(text)`` for a plain number, else ``None``."""
try:
return float(text)
except (TypeError, ValueError):
return None
[docs]
def scorecard_export_rows(entries: Optional[Iterable[Any]] = None) -> List[Dict[str, Any]]:
"""Every model's scorecard as long-format rows, one per metric.
The matched-stock export: the measurement, the stock model's value on
the same fields and their difference, with the evaluation set and its
counts. Values are copied from the entries, never recomputed. A model
with no measurement still gets one row, whose ``state`` says so.
:param entries: zoo rows; defaults to every builtin declaration, bound
as the Model Zoo binds it.
:returns: dicts with the keys ``key``, ``name``, ``state``, ``metric``,
``this_model``, ``stock``, ``difference``, ``baseline``,
``evaluation_set``, ``evaluation_version``, ``n_fields``,
``n_objects``, ``boundary_tolerance_px`` and ``sha256``.
"""
if entries is None:
entries = [_builtin_entry_from_mapping(record) for record in BUNDLED_REMOTE_MODELS]
rows: List[Dict[str, Any]] = []
for entry in entries:
metrics = dict(getattr(entry, "metrics", None) or {})
state = scorecard_state(entry)
base = {"key": str(getattr(entry, "key", "") or ""),
"name": str(getattr(entry, "name", "") or ""), "state": state}
tail = {"sha256": str(getattr(entry, "sha256", "") or "")}
provenance = metrics.get("scorecard_provenance") or {}
evaluation = {"baseline": str(provenance.get("stock_model") or ""),
"evaluation_set": str(metrics.get("holdout") or ""),
"evaluation_version": str(metrics.get("holdout_version") or ""),
"n_fields": metrics.get("n_fields", ""),
"n_objects": metrics.get("n_objects", "")}
found: List[Dict[str, Any]] = []
if state in ("reviewed", "held_out"):
names = ([key for _label, key in _curated_scorecard_rows()] if state == "reviewed"
else [key for key, value in metrics.items() if isinstance(value, Mapping)
and value.get("finetuned") is not None])
for name in names:
value = metrics[name]
found.append({"metric": name, "this_model": value.get("finetuned"),
"stock": value.get("vanilla"), "difference": value.get("delta"),
"boundary_tolerance_px": (provenance.get("boundary_tolerance_px", "")
if name == "boundary_f1_2px" else "")})
elif state == "recorded":
evaluation["n_fields"] = _recorded_text(metrics.get("n_test"))
evaluation["n_objects"] = _recorded_text(metrics.get("test_objects"))
for name in ("f1", "aji", "dice", "mAP50", "mAP50_95"):
mine = _recorded_text(metrics.get(name))
if not mine:
continue
stock = _recorded_text(metrics.get(f"stock_{name}"))
a, b = _number_or_none(mine), _number_or_none(stock)
found.append({"metric": name, "this_model": a if a is not None else mine,
"stock": b if b is not None else stock,
"difference": (round(a - b, 6) if a is not None and b is not None
else None),
"boundary_tolerance_px": ""})
if not found:
found.append({"metric": "", "this_model": None, "stock": None,
"difference": None, "boundary_tolerance_px": ""})
rows.extend({**base, **row, **evaluation, **tail} for row in found)
return rows
_BENCH_COLUMNS = (
("field", lambda r: r.field),
("objects", lambda r: str(r.n_objects)),
("seg_qc", lambda r: r.severity),
("flags", lambda r: ", ".join(r.flags) if r.flags else "-"),
)
MITONET_KEY = "mitonet_mini_v1"
_MITONET_MEAN = 0.57571
_MITONET_STD = 0.12765
_MITONET_STRIDE = 128
_MITONET_TILE = 1024
_MITONET_HALO = 128
_MITONET_MODELS: Dict[str, Any] = {}
_MITONET_VERIFIED: Dict[str, Tuple[int, int]] = {}
_MITONET_LOCK = threading.Lock()
def _mitonet_record() -> Mapping[str, Any]:
"""The :data:`BUNDLED_REMOTE_MODELS` record of :data:`MITONET_KEY`."""
return next(record for record in BUNDLED_REMOTE_MODELS
if record.get("key") == MITONET_KEY)
def _verified_once(path: Path, sha256: str) -> bool:
"""Whether ``path`` holds the bytes ``sha256`` names, hashing it once.
The digest of a file that matched is remembered by its size and change
time, so a second segmentation does not read 110 MB again.
"""
try:
stat = path.stat()
except OSError:
return False
stamp = (int(stat.st_size), int(stat.st_mtime_ns))
if _MITONET_VERIFIED.get(str(path)) == stamp:
return True
if sha256_file(path) != sha256:
return False
_MITONET_VERIFIED[str(path)] = stamp
return True
[docs]
def mitonet_weights(dest: Any = None, download: bool = True,
opener: Optional[Callable[[str], Any]] = None,
progress: Optional[Callable[[int, int], None]] = None,
cancel: Optional[Callable[[], bool]] = None) -> Path:
"""The MitoNet v1 mini weights on this machine, downloaded on first use.
The file is the TorchScript checkpoint of Zenodo record 6861565
(CC-BY-4.0; cite Conrad and Narayan, Cell Systems 2023). It is cached
in the folder the Model Zoo downloads into, ``~/.spacr/models`` unless
``dest`` names another, and a cached copy is used only when its sha256
is the published one. A download goes through :func:`fetch`, so it is
checksummed and written atomically.
:param dest: the folder to look in and download into.
:param download: fetch the weights when no verified copy is here.
:param opener: passed to :func:`fetch`; tests replace the network.
:param progress: passed to :func:`fetch`.
:param cancel: passed to :func:`fetch`.
:returns: the path of the verified weights.
:raises ModelUnreadable: when no verified copy is here and ``download``
is False.
:raises ChecksumMismatch: when the downloaded bytes are not the
published ones.
"""
entry = _entry_from_mapping(_mitonet_record())
folder = Path(dest).expanduser() if dest is not None else _model_cache_root()
cached = folder / entry.name
candidates = [cached]
if folder.is_dir():
candidates.extend(sorted(folder.glob(
f"{cached.stem}_v*{cached.suffix}")))
for path in candidates:
if path.is_file() and _verified_once(path, entry.sha256):
return path
if not download:
raise ModelUnreadable(
f"the MitoNet weights are not in {folder}; they download from "
f"{entry.uri} on first use")
LOG.info("Downloading %s into %s", entry.name, folder)
path = fetch(entry, folder, opener=opener, progress=progress,
cancel=cancel)
_verified_once(path, entry.sha256)
return path
def _mitonet_model(path: Path, device: str) -> Any:
"""The TorchScript network at ``path`` on ``device``, loaded once."""
import torch
key = f"{path}|{device}"
with _MITONET_LOCK:
model = _MITONET_MODELS.get(key)
if model is None:
model = torch.jit.load(str(path), map_location=device).eval()
_MITONET_MODELS[key] = model
return model
def _mitonet_input(image: Any) -> np.ndarray:
"""``image`` as one float32 plane on MitoNet's 0..1 scale.
An 8-bit image is divided by 255, as the network was trained; a float
image already within 0..1 is kept; anything else is stretched from its
0.5th to its 99.5th percentile. A colour or multichannel image is
averaged over its last axis.
"""
array = np.asarray(image)
if array.ndim == 3 and array.shape[-1] <= 4:
array = array.mean(axis=-1) if array.dtype != np.uint8 else (
array.mean(axis=-1).astype(np.uint8))
if array.ndim != 2:
raise ValueError(
f"MitoNet segments one 2-D image, not shape {array.shape}")
if array.dtype == np.uint8:
return array.astype(np.float32) / 255.0
plane = array.astype(np.float32)
finite = plane[np.isfinite(plane)]
if not finite.size:
return np.zeros(plane.shape, np.float32)
if np.issubdtype(array.dtype, np.floating) and finite.min() >= 0.0 \
and finite.max() <= 1.0:
return np.nan_to_num(plane, nan=0.0)
low, high = np.percentile(finite, (0.5, 99.5))
if high <= low:
return np.zeros(plane.shape, np.float32)
return np.clip((np.nan_to_num(plane, nan=low) - low) / (high - low),
0.0, 1.0).astype(np.float32)
def _mitonet_heads(model: Any, plane: np.ndarray, device: str
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Run the network once over ``plane``: probability, centres, offsets.
The plane is padded by reflection to a multiple of 128 pixels, the
network's stride, and the outputs are cut back to the plane's size.
"""
import torch
height, width = plane.shape
pad_h, pad_w = -height % _MITONET_STRIDE, -width % _MITONET_STRIDE
mode = "reflect" if min(height, width) > 1 else "edge"
padded = np.pad((plane - _MITONET_MEAN) / _MITONET_STD,
((0, pad_h), (0, pad_w)), mode=mode).astype(np.float32)
with torch.no_grad():
out = model(torch.from_numpy(padded)[None, None].to(device))
logits = out["sem_logits"]
if logits.shape[1] == 1:
semantic = torch.sigmoid(logits)[0, 0]
else:
semantic = torch.softmax(logits, dim=1)[0, 1]
semantic = semantic[:height, :width].float().cpu().numpy()
centres = out["ctr_hmp"][0, 0, :height, :width].float().cpu().numpy()
offsets = out["offsets"][0, :, :height, :width].float().cpu().numpy()
return semantic, centres, offsets
def _tile_starts(size: int, tile: int) -> List[int]:
"""Where the tiles along an axis of ``size`` pixels begin."""
return list(range(0, max(1, size), tile)) if size > tile else [0]
[docs]
def mitonet_instances(semantic: Any, centres: Any, offsets: Any, *,
threshold: float = 0.5,
centre_threshold: float = 0.1,
centre_window: int = 7,
min_size: int = 64,
chunk: int = 1 << 16) -> np.ndarray:
"""Turn MitoNet's three output heads into labelled mitochondria.
The post-processing of empanada's Panoptic-DeepLab models, ported so
spaCR does not depend on empanada-dl. A pixel is mitochondrion when
``semantic`` reaches ``threshold``. Instance centres are the local
maxima of ``centres`` in a ``centre_window`` square that reach
``centre_threshold``. Each mitochondrion pixel votes, through its
``offsets`` (row, column), for the place its centre is, and joins the
nearest centre found. Instances smaller than ``min_size`` pixels are
dropped and the rest are numbered 1..n in raster order of their
centres.
:param semantic: (rows, columns) foreground probability.
:param centres: (rows, columns) centre heatmap.
:param offsets: (2, rows, columns) offsets from each pixel to its
centre, in pixels, rows first.
:param threshold: the foreground probability cut.
:param centre_threshold: the lowest heatmap value a centre may have.
:param centre_window: the side of the square a centre is the maximum of.
:param min_size: the smallest instance kept, in pixels.
:param chunk: foreground pixels assigned per step, which bounds memory.
:returns: an int32 label image shaped like ``semantic``.
"""
from scipy import ndimage
semantic = np.asarray(semantic, dtype=np.float32)
centres = np.asarray(centres, dtype=np.float32)
offsets = np.asarray(offsets, dtype=np.float32)
if semantic.ndim != 2 or centres.shape != semantic.shape or \
offsets.shape != (2,) + semantic.shape:
raise ValueError(
"semantic and centres must be (rows, columns) and offsets "
"(2, rows, columns) of the same size")
labels = np.zeros(semantic.shape, np.int32)
foreground = semantic >= float(threshold)
pooled = ndimage.maximum_filter(
centres, size=max(1, int(centre_window)), mode="constant",
cval=-np.inf)
peaks = np.argwhere((centres == pooled)
& (centres >= float(centre_threshold)))
if not len(peaks) or not foreground.any():
return labels
peaks = peaks.astype(np.float32)
rows, cols = np.nonzero(foreground)
owner = np.empty(rows.size, np.int64)
peak_norm = (peaks ** 2).sum(axis=1)
step = max(1, int(chunk))
for start in range(0, rows.size, step):
r = rows[start:start + step]
c = cols[start:start + step]
votes = np.stack([r + offsets[0, r, c], c + offsets[1, r, c]], 1)
distance = (peak_norm[None, :] - 2.0 * votes @ peaks.T)
owner[start:start + step] = distance.argmin(axis=1)
counts = np.bincount(owner, minlength=len(peaks))
keep = counts >= max(1, int(min_size))
number = np.zeros(len(peaks), np.int32)
number[keep] = np.arange(1, int(keep.sum()) + 1, dtype=np.int32)
labels[rows, cols] = number[owner]
return labels
[docs]
def segment_mitonet(image: Any, weights: Any = None, *,
threshold: float = 0.5,
centre_threshold: float = 0.1,
min_size: int = 64,
device: Optional[str] = None,
tile: int = _MITONET_TILE,
halo: int = _MITONET_HALO,
on_tile: Optional[Callable[[int, int], None]] = None
) -> np.ndarray:
"""Segment the mitochondria of a TEM image with MitoNet v1 mini.
The network runs over tiles of ``tile`` pixels with ``halo`` pixels of
context on every side, so a large field never needs the whole image on
the network at once, and the three heads are stitched before
:func:`mitonet_instances` turns them into labels. MitoNet was trained
on electron micrographs with dark membranes; it is not a fluorescence
model.
:param image: a 2-D image (8-bit is read as the network was trained;
see :func:`mitonet_weights` for where the weights come from).
:param weights: the TorchScript file; None uses :func:`mitonet_weights`,
which downloads it on first use.
:param threshold: passed to :func:`mitonet_instances`.
:param centre_threshold: passed to :func:`mitonet_instances`.
:param min_size: passed to :func:`mitonet_instances`.
:param device: a torch device; None uses CUDA when torch sees it and
the CPU otherwise.
:param tile: the side of the core of one tile, in pixels.
:param halo: the context around each tile's core, in pixels.
:param on_tile: called as ``on_tile(done, total)`` after each tile; an
exception it raises stops the run.
:returns: an int32 label image shaped like the image's plane.
"""
import torch
plane = _mitonet_input(image)
if device is None:
device = "cuda" if torch.cuda.is_available() else "cpu"
path = Path(weights) if weights is not None else mitonet_weights()
model = _mitonet_model(path, str(device))
height, width = plane.shape
tile = max(_MITONET_STRIDE, int(tile))
halo = max(0, int(halo))
starts = [(y, x) for y in _tile_starts(height, tile)
for x in _tile_starts(width, tile)]
if len(starts) == 1:
semantic, centres, offsets = _mitonet_heads(model, plane, device)
if on_tile is not None:
on_tile(1, 1)
else:
semantic = np.zeros((height, width), np.float32)
centres = np.zeros((height, width), np.float32)
offsets = np.zeros((2, height, width), np.float32)
for done, (y0, x0) in enumerate(starts, 1):
y1, x1 = min(height, y0 + tile), min(width, x0 + tile)
top, left = max(0, y0 - halo), max(0, x0 - halo)
bottom, right = min(height, y1 + halo), min(width, x1 + halo)
part = _mitonet_heads(model, plane[top:bottom, left:right],
device)
core = (slice(y0 - top, y1 - top), slice(x0 - left, x1 - left))
semantic[y0:y1, x0:x1] = part[0][core]
centres[y0:y1, x0:x1] = part[1][core]
offsets[:, y0:y1, x0:x1] = part[2][(slice(None),) + core]
if on_tile is not None:
on_tile(done, len(starts))
return mitonet_instances(semantic, centres, offsets,
threshold=threshold,
centre_threshold=centre_threshold,
min_size=min_size)
TIMEPOSE_CONTRACT = "spacr.timepose/1"
_TIMEPOSE_VERIFIED: Dict[str, Tuple[int, int]] = {}
_TIMEPOSE_MODELS: Dict[Tuple[str, str, str], Any] = {}
_TIMEPOSE_LOCK = threading.Lock()
class _TimeposeAccessError(ModelZooError):
"""No Hugging Face token, or the token's account cannot read the repository."""
def _timepose_record(key: str) -> Mapping[str, Any]:
"""The :data:`BUNDLED_REMOTE_MODELS` record of a Timepose release.
Every Timepose version is one record of kind ``timepose``: its repository,
pinned revision, the sha256 of every file and the adapter module that
exposes ``load_model`` and ``segment_track_movie``. A new version needs
only a new record.
:param key: the zoo key, also the ``timelapse_mode`` choice it is
offered as.
:raises ModelZooError: for a key that is not a Timepose release.
"""
for record in BUNDLED_REMOTE_MODELS:
if record.get("key") == key and record.get("kind") == "timepose":
return record
from .qt.i18n import tr
raise ModelZooError(tr("{key} is not a Timepose model in the model zoo.",
key=key))
def _timepose_verified(path: Path, sha256: str) -> bool:
"""Whether ``path`` holds the pinned bytes, hashing each file once."""
try:
stat = path.stat()
except OSError:
return False
stamp = (int(stat.st_size), int(stat.st_mtime_ns))
if _TIMEPOSE_VERIFIED.get(str(path)) == stamp:
return True
if sha256_file(path) != sha256:
return False
_TIMEPOSE_VERIFIED[str(path)] = stamp
return True
def _timepose_access_failure(error: BaseException) -> bool:
"""Whether a Hugging Face error means no token or no access."""
if type(error).__name__ in ("GatedRepoError", "RepositoryNotFoundError",
"LocalTokenNotFoundError"):
return True
response = getattr(error, "response", None)
return getattr(response, "status_code", None) in (401, 403, 404)
def _timepose_token(token: Optional[str] = None) -> Optional[str]:
"""The Hugging Face token: ``token``, else the one huggingface_hub finds."""
if token:
return token
try:
from huggingface_hub import get_token
except ImportError:
return os.environ.get("HF_TOKEN") or None
return get_token()
def _timepose_files(key: str, dest: Any = None,
token: Optional[str] = None,
downloader: Optional[Callable[..., Any]] = None,
progress: Optional[Callable[[int, int], None]] = None
) -> Path:
"""Download a Timepose release once and verify every file against its pin.
The files go to ``<cache>/timepose/<key>/<revision>``. A file whose sha256
is not the pinned one is deleted and nothing is loaded from the release.
:param key: the zoo key.
:param dest: the cache root; spaCR's model cache by default.
:param token: a Hugging Face token; the logged-in one by default.
:param downloader: ``hf_hub_download``-compatible callable; tests pass a
fake one.
:param progress: ``fn(done_files, total_files)``.
:returns: the folder holding the verified release.
:raises _TimeposeAccessError: with no token, or no access to the private
repository.
:raises ChecksumMismatch: when a downloaded file is not the pinned one.
"""
from .qt.i18n import tr
record = _timepose_record(key)
root = Path(dest).expanduser() if dest is not None else _model_cache_root() / "timepose"
folder = root / key / str(record["revision"])[:12]
files = dict(record["files"])
missing = [name for name, sha in sorted(files.items())
if not _timepose_verified(folder / name, sha)]
if not missing:
return folder
found = _timepose_token(token)
if downloader is None:
if not found:
raise _TimeposeAccessError(tr(
"{name} is in a private Hugging Face repository ({repo}). Sign "
"in with an account that has been given access: run `hf auth "
"login` or set HF_TOKEN, then run again.",
name=record.get("display_name", key), repo=record["repo_id"]))
from huggingface_hub import hf_hub_download
downloader = hf_hub_download
folder.mkdir(parents=True, exist_ok=True)
for done, name in enumerate(missing, 1):
try:
downloader(repo_id=record["repo_id"], filename=name,
revision=record["revision"], repo_type="model",
local_dir=str(folder), token=found)
except Exception as error:
if _timepose_access_failure(error):
raise _TimeposeAccessError(tr(
"Your Hugging Face account cannot read {repo}, the private "
"repository of {name}. Ask the spaCR maintainer for "
"access, or choose another timelapse_mode.",
name=record.get("display_name", key),
repo=record["repo_id"])) from error
raise
path = folder / name
if not _timepose_verified(path, files[name]):
try:
path.unlink()
except OSError:
pass
raise ChecksumMismatch(tr(
"{file} of {name} does not match its pinned sha256 and was "
"deleted; nothing from this download was run.",
file=name, name=record.get("display_name", key)))
if progress is not None:
progress(done, len(missing))
return folder
def _timepose_device(device: Any = None) -> str:
"""``'cuda'`` when asked for or automatic with a GPU present, else ``'cpu'``."""
text = str(device or "automatic").lower()
if text not in ("automatic", "auto", "none"):
return text
try:
import torch
return "cuda" if torch.cuda.is_available() else "cpu"
except ImportError:
return "cpu"
def _load_timepose(key: str, device: Any = None,
token: Optional[str] = None, dest: Any = None,
downloader: Optional[Callable[..., Any]] = None
) -> Tuple[Any, Any]:
"""The adapter module and loaded model of a Timepose release.
The adapter ships with the weights and is imported only after every file
matched its pin. It must declare :data:`TIMEPOSE_CONTRACT` and expose
``load_model(folder, device)`` and ``segment_track_movie(model, images,
invert=..., masks=...)``, which returns ``masks`` (T, Y, X) and ``edges``
``[t, parent_label, child_label]``.
:returns: ``(adapter, model)``, cached per key, revision and device.
"""
import importlib
import importlib.util
import sys
from .qt.i18n import tr
record = _timepose_record(key)
device = _timepose_device(device)
cache_key = (key, str(record["revision"]), device)
with _TIMEPOSE_LOCK:
if cache_key in _TIMEPOSE_MODELS:
return _TIMEPOSE_MODELS[cache_key]
folder = _timepose_files(key, dest=dest, token=token,
downloader=downloader)
package, _, module_name = str(record["entry"]).partition(".")
alias = f"_spacr_{key}_{str(record['revision'])[:12]}"
if alias not in sys.modules:
spec = importlib.util.spec_from_file_location(
alias, folder / package / "__init__.py",
submodule_search_locations=[str(folder / package)])
module = importlib.util.module_from_spec(spec)
sys.modules[alias] = module
try:
spec.loader.exec_module(module)
except Exception:
sys.modules.pop(alias, None)
raise
adapter = (importlib.import_module(f"{alias}.{module_name}")
if module_name else sys.modules[alias])
if getattr(adapter, "CONTRACT", None) != record["contract"]:
raise ModelZooError(tr(
"{name} does not provide the interface this spaCR expects "
"({contract}); update spaCR.", name=record.get("display_name", key),
contract=record["contract"]))
model = adapter.load_model(str(folder), device)
_TIMEPOSE_MODELS[cache_key] = (adapter, model)
return adapter, model