diff --git a/README.md b/README.md index 0d14e65..1a283d9 100644 --- a/README.md +++ b/README.md @@ -162,6 +162,42 @@ stgpt spatho-embed --model outputs/atera_wta_breast/model --config configs/atera The spatho export writes `region_embeddings.parquet`, `region_cell_membership.parquet`, `region_molecular_summary.parquet`, `region_image_manifest.json`, `region_qc_report.json`, `evidence_manifest.json`, `structure_summary.parquet`, and `structure_embedding_summary.csv`. +## Region Auto-Annotation + +When part of a slide is annotated and the rest is not, propagate the labels +with `stgpt annotate-regions`. The command consumes a seed CSV +(`region_id, structure_label[, confidence]`), reuses the trained +`structure_head` calibrated by temperature scaling, and writes +`region_predictions.parquet` plus `auto_annotation_report.json`: + +```bash +stgpt annotate-regions \ + --config configs/atera_wta_breast_slide.yaml \ + --checkpoint outputs/.../checkpoints/last.pt \ + --seed-labels seed_labels.csv \ + --classifier both \ + --output outputs/.../auto_annotation +``` + +Same call from Python: + +```python +from stgpt.runtime import annotate_regions + +result = annotate_regions( + config="configs/atera_wta_breast_slide.yaml", + checkpoint="outputs/.../checkpoints/last.pt", + seed_labels="seed_labels.csv", + output_dir="outputs/.../auto_annotation", + classifier="both", + write_probabilities=True, +) +``` + +The contract is RFC 0002 (`docs/rfcs/0002-region-auto-annotation.md`). +Predictions falling below the probability/entropy/distance thresholds are +labelled `__abstain__` rather than guessed. + ## H&E Evidence Workflow stGPT consumes contour H&E assets built by pyXenium/spatho. It can train with the default lightweight CNN, a frozen `timm`/Hugging Face pathology encoder, or a precomputed embedding store: diff --git a/docs/rfcs/0002-region-auto-annotation.md b/docs/rfcs/0002-region-auto-annotation.md new file mode 100644 index 0000000..c858876 --- /dev/null +++ b/docs/rfcs/0002-region-auto-annotation.md @@ -0,0 +1,362 @@ +# RFC 0002: Region Auto-Annotation from Sparse Labels + +**Status:** Draft +**Author:** stGPT Team +**Created:** 2026-05-20 + +## 1. Motivation + +Pathologists annotate a small number of structurally meaningful regions on +one slide. Marking every region by hand on the same slide, and on every +future slide, is not feasible — especially for structures that are hard to +recognize visually without molecular context. + +stGPT already trains a multimodal region encoder over H&E, gene +expression, spatial coordinates, and structure-context tokens, and emits +per-region embeddings plus an auxiliary structure-classification head. +This RFC turns that latent capability into a stable user-facing workflow: + +> Given an annotated subset of regions on a slide (or set of slides), +> produce confidence-scored structure labels for the remaining regions +> on the same slide, and later for regions on a held-out slide. + +The product claim is intentionally narrow: + +```text +stGPT propagates sparse expert structure labels to unannotated regions +via measured morpho-molecular embeddings, with per-region confidence, +abstain semantics, and full provenance. +``` + +This RFC is annotation propagation, not online learning. The trained +checkpoint is frozen; only a lightweight classifier on top of region +embeddings is fit per case. + +## 2. Non-goals + +- Not an unsupervised structure-discovery method. The 5 classes (or + whatever the user provides) come from a labeled `structure_assignments` + contract. +- Not a clinical classifier. Output is auditable evidence with confidence + and abstain flags, never a diagnosis. +- Not a replacement for `structure_head` retraining. This is a + per-case calibration layer; full retraining still happens through + `stgpt train`. +- Not a cross-slide H&E-only path in v0. Phase 2 covers that, and + needs a separate image-only inference contract (see §7). +- Not a stain-normalization or contour-segmentation tool. Both are + upstream of stGPT and must be settled in pyXenium / spatho first. + +## 3. Feasibility + +stGPT today already carries the building blocks: + +- `ImageGeneSTGPT.structure_head` (`src/stgpt/models.py:564`) is a + linear head over the region embedding, trained jointly with masked + gene reconstruction, neighborhood reconstruction, and the image-gene + contrastive loss when `structure_loss_weight > 0` and + `n_structures > 1` (`src/stgpt/training.py:445`, + `src/stgpt/config.py:165`). +- `ImageGeneSTGPT.forward` emits `structure_logits` per region + (`src/stgpt/models.py:716`), so both the embedding and a 5-way logit + vector are available at inference with no architectural change. +- `embed-regions` already writes `region_embeddings.parquet` with + `structure_label`, `qc_flag`, and `emb_*` columns + (`README.md:75-82`), giving us a stable seed for the unannotated + pool. +- `evaluate` already computes `label_retrieval_metrics.csv` and + `image_gene_retrieval_metrics.csv`, which is the same machinery a + nearest-neighbor propagator needs. + +What is missing is a runtime entry point that consumes the seed labels +and writes per-region predictions for the unlabeled regions. This RFC +adds exactly that. + +## 4. Scope, in order + +**Phase 1 — same-slide propagation.** +One Xenium case, partial expert labels, predict the rest. Reuses the +existing `xenium_slide.zarr` pipeline end-to-end. This is the v0 +deliverable. + +**Phase 2 — cross-slide propagation (Xenium + H&E + Xenium).** +A second labeled or unlabeled Xenium case in the same panel. Same +artifacts, but with a domain-shift QC gate and explicit warnings. + +**Phase 3 — cross-slide propagation (H&E only).** +A target slide with no Xenium molecules. Requires an `image_only` +inference path that builds region tokens from H&E contour segmentation +alone and masks gene tokens. Separately RFC'd; flagged here so the v0 +schema reserves room for `expression_present: bool`. + +This RFC scopes implementation to Phase 1. + +## 5. Data Contract + +### 5.1 Seed labels (input) + +The annotator supplies a CSV with at minimum: + +``` +region_id,structure_label,confidence +R000123,gland,1.0 +R000456,stroma,1.0 +R000789,necrosis,0.8 +... +``` + +- `region_id` must resolve against the case's `region_embeddings.parquet` + (built by `stgpt embed-regions`). Unknown ids are rejected before + fitting. +- `structure_label` must be one of the strings declared in the + case's `structure_assignments.csv`, or a superset declared in the + config. Unknown labels are rejected. +- `confidence` is optional (defaults to 1.0). Values below a + config threshold become sample weights when fitting, never become + pseudo-labels for unlabeled regions. + +A typo or unknown label is a fatal error, not a warning. We do not +silently drop expert annotations. + +### 5.2 Region pool (input) + +The unlabeled pool defaults to "all regions in this case where the seed +file does not provide a label and `qc_flag == 'ok'`". Two opt-ins: + +- `--include-no-image`: also score regions where the H&E patch was + missing (`qc_flag == 'no_image'`). Off by default; predictions on + these regions are marked `qc_flag: no_image` in the output and the + classifier never sees them at fit time. +- `--region-ids `: restrict the pool to a caller-provided list, + useful for "score these specific regions only" workflows. + +### 5.3 Predictions (output) + +`outputs//auto_annotation/region_predictions.parquet` with one +row per region in the pool: + +| Column | Type | Description | +|---|---|---| +| `region_id` | str | Joins back to `region_embeddings.parquet` | +| `predicted_label` | str | Top-1 label, or `__abstain__` if confidence below threshold | +| `predicted_prob` | float | Top-1 calibrated probability | +| `entropy` | float | Predictive entropy across labels | +| `nearest_seed_region_id` | str | Closest annotated region in embedding space | +| `nearest_seed_distance` | float | Cosine distance to that seed | +| `qc_flag` | str | Forwarded from `region_embeddings.parquet`, plus `seed` for seed regions | +| `classifier` | str | `structure_head` or `prototype_knn` | +| `evidence_id` | str | Deterministic hash of inputs | + +Alongside it, two summary files: + +- `region_predictions_per_class.parquet` — one row per (`region_id`, + `class`) with the full probability vector. Optional, gated by + `--write-probabilities`. +- `auto_annotation_report.json` — fit/inference fingerprints, seed + counts per class, abstain rate, confusion matrix on held-out seed + folds, QC gates, warnings. + +Provenance fields on every row (model checkpoint fingerprint, config +fingerprint, seed file hash, classifier kind, threshold) live in +`auto_annotation_report.json`, not duplicated per row. + +## 6. Algorithm + +Two paths, both emitted by default for v0 so they can be compared on +the same case. + +### 6.1 Path A — `structure_head` direct (recommended primary) + +1. Load checkpoint via `ImageGeneSTGPT.from_pretrained`. +2. Run `embed-regions` if `region_embeddings.parquet` is not already + present; otherwise reuse it. +3. Re-run a forward pass over the pool to collect `structure_logits` + (already produced by the head, see `src/stgpt/models.py:716`). +4. **Temperature-scale the logits on the seed set.** Single scalar + temperature, fit by minimizing NLL on held-out seed folds. This is + the standard calibration step and avoids over-confidence on classes + the head was already biased toward at training time. +5. Probabilities → top-1 + entropy → abstain rule. + +Pros: zero new training, exactly the supervision signal we already +trained on, fast. + +Cons: only works if the checkpoint was trained with the same label +vocabulary. Adding a new class needs Path B. + +### 6.2 Path B — prototype/k-NN on the embedding (sanity check & cold start) + +1. Pull `emb_*` columns for seed and pool regions. +2. L2-normalize embeddings. +3. Compute per-class prototypes (mean of seed embeddings, weighted by + seed `confidence`). For very low seed counts (< 5 per class), use + k-NN with k=min(5, seed_count) instead. +4. Score pool regions by cosine similarity → softmax over similarities + with a fitted temperature → top-1 + entropy. + +Pros: lets the user redefine label vocabulary without retraining, +robust to small seed sets via k-NN. + +Cons: ignores `structure_head` weights; quality depends purely on how +discriminative the embedding is for the label set. + +### 6.3 Abstain rule + +A prediction becomes `__abstain__` if any of the following are true: + +- `predicted_prob < tau_p` (default 0.5, configurable). +- `entropy > tau_h` (default `log(n_classes) - 0.5`). +- `nearest_seed_distance > tau_d` measured against seed-distance + quantiles fit on the seed set itself (default 99th percentile of + seed-to-nearest-other-seed distances). + +Abstain is not a failure; it is the model saying "ask a human." The +fraction of abstained regions is logged in the report. + +### 6.4 Seed validation + +Before fitting, run a `k`-fold cross-validation on seeds alone +(default k=5, capped at seed_count_per_class - 1). Write: + +- per-class precision/recall/F1 +- macro F1 +- a small confusion matrix + +If any class has fewer than 5 seeds, the report flags `low_seed_count` +and recommends the user add more before trusting Phase 2 cross-slide +runs. + +## 7. Reservations for Phase 2 / 3 + +The v0 schema reserves these fields so Phase 2 and 3 do not break the +contract: + +- `expression_present: bool` — false for Phase 3 image-only regions. +- `source_case_id: str` — the case that produced the region, for + cross-slide propagation reports. +- `propagation_kind: enum { same_slide, cross_slide_full, cross_slide_image_only }`. + +In Phase 1 these are filled with `True`, the current case id, and +`same_slide`. + +## 8. Runtime and CLI Surface + +New runtime entry point in `stgpt.runtime`: + +```python +def annotate_regions( + *, + config: StGPTConfig | str | Path, + checkpoint: str | Path, + seed_labels: str | Path, + output_dir: str | Path, + region_ids: str | Path | None = None, + include_no_image: bool = False, + classifier: Literal["structure_head", "prototype_knn", "both"] = "both", + abstain_prob: float = 0.5, + write_probabilities: bool = False, + batch_size: int = 32, + device: str = "auto", +) -> dict[str, Any]: + """Propagate sparse expert structure labels to unannotated regions.""" +``` + +New CLI command (in `src/stgpt/cli.py`): + +``` +stgpt annotate-regions \ + --config configs/atera_wta_breast_slide.yaml \ + --checkpoint outputs/.../checkpoints/last.pt \ + --seed-labels seed_labels.csv \ + --classifier both \ + --output outputs/.../auto_annotation +``` + +`stgpt embed-regions` is invoked transparently if its outputs are not +present, so the user can run `annotate-regions` cold. + +## 9. QC Gates + +Fatal (run aborts): + +- Any seed `region_id` not present in the case. +- Any seed `structure_label` not in the declared vocabulary. +- Fewer than 2 distinct seed classes (nothing to propagate). +- Checkpoint label vocabulary mismatches the seed vocabulary under + `--classifier structure_head` (Path A is unavailable; the run + proceeds with Path B only and the report records the downgrade). + +Warning (run continues, report flags it): + +- Any class with fewer than 5 seeds. +- Abstain rate over 50% of the pool. +- The held-out seed macro-F1 below 0.6. +- Embedding QC: any `emb_*` column with > 5% NaN or > 1% zero-vectors. + +## 10. Reproducibility and Provenance + +Every artifact carries: + +- `model_fingerprint`: sha256 of the checkpoint payload bytes. +- `config_fingerprint`: sha256 of the resolved StGPTConfig as JSON. +- `seed_fingerprint`: sha256 of the seed CSV. +- `evidence_id`: sha256 of the above three concatenated, used as the + per-run identifier in spatho. + +Re-running with identical inputs must produce identical outputs; +random seed for k-fold splits is fixed (default 42, configurable). + +## 11. Testing Plan + +- A unit test that builds a tiny synthetic case (extending the smoke + fixtures), assigns labels to 30% of regions, and asserts that: + - `region_predictions.parquet` has the right columns and dtypes; + - all seed regions appear with `qc_flag == 'seed'`; + - the report's seed-fold macro F1 is above 0.5 on the synthetic + data; + - re-running with the same inputs yields a byte-identical output + file. +- A smoke CLI test: `stgpt annotate-regions --config configs/smoke.yaml + --seed-labels ` runs to completion in under 30 seconds + on CPU. + +No tests against real Atera data live in the public repo, matching +the existing convention. + +## 12. Out-of-Scope Decisions Logged Here + +- **Active learning hooks** (next batch of regions to label) are + intentionally deferred. Once the abstain set exists, picking + high-uncertainty + high-coverage seeds is straightforward, but it + is its own RFC because it affects the spatho UI. +- **Stain normalization for Phase 2/3** stays in pyXenium/spatho. The + stGPT side records which normalization was applied as a + fingerprint, but does not pick one. +- **A "merge labels" or "rename label" step** is out of scope. If the + user wants `tumor + necrosis -> tumor_or_necrosis`, they do it + upstream in the seed CSV. + +## 13. Open Questions + +1. Should `nearest_seed_distance` use cosine on the trained embedding + or on the post-`structure_head` logits? Current proposal: cosine + on embedding, because it remains meaningful when Path A is + downgraded. +2. How do we surface "two paths disagree" in `auto_annotation_report.json`? + Proposal: a `path_agreement.csv` side file with one row per + region and `agree: bool` plus both top-1s, gated by `--classifier both`. +3. For Phase 2, do we require the second slide to share the same + contour segmentation parameters with the seed slide, or do we + accept any pyXenium contour set and rely on QC? Proposal: + require, until cross-segmentation calibration is studied. + +## 14. Status + +This RFC is design-only. No code lands until the contract above is +signed off. Implementation order, once approved: + +1. `annotate_regions` Python entry point with Path A and Path B. +2. CLI command and smoke test. +3. Seed-fold report and abstain rule. +4. Real-data dry run on one Atera breast case, paired with a short + evidence note in `docs/results/`. diff --git a/src/stgpt/annotation.py b/src/stgpt/annotation.py new file mode 100644 index 0000000..3cec317 --- /dev/null +++ b/src/stgpt/annotation.py @@ -0,0 +1,881 @@ +"""Region auto-annotation: propagate sparse expert structure labels to the rest. + +This module implements RFC 0002. Given an annotated subset of regions on a +case (the "seed set"), it produces per-region predicted structure labels for +the remaining regions ("the pool"), with calibrated probability, entropy, +nearest-seed distance, and an explicit abstain flag. + +Two classifier paths are exposed and can be emitted together for comparison: + +* ``structure_head`` – reuse the trained classification head from the + checkpoint, with a single-scalar temperature fit on the seed labels. +* ``prototype_knn`` – class prototypes (or k-NN for very small seed counts) + computed on the region embedding. Independent of the head, so the same + user can supply labels that were not in the training vocabulary. +""" +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal + +import numpy as np +import pandas as pd +import torch +from torch.utils.data import DataLoader + +from .config import StGPTConfig +from .data import RegionDataset, build_training_case +from .foundation.packaging import resolve_model_checkpoint +from .models import ImageGeneSTGPT + +ABSTAIN_LABEL = "__abstain__" +SCHEMA_VERSION = "stgpt.region_auto_annotation.v0.1" + +ClassifierKind = Literal["structure_head", "prototype_knn", "both"] + + +@dataclass(frozen=True) +class AutoAnnotationResult: + predictions: Path + report: Path + probabilities: Path | None + path_agreement: Path | None + n_seed_regions: int + n_pool_regions: int + n_abstain: int + + def to_dict(self) -> dict[str, Any]: + payload: dict[str, Any] = { + "predictions": str(self.predictions), + "report": str(self.report), + "n_seed_regions": int(self.n_seed_regions), + "n_pool_regions": int(self.n_pool_regions), + "n_abstain": int(self.n_abstain), + } + if self.probabilities is not None: + payload["probabilities"] = str(self.probabilities) + if self.path_agreement is not None: + payload["path_agreement"] = str(self.path_agreement) + return payload + + +def annotate_regions( + *, + config: StGPTConfig | str | Path, + checkpoint: str | Path, + seed_labels: str | Path, + output_dir: str | Path, + region_ids: str | Path | None = None, + include_no_image: bool = False, + classifier: ClassifierKind = "both", + abstain_prob: float = 0.5, + write_probabilities: bool = False, + seed_folds: int = 5, + rng_seed: int = 42, + batch_size: int = 32, + device: str = "auto", +) -> dict[str, Any]: + """Propagate sparse expert structure labels to unannotated regions. + + Writes ``region_predictions.parquet`` and ``auto_annotation_report.json`` + under *output_dir*. See RFC 0002 for the field-by-field contract. + """ + cfg = StGPTConfig.from_file(config) if isinstance(config, (str, Path)) else config + checkpoint_path = resolve_model_checkpoint(checkpoint) + out_dir = Path(output_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + inference = _run_inference(cfg, checkpoint_path, batch_size=batch_size, device=device) + label_vocab = inference.label_vocab + region_df = inference.region_table.copy().reset_index(drop=True) + region_ids_all = region_df["region_id"].astype(str).to_numpy() + embeddings = inference.embeddings + logits = inference.structure_logits # may be None + + seeds_path = Path(seed_labels) + seed_frame = _read_seed_labels(seeds_path, region_ids_all) + label_vocab, mode = _resolve_label_vocabulary(seed_frame, label_vocab, classifier) + + seed_index_map = _build_index_map(region_ids_all) + seed_indices = np.array([seed_index_map[rid] for rid in seed_frame["region_id"].to_numpy()], dtype=np.int64) + seed_labels_int = np.array([label_vocab.index(lbl) for lbl in seed_frame["structure_label"].to_numpy()], dtype=np.int64) + seed_weights = seed_frame["confidence"].astype(np.float64).to_numpy() + + pool_mask = np.ones(len(region_df), dtype=bool) + pool_mask[seed_indices] = False + if not include_no_image and "qc_flag" in region_df.columns: + pool_mask &= region_df["qc_flag"].astype(str).to_numpy() == "ok" + if region_ids is not None: + restrict = _read_region_id_list(Path(region_ids)) + pool_mask &= np.isin(region_ids_all, list(restrict)) + + pool_indices = np.where(pool_mask)[0] + n_classes = len(label_vocab) + if n_classes < 2: + raise ValueError("Need at least 2 distinct seed structure labels to propagate.") + + selected_paths: list[str] = [] + if classifier in {"structure_head", "both"}: + if logits is None or mode != "checkpoint_vocab": + if classifier == "structure_head": + raise ValueError( + "Path A (structure_head) is unavailable: checkpoint vocabulary does not match seed labels, " + "or the checkpoint has no structure_head." + ) + else: + selected_paths.append("structure_head") + if classifier in {"prototype_knn", "both"}: + selected_paths.append("prototype_knn") + if not selected_paths: + raise ValueError("No classifier path is available for this run.") + + rng = np.random.default_rng(rng_seed) + + path_predictions: dict[str, _PathPrediction] = {} + if "structure_head" in selected_paths: + path_predictions["structure_head"] = _fit_predict_structure_head( + logits=logits, + seed_indices=seed_indices, + seed_labels=seed_labels_int, + seed_weights=seed_weights, + pool_indices=pool_indices, + n_classes=n_classes, + rng=rng, + seed_folds=seed_folds, + ) + if "prototype_knn" in selected_paths: + path_predictions["prototype_knn"] = _fit_predict_prototype( + embeddings=embeddings, + seed_indices=seed_indices, + seed_labels=seed_labels_int, + seed_weights=seed_weights, + pool_indices=pool_indices, + n_classes=n_classes, + rng=rng, + seed_folds=seed_folds, + ) + + primary_path = "structure_head" if "structure_head" in path_predictions else "prototype_knn" + primary = path_predictions[primary_path] + seed_emb_distances = _seed_to_seed_distances(embeddings, seed_indices) + distance_cutoff = float(np.quantile(seed_emb_distances, 0.99)) if seed_emb_distances.size else float("inf") + entropy_cutoff = float(np.log(max(n_classes, 2)) - 0.5) + + nearest_seed_idx, nearest_seed_dist = _nearest_seed(embeddings, seed_indices, pool_indices) + + predictions_table = _build_predictions_table( + region_ids_all=region_ids_all, + region_df=region_df, + seed_indices=seed_indices, + seed_labels_int=seed_labels_int, + pool_indices=pool_indices, + primary=primary, + primary_path=primary_path, + nearest_seed_idx=nearest_seed_idx, + nearest_seed_dist=nearest_seed_dist, + label_vocab=label_vocab, + abstain_prob=float(abstain_prob), + entropy_cutoff=entropy_cutoff, + distance_cutoff=distance_cutoff, + evidence_prefix=_evidence_prefix(cfg, checkpoint_path, seeds_path), + ) + predictions_path = out_dir / "region_predictions.parquet" + predictions_table.to_parquet(predictions_path, index=False) + + probabilities_path: Path | None = None + if write_probabilities: + probabilities_path = out_dir / "region_predictions_per_class.parquet" + _build_probabilities_table( + region_ids_all=region_ids_all, + pool_indices=pool_indices, + primary=primary, + label_vocab=label_vocab, + ).to_parquet(probabilities_path, index=False) + + path_agreement_path: Path | None = None + if classifier == "both" and len(path_predictions) == 2: + path_agreement_path = out_dir / "path_agreement.csv" + _build_path_agreement_table( + region_ids_all=region_ids_all, + pool_indices=pool_indices, + paths=path_predictions, + label_vocab=label_vocab, + ).to_csv(path_agreement_path, index=False) + + seed_cv = _seed_cross_validation( + path=primary_path, + embeddings=embeddings, + logits=logits, + seed_indices=seed_indices, + seed_labels=seed_labels_int, + seed_weights=seed_weights, + n_classes=n_classes, + seed_folds=seed_folds, + rng=rng, + ) + n_abstain = int((predictions_table["predicted_label"] == ABSTAIN_LABEL).sum()) + report = { + "schema_version": SCHEMA_VERSION, + "case_name": cfg.case_name, + "checkpoint": str(checkpoint_path), + "config_fingerprint": _config_hash(cfg), + "checkpoint_fingerprint": _sha256_path(checkpoint_path), + "seed_fingerprint": _sha256_path(seeds_path), + "label_vocab": list(label_vocab), + "label_vocab_mode": mode, + "classifier_requested": classifier, + "classifiers_used": list(path_predictions.keys()), + "primary_classifier": primary_path, + "abstain_rule": { + "abstain_prob": float(abstain_prob), + "entropy_cutoff": entropy_cutoff, + "distance_cutoff": distance_cutoff, + }, + "seed_counts_per_class": _label_counts(seed_labels_int, label_vocab), + "n_seed_regions": int(seed_indices.size), + "n_pool_regions": int(pool_indices.size), + "n_abstain": n_abstain, + "abstain_rate": float(n_abstain / max(1, pool_indices.size)), + "warnings": _collect_warnings(seed_labels_int, label_vocab, n_abstain, pool_indices.size, seed_cv), + "seed_cross_validation": seed_cv, + "rng_seed": int(rng_seed), + } + report_path = out_dir / "auto_annotation_report.json" + report_path.write_text(json.dumps(report, indent=2), encoding="utf-8") + + result = AutoAnnotationResult( + predictions=predictions_path, + report=report_path, + probabilities=probabilities_path, + path_agreement=path_agreement_path, + n_seed_regions=int(seed_indices.size), + n_pool_regions=int(pool_indices.size), + n_abstain=n_abstain, + ) + return result.to_dict() + + +# --------------------------------------------------------------------------- +# Inference: collect embeddings AND structure_logits in one forward pass +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class _InferenceOutputs: + region_table: pd.DataFrame + embeddings: np.ndarray + structure_logits: np.ndarray | None + label_vocab: tuple[str, ...] + + +def _run_inference( + cfg: StGPTConfig, + checkpoint_path: Path, + *, + batch_size: int, + device: str, +) -> _InferenceOutputs: + payload: dict[str, Any] = torch.load(checkpoint_path, map_location="cpu") + inner_cfg = cfg.model_copy(deep=True) + inner_cfg = StGPTConfig.model_validate({**inner_cfg.model_dump(), "training": {**inner_cfg.training.model_dump(), "batch_size": int(batch_size)}}) + case = build_training_case(inner_cfg) + dataset = RegionDataset(case, inner_cfg, for_inference=True) + checkpoint_genes = tuple(str(item) for item in payload.get("vocab", {}).get("genes", [])) + if checkpoint_genes and checkpoint_genes != dataset.vocab.genes: + raise ValueError("Embedding data gene vocabulary does not match the checkpoint vocabulary.") + + target = _resolve_device(device) + checkpoint_cfg = StGPTConfig.model_validate(payload.get("config", inner_cfg.model_dump())) + n_structures = int(payload.get("n_structures", dataset.n_structures)) + label_vocab = tuple(payload.get("structure_names") or dataset.structure_names) + if len(label_vocab) < n_structures: + # Defensive: align vocab length to head width by padding with placeholders. + label_vocab = tuple(list(label_vocab) + [f"__unknown_{i}__" for i in range(n_structures - len(label_vocab))]) + model = ImageGeneSTGPT( + n_genes=dataset.vocab.size - 1, + n_structures=n_structures, + d_model=inner_cfg.model.d_model, + n_heads=inner_cfg.model.n_heads, + n_layers=inner_cfg.model.n_layers, + dim_feedforward=inner_cfg.model.dim_feedforward, + n_expression_bins=inner_cfg.model.n_expression_bins, + image_channels=inner_cfg.model.image_channels, + patch_scales=inner_cfg.model.patch_scales, + image_encoder_backend="precomputed" if inner_cfg.data.image_embedding_store else inner_cfg.model.image_encoder_backend, + image_encoder_preset=inner_cfg.model.image_encoder_preset, + image_encoder_name=inner_cfg.model.image_encoder_name, + image_encoder_frozen=inner_cfg.model.image_encoder_frozen, + image_embedding_dim=inner_cfg.model.image_embedding_dim or dataset.image_embedding_dim or None, + n_prototypes=checkpoint_cfg.model.n_prototypes, + prototype_temperature=checkpoint_cfg.model.prototype_temperature, + use_expression_values=inner_cfg.model.use_expression_values, + use_image_context=inner_cfg.model.use_image_context, + use_spatial_context=inner_cfg.model.use_spatial_context, + use_structure_context=inner_cfg.model.use_structure_context and inner_cfg.data.include_structure_context, + use_cell_context=inner_cfg.model.use_cell_context, + dropout=inner_cfg.model.dropout, + ) + model.load_state_dict(payload["model_state"], strict=False) + model.to(target) + model.eval() + + loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, collate_fn=dataset.collate, num_workers=0) + embeddings_chunks: list[np.ndarray] = [] + logits_chunks: list[np.ndarray] = [] + have_logits = True + with torch.no_grad(): + for batch in loader: + batch_on_device = {k: v.to(target) if isinstance(v, torch.Tensor) else v for k, v in batch.items()} + output = model( + gene_ids=batch_on_device["gene_ids"], + expr_values=batch_on_device["expr_values"], + expr_bins=batch_on_device["expr_bins"], + image=batch_on_device["image"], + spatial=batch_on_device["spatial"], + context_ids=batch_on_device["context_ids"], + gene_padding_mask=batch_on_device["gene_padding_mask"], + cell_expr_values=batch_on_device["cell_expr_values"], + cell_token_mask=batch_on_device["cell_token_mask"], + object_image=batch_on_device.get("object_image"), + context_image=batch_on_device.get("context_image"), + contour_mask=batch_on_device.get("contour_mask"), + contour_geometry=batch_on_device.get("contour_geometry"), + precomputed_image_embedding=batch_on_device.get("precomputed_image_embedding"), + ) + embeddings_chunks.append(output.region_emb.cpu().numpy().astype(np.float32)) + if output.structure_logits is None: + have_logits = False + else: + logits_chunks.append(output.structure_logits.cpu().numpy().astype(np.float32)) + embeddings = np.vstack(embeddings_chunks) if embeddings_chunks else np.zeros((0, inner_cfg.model.d_model), dtype=np.float32) + logits = np.vstack(logits_chunks) if (have_logits and logits_chunks) else None + return _InferenceOutputs( + region_table=dataset.region_table.copy(), + embeddings=embeddings, + structure_logits=logits, + label_vocab=label_vocab, + ) + + +# --------------------------------------------------------------------------- +# Seed CSV parsing and validation +# --------------------------------------------------------------------------- + + +def _read_seed_labels(path: Path, all_region_ids: np.ndarray) -> pd.DataFrame: + frame = pd.read_csv(path) if path.suffix.lower() in {".csv", ".tsv"} else pd.read_parquet(path) + required = {"region_id", "structure_label"} + missing = required - set(frame.columns) + if missing: + raise ValueError(f"Seed labels file is missing required columns: {sorted(missing)}") + frame = frame[list(required) + (["confidence"] if "confidence" in frame.columns else [])].copy() + frame["region_id"] = frame["region_id"].astype(str) + frame["structure_label"] = frame["structure_label"].astype(str) + if "confidence" not in frame.columns: + frame["confidence"] = 1.0 + else: + frame["confidence"] = pd.to_numeric(frame["confidence"], errors="coerce").fillna(1.0).clip(lower=0.0) + frame = frame.drop_duplicates(subset=["region_id"], keep="first") + region_set = set(all_region_ids.tolist()) + unknown = sorted(set(frame["region_id"].tolist()) - region_set) + if unknown: + raise ValueError(f"Seed labels reference unknown region_ids: {unknown[:5]}{'...' if len(unknown) > 5 else ''}") + return frame.reset_index(drop=True) + + +def _read_region_id_list(path: Path) -> list[str]: + if path.suffix.lower() == ".csv": + frame = pd.read_csv(path) + column = "region_id" if "region_id" in frame.columns else frame.columns[0] + return [str(item) for item in frame[column].tolist()] + return [line.strip() for line in path.read_text(encoding="utf-8").splitlines() if line.strip()] + + +def _resolve_label_vocabulary( + seed_frame: pd.DataFrame, + checkpoint_vocab: tuple[str, ...], + classifier: ClassifierKind, +) -> tuple[tuple[str, ...], str]: + seed_set = sorted(set(seed_frame["structure_label"].tolist())) + if all(label in checkpoint_vocab for label in seed_set): + return checkpoint_vocab, "checkpoint_vocab" + if classifier == "structure_head": + unknown = [label for label in seed_set if label not in checkpoint_vocab] + raise ValueError( + f"Path A (structure_head) was requested but seed labels include " + f"classes not in the checkpoint vocabulary: {unknown}" + ) + return tuple(seed_set), "seed_vocab" + + +def _build_index_map(region_ids: np.ndarray) -> dict[str, int]: + return {str(rid): idx for idx, rid in enumerate(region_ids)} + + +# --------------------------------------------------------------------------- +# Path A: structure_head with temperature scaling +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class _PathPrediction: + pool_probs: np.ndarray # (n_pool, n_classes) + pool_top1_idx: np.ndarray # (n_pool,) + pool_top1_prob: np.ndarray # (n_pool,) + pool_entropy: np.ndarray # (n_pool,) + temperature: float + + +def _fit_predict_structure_head( + *, + logits: np.ndarray, + seed_indices: np.ndarray, + seed_labels: np.ndarray, + seed_weights: np.ndarray, + pool_indices: np.ndarray, + n_classes: int, + rng: np.random.Generator, + seed_folds: int, +) -> _PathPrediction: + head_width = logits.shape[1] + if n_classes != head_width: + # Pad logits with -inf for any class we did not train on; should not happen + # because checkpoint_vocab mode aligns the two. + raise ValueError("Structure head width does not match label vocabulary size.") + seed_logits = logits[seed_indices] + temperature = _fit_temperature(seed_logits, seed_labels, seed_weights) + pool_logits = logits[pool_indices] / temperature + pool_probs = _softmax(pool_logits) + return _summarize(pool_probs, temperature) + + +def _fit_temperature( + logits: np.ndarray, + labels: np.ndarray, + weights: np.ndarray, +) -> float: + if logits.size == 0: + return 1.0 + logits_t = torch.from_numpy(logits.astype(np.float64)) + labels_t = torch.from_numpy(labels.astype(np.int64)) + weights_t = torch.from_numpy(weights.astype(np.float64)) + log_temp = torch.zeros((), dtype=torch.float64, requires_grad=True) + optimizer = torch.optim.LBFGS([log_temp], lr=0.1, max_iter=50, line_search_fn="strong_wolfe") + + def closure() -> torch.Tensor: + optimizer.zero_grad() + temperature = torch.exp(log_temp).clamp(min=1e-3, max=1e3) + scaled = logits_t / temperature + log_probs = torch.log_softmax(scaled, dim=1) + nll = -log_probs.gather(1, labels_t.unsqueeze(1)).squeeze(1) + loss = (nll * weights_t).sum() / weights_t.sum().clamp(min=1e-6) + loss.backward() + return loss + + try: + optimizer.step(closure) + except RuntimeError: + # LBFGS occasionally fails on degenerate inputs; fall back to T=1.0. + return 1.0 + temperature = float(torch.exp(log_temp).clamp(min=1e-3, max=1e3).item()) + if not np.isfinite(temperature) or temperature <= 0: + return 1.0 + return temperature + + +# --------------------------------------------------------------------------- +# Path B: per-class prototypes on the region embedding +# --------------------------------------------------------------------------- + + +def _fit_predict_prototype( + *, + embeddings: np.ndarray, + seed_indices: np.ndarray, + seed_labels: np.ndarray, + seed_weights: np.ndarray, + pool_indices: np.ndarray, + n_classes: int, + rng: np.random.Generator, + seed_folds: int, +) -> _PathPrediction: + normalized = _l2_normalize(embeddings) + prototypes = _class_prototypes(normalized, seed_indices, seed_labels, seed_weights, n_classes) + seed_logits = normalized[seed_indices] @ prototypes.T # cosine similarities + temperature = _fit_temperature(seed_logits, seed_labels, seed_weights) + pool_logits = (normalized[pool_indices] @ prototypes.T) / temperature + pool_probs = _softmax(pool_logits) + return _summarize(pool_probs, temperature) + + +def _class_prototypes( + normalized: np.ndarray, + seed_indices: np.ndarray, + seed_labels: np.ndarray, + seed_weights: np.ndarray, + n_classes: int, +) -> np.ndarray: + d_model = normalized.shape[1] + prototypes = np.zeros((n_classes, d_model), dtype=np.float32) + for cls in range(n_classes): + mask = seed_labels == cls + if not mask.any(): + prototypes[cls] = 0.0 + continue + weights = seed_weights[mask].astype(np.float64) + weights = weights / max(weights.sum(), 1e-6) + proto = (normalized[seed_indices[mask]].astype(np.float64) * weights[:, None]).sum(axis=0) + norm = float(np.linalg.norm(proto)) + prototypes[cls] = (proto / norm).astype(np.float32) if norm > 1e-6 else proto.astype(np.float32) + return prototypes + + +# --------------------------------------------------------------------------- +# Shared helpers +# --------------------------------------------------------------------------- + + +def _softmax(x: np.ndarray) -> np.ndarray: + x = np.asarray(x, dtype=np.float64) + shifted = x - x.max(axis=1, keepdims=True) + exp = np.exp(shifted) + return (exp / exp.sum(axis=1, keepdims=True)).astype(np.float64) + + +def _summarize(pool_probs: np.ndarray, temperature: float) -> _PathPrediction: + if pool_probs.size == 0: + return _PathPrediction( + pool_probs=pool_probs, + pool_top1_idx=np.zeros(0, dtype=np.int64), + pool_top1_prob=np.zeros(0, dtype=np.float64), + pool_entropy=np.zeros(0, dtype=np.float64), + temperature=temperature, + ) + top1 = pool_probs.argmax(axis=1) + top1_prob = pool_probs[np.arange(pool_probs.shape[0]), top1] + entropy = -np.sum(pool_probs * np.log(np.clip(pool_probs, 1e-12, 1.0)), axis=1) + return _PathPrediction( + pool_probs=pool_probs, + pool_top1_idx=top1.astype(np.int64), + pool_top1_prob=top1_prob.astype(np.float64), + pool_entropy=entropy.astype(np.float64), + temperature=temperature, + ) + + +def _l2_normalize(matrix: np.ndarray) -> np.ndarray: + norms = np.linalg.norm(matrix, axis=1, keepdims=True) + norms = np.where(norms < 1e-12, 1.0, norms) + return (matrix / norms).astype(np.float32) + + +def _nearest_seed( + embeddings: np.ndarray, + seed_indices: np.ndarray, + pool_indices: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + if pool_indices.size == 0 or seed_indices.size == 0: + return np.zeros(pool_indices.size, dtype=np.int64), np.zeros(pool_indices.size, dtype=np.float64) + normalized = _l2_normalize(embeddings) + seed_emb = normalized[seed_indices] + pool_emb = normalized[pool_indices] + sim = pool_emb @ seed_emb.T + nearest = sim.argmax(axis=1) + cosine_dist = 1.0 - sim[np.arange(sim.shape[0]), nearest] + return seed_indices[nearest].astype(np.int64), np.clip(cosine_dist.astype(np.float64), 0.0, 2.0) + + +def _seed_to_seed_distances(embeddings: np.ndarray, seed_indices: np.ndarray) -> np.ndarray: + if seed_indices.size < 2: + return np.zeros(0, dtype=np.float64) + normalized = _l2_normalize(embeddings) + sub = normalized[seed_indices] + sim = sub @ sub.T + np.fill_diagonal(sim, -np.inf) + nearest_sim = sim.max(axis=1) + return np.clip(1.0 - nearest_sim, 0.0, 2.0).astype(np.float64) + + +def _build_predictions_table( + *, + region_ids_all: np.ndarray, + region_df: pd.DataFrame, + seed_indices: np.ndarray, + seed_labels_int: np.ndarray, + pool_indices: np.ndarray, + primary: _PathPrediction, + primary_path: str, + nearest_seed_idx: np.ndarray, + nearest_seed_dist: np.ndarray, + label_vocab: tuple[str, ...], + abstain_prob: float, + entropy_cutoff: float, + distance_cutoff: float, + evidence_prefix: str, +) -> pd.DataFrame: + rows: list[dict[str, Any]] = [] + qc_col = region_df["qc_flag"].astype(str).to_numpy() if "qc_flag" in region_df.columns else np.array(["unknown"] * len(region_df)) + for rank, idx in enumerate(seed_indices.tolist()): + rid = str(region_ids_all[idx]) + rows.append({ + "region_id": rid, + "predicted_label": label_vocab[int(seed_labels_int[rank])], + "predicted_prob": 1.0, + "entropy": 0.0, + "nearest_seed_region_id": rid, + "nearest_seed_distance": 0.0, + "qc_flag": "seed", + "classifier": primary_path, + "evidence_id": _evidence_id(evidence_prefix, rid), + "propagation_kind": "same_slide", + "expression_present": True, + "source_case_id": evidence_prefix.split("|", 1)[0], + }) + for pool_rank, idx in enumerate(pool_indices.tolist()): + rid = str(region_ids_all[idx]) + top_idx = int(primary.pool_top1_idx[pool_rank]) + prob = float(primary.pool_top1_prob[pool_rank]) + entropy = float(primary.pool_entropy[pool_rank]) + nearest_idx = int(nearest_seed_idx[pool_rank]) + nearest_distance = float(nearest_seed_dist[pool_rank]) + abstain = ( + prob < float(abstain_prob) + or entropy > float(entropy_cutoff) + or nearest_distance > float(distance_cutoff) + ) + rows.append({ + "region_id": rid, + "predicted_label": ABSTAIN_LABEL if abstain else label_vocab[top_idx], + "predicted_prob": prob, + "entropy": entropy, + "nearest_seed_region_id": str(region_ids_all[nearest_idx]) if nearest_seed_idx.size else "", + "nearest_seed_distance": nearest_distance, + "qc_flag": str(qc_col[idx]), + "classifier": primary_path, + "evidence_id": _evidence_id(evidence_prefix, rid), + "propagation_kind": "same_slide", + "expression_present": True, + "source_case_id": evidence_prefix.split("|", 1)[0], + }) + return pd.DataFrame(rows) + + +def _build_probabilities_table( + *, + region_ids_all: np.ndarray, + pool_indices: np.ndarray, + primary: _PathPrediction, + label_vocab: tuple[str, ...], +) -> pd.DataFrame: + if pool_indices.size == 0 or primary.pool_probs.size == 0: + return pd.DataFrame(columns=["region_id", "structure_label", "probability"]) + rows: list[dict[str, Any]] = [] + for pool_rank, idx in enumerate(pool_indices.tolist()): + rid = str(region_ids_all[idx]) + for cls_idx, label in enumerate(label_vocab): + rows.append({ + "region_id": rid, + "structure_label": label, + "probability": float(primary.pool_probs[pool_rank, cls_idx]), + }) + return pd.DataFrame(rows) + + +def _build_path_agreement_table( + *, + region_ids_all: np.ndarray, + pool_indices: np.ndarray, + paths: dict[str, _PathPrediction], + label_vocab: tuple[str, ...], +) -> pd.DataFrame: + head = paths["structure_head"] + proto = paths["prototype_knn"] + rows: list[dict[str, Any]] = [] + for pool_rank, idx in enumerate(pool_indices.tolist()): + head_top = label_vocab[int(head.pool_top1_idx[pool_rank])] + proto_top = label_vocab[int(proto.pool_top1_idx[pool_rank])] + rows.append({ + "region_id": str(region_ids_all[idx]), + "structure_head_top1": head_top, + "structure_head_prob": float(head.pool_top1_prob[pool_rank]), + "prototype_knn_top1": proto_top, + "prototype_knn_prob": float(proto.pool_top1_prob[pool_rank]), + "agree": bool(head_top == proto_top), + }) + return pd.DataFrame(rows) + + +# --------------------------------------------------------------------------- +# Seed cross-validation QC +# --------------------------------------------------------------------------- + + +def _seed_cross_validation( + *, + path: str, + embeddings: np.ndarray, + logits: np.ndarray | None, + seed_indices: np.ndarray, + seed_labels: np.ndarray, + seed_weights: np.ndarray, + n_classes: int, + seed_folds: int, + rng: np.random.Generator, +) -> dict[str, Any]: + n_seeds = seed_indices.size + per_class_counts = [int((seed_labels == c).sum()) for c in range(n_classes)] + min_per_class = min(per_class_counts) if per_class_counts else 0 + effective_folds = max(2, min(int(seed_folds), n_seeds, min_per_class)) if min_per_class >= 2 else 0 + if effective_folds < 2 or n_seeds < 2: + return { + "status": "skipped", + "reason": "insufficient_seeds_for_cv", + "n_seeds": int(n_seeds), + "n_classes": int(n_classes), + "per_class_counts": per_class_counts, + } + folds = _stratified_fold_indices(seed_labels, n_folds=effective_folds, rng=rng) + oof_pred = np.full(n_seeds, -1, dtype=np.int64) + for fold_idx in range(effective_folds): + train_mask = folds != fold_idx + valid_mask = ~train_mask + if path == "structure_head" and logits is not None: + train_logits = logits[seed_indices[train_mask]] + valid_logits = logits[seed_indices[valid_mask]] + temperature = _fit_temperature(train_logits, seed_labels[train_mask], seed_weights[train_mask]) + probs = _softmax(valid_logits / temperature) + else: + normalized = _l2_normalize(embeddings) + prototypes = _class_prototypes( + normalized, + seed_indices[train_mask], + seed_labels[train_mask], + seed_weights[train_mask], + n_classes, + ) + valid_sim = normalized[seed_indices[valid_mask]] @ prototypes.T + temperature = _fit_temperature( + normalized[seed_indices[train_mask]] @ prototypes.T, + seed_labels[train_mask], + seed_weights[train_mask], + ) + probs = _softmax(valid_sim / temperature) + oof_pred[valid_mask] = probs.argmax(axis=1) + return _classification_report(seed_labels, oof_pred, n_classes=n_classes, effective_folds=int(effective_folds)) + + +def _stratified_fold_indices(labels: np.ndarray, *, n_folds: int, rng: np.random.Generator) -> np.ndarray: + folds = np.full(labels.size, -1, dtype=np.int64) + for cls in np.unique(labels): + idx = np.where(labels == cls)[0] + rng.shuffle(idx) + for j, pos in enumerate(idx): + folds[pos] = j % n_folds + return folds + + +def _classification_report(true: np.ndarray, pred: np.ndarray, *, n_classes: int, effective_folds: int) -> dict[str, Any]: + confusion = np.zeros((n_classes, n_classes), dtype=np.int64) + for t, p in zip(true.tolist(), pred.tolist()): + if p < 0: + continue + confusion[int(t), int(p)] += 1 + per_class: list[dict[str, float]] = [] + f1_scores: list[float] = [] + for cls in range(n_classes): + tp = int(confusion[cls, cls]) + fp = int(confusion[:, cls].sum() - tp) + fn = int(confusion[cls, :].sum() - tp) + precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0 + recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0 + f1 = (2 * precision * recall / (precision + recall)) if (precision + recall) > 0 else 0.0 + per_class.append({ + "class_index": cls, + "precision": float(precision), + "recall": float(recall), + "f1": float(f1), + "support": int(confusion[cls, :].sum()), + }) + f1_scores.append(float(f1)) + return { + "status": "ok", + "n_folds": int(effective_folds), + "macro_f1": float(np.mean(f1_scores)) if f1_scores else 0.0, + "accuracy": float((true == pred).mean()) if pred.size else 0.0, + "confusion_matrix": confusion.tolist(), + "per_class": per_class, + } + + +def _label_counts(seed_labels: np.ndarray, label_vocab: tuple[str, ...]) -> dict[str, int]: + counts: dict[str, int] = {label: 0 for label in label_vocab} + for cls_idx in seed_labels.tolist(): + counts[label_vocab[int(cls_idx)]] += 1 + return counts + + +def _collect_warnings( + seed_labels: np.ndarray, + label_vocab: tuple[str, ...], + n_abstain: int, + n_pool: int, + seed_cv: dict[str, Any], +) -> list[str]: + warnings: list[str] = [] + counts = _label_counts(seed_labels, label_vocab) + low = [label for label, count in counts.items() if count < 5] + if low: + warnings.append(f"low_seed_count: classes with fewer than 5 seeds: {low}") + if n_pool > 0 and (n_abstain / n_pool) > 0.5: + warnings.append(f"high_abstain_rate: {n_abstain}/{n_pool} regions abstained") + if seed_cv.get("status") == "ok" and seed_cv.get("macro_f1", 0.0) < 0.6: + warnings.append(f"low_seed_macro_f1: {seed_cv['macro_f1']:.3f}") + return warnings + + +# --------------------------------------------------------------------------- +# Provenance helpers +# --------------------------------------------------------------------------- + + +def _evidence_prefix(cfg: StGPTConfig, checkpoint_path: Path, seeds_path: Path) -> str: + return f"{cfg.case_name}|{_config_hash(cfg)}|{_sha256_path(checkpoint_path) or 'no-checkpoint'}|{_sha256_path(seeds_path) or 'no-seeds'}" + + +def _evidence_id(prefix: str, region_id: str) -> str: + return "ann_" + hashlib.sha256(f"{prefix}|{region_id}".encode("utf-8")).hexdigest()[:16] + + +def _config_hash(cfg: StGPTConfig) -> str: + payload = json.dumps(cfg.model_dump(mode="json"), sort_keys=True).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def _sha256_path(path: Path) -> str | None: + if not path.exists() or not path.is_file(): + return None + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _resolve_device(name: str) -> torch.device: + normalized = str(name).lower() + if normalized == "auto": + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + if normalized == "cuda" and not torch.cuda.is_available(): + return torch.device("cpu") + return torch.device(normalized) + + +__all__ = [ + "ABSTAIN_LABEL", + "AutoAnnotationResult", + "SCHEMA_VERSION", + "annotate_regions", +] diff --git a/src/stgpt/cli.py b/src/stgpt/cli.py index c63938c..dd4d135 100644 --- a/src/stgpt/cli.py +++ b/src/stgpt/cli.py @@ -11,6 +11,7 @@ import typer from . import __version__ +from .annotation import annotate_regions as annotate_regions_backend from .config import StGPTConfig from .contour_store import pack_contour_patches from .data import build_training_manifest @@ -339,6 +340,40 @@ def embed_regions_command( typer.echo(json.dumps(result.to_dict(), indent=2)) +@app.command("annotate-regions") +def annotate_regions_command( + config: Annotated[Path, typer.Option("--config", "-c", exists=True)], + checkpoint: Annotated[Path, typer.Option("--checkpoint", "-k", exists=True)], + seed_labels: Annotated[Path, typer.Option("--seed-labels", "-s", exists=True)], + output: Annotated[Path, typer.Option("--output", "-o")], + region_ids: Annotated[Path | None, typer.Option("--region-ids", exists=True)] = None, + include_no_image: Annotated[bool, typer.Option("--include-no-image/--require-image")] = False, + classifier: Annotated[str, typer.Option("--classifier")] = "both", + abstain_prob: Annotated[float, typer.Option("--abstain-prob", min=0.0, max=1.0)] = 0.5, + write_probabilities: Annotated[bool, typer.Option("--write-probabilities/--no-probabilities")] = False, + seed_folds: Annotated[int, typer.Option("--seed-folds", min=2)] = 5, + rng_seed: Annotated[int, typer.Option("--rng-seed")] = 42, + batch_size: Annotated[int, typer.Option("--batch-size", min=1)] = 32, + device: Annotated[str, typer.Option("--device")] = "auto", +) -> None: + result = annotate_regions_backend( + config=config, + checkpoint=checkpoint, + seed_labels=seed_labels, + output_dir=output, + region_ids=region_ids, + include_no_image=include_no_image, + classifier=classifier, # type: ignore[arg-type] + abstain_prob=abstain_prob, + write_probabilities=write_probabilities, + seed_folds=seed_folds, + rng_seed=rng_seed, + batch_size=batch_size, + device=device, + ) + typer.echo(json.dumps(result, indent=2)) + + @app.command("export-spatho") def export_spatho( config: Annotated[Path, typer.Option("--config", "-c", exists=True)], diff --git a/src/stgpt/runtime/__init__.py b/src/stgpt/runtime/__init__.py index 67ba50b..f95eeca 100644 --- a/src/stgpt/runtime/__init__.py +++ b/src/stgpt/runtime/__init__.py @@ -1,8 +1,9 @@ from __future__ import annotations from pathlib import Path -from typing import Any +from typing import Any, Literal +from ..annotation import annotate_regions as _annotate_regions from ..config import StGPTConfig from ..evaluation import evaluate from ..foundation import package_model @@ -85,7 +86,42 @@ def export_spatho_artifacts( ) +def annotate_regions( + *, + config: StGPTConfig | str | Path, + checkpoint: str | Path, + seed_labels: str | Path, + output_dir: str | Path, + region_ids: str | Path | None = None, + include_no_image: bool = False, + classifier: Literal["structure_head", "prototype_knn", "both"] = "both", + abstain_prob: float = 0.5, + write_probabilities: bool = False, + seed_folds: int = 5, + rng_seed: int = 42, + batch_size: int = 32, + device: str = "auto", +) -> dict[str, Any]: + """Propagate sparse expert structure labels to unannotated regions.""" + return _annotate_regions( + config=config, + checkpoint=checkpoint, + seed_labels=seed_labels, + output_dir=output_dir, + region_ids=region_ids, + include_no_image=include_no_image, + classifier=classifier, + abstain_prob=abstain_prob, + write_probabilities=write_probabilities, + seed_folds=seed_folds, + rng_seed=rng_seed, + batch_size=batch_size, + device=device, + ) + + __all__ = [ + "annotate_regions", "embed_cells", "embed_regions", "evaluate_checkpoint", diff --git a/tests/test_annotation.py b/tests/test_annotation.py new file mode 100644 index 0000000..2a1cbb0 --- /dev/null +++ b/tests/test_annotation.py @@ -0,0 +1,218 @@ +"""Smoke tests for the region auto-annotation runtime.""" +from __future__ import annotations + +import json +from pathlib import Path + +import pandas as pd +import pytest +from typer.testing import CliRunner + +from stgpt.annotation import ABSTAIN_LABEL, annotate_regions +from stgpt.cli import app +from stgpt.config import DataConfig, ModelConfig, SplitConfig, StGPTConfig, TrainingConfig +from stgpt.runtime import annotate_regions as runtime_annotate_regions +from stgpt.spatho import run_spatho_export +from stgpt.training import train + + +def _small_config(tmp_path: Path) -> StGPTConfig: + return StGPTConfig( + case_name="annotation_smoke", + data=DataConfig( + mode="synthetic", + output_dir=str(tmp_path / "case"), + n_cells=24, + n_genes=16, + n_structures=3, + image_size=32, + seed=11, + ), + model=ModelConfig( + d_model=32, + n_heads=4, + n_layers=1, + max_genes=12, + image_size=32, + n_expression_bins=8, + n_prototypes=3, + ), + training=TrainingConfig(batch_size=4, max_steps=2, output_dir=str(tmp_path / "train"), device="cpu", seed=7), + split=SplitConfig(seed=3), + ) + + +def _train_and_seed(tmp_path: Path) -> tuple[StGPTConfig, Path, Path]: + cfg = _small_config(tmp_path) + checkpoint = Path(train(cfg, preset="smoke", max_steps=2)["checkpoint"]) + spatho_out = tmp_path / "spatho" + export = run_spatho_export(cfg, checkpoint=checkpoint, output_dir=spatho_out, batch_size=4, device="cpu") + regions = pd.read_parquet(export.region_embeddings) + # Pick 60% of regions per class as seeds so each class has >= 2 seeds for k-fold. + seed_rows_list = [] + for _, group in regions.groupby("structure_label", sort=True): + keep = max(2, int(len(group) * 0.6)) + seed_rows_list.append(group.head(keep)[["region_id", "structure_label"]]) + seeds = pd.concat(seed_rows_list, ignore_index=True) + seeds["confidence"] = 1.0 + seeds_path = tmp_path / "seeds.csv" + seeds.to_csv(seeds_path, index=False) + return cfg, checkpoint, seeds_path + + +def test_annotate_regions_writes_predictions(tmp_path: Path) -> None: + cfg, checkpoint, seeds_path = _train_and_seed(tmp_path) + out_dir = tmp_path / "annotate" + result = annotate_regions( + config=cfg, + checkpoint=checkpoint, + seed_labels=seeds_path, + output_dir=out_dir, + write_probabilities=True, + batch_size=4, + device="cpu", + ) + predictions_path = Path(result["predictions"]) + report_path = Path(result["report"]) + probs_path = Path(result["probabilities"]) + agreement_path = Path(result["path_agreement"]) + assert predictions_path.exists() + assert report_path.exists() + assert probs_path.exists() + assert agreement_path.exists() + + predictions = pd.read_parquet(predictions_path) + expected_cols = { + "region_id", + "predicted_label", + "predicted_prob", + "entropy", + "nearest_seed_region_id", + "nearest_seed_distance", + "qc_flag", + "classifier", + "evidence_id", + "propagation_kind", + "expression_present", + } + assert expected_cols.issubset(set(predictions.columns)) + + seeds = pd.read_csv(seeds_path) + seed_rows = predictions[predictions["region_id"].isin(seeds["region_id"])] + assert (seed_rows["qc_flag"] == "seed").all() + assert set(seed_rows["predicted_label"].tolist()) == set(seeds["structure_label"].tolist()) + assert (seed_rows["predicted_prob"] == 1.0).all() + + pool_rows = predictions[predictions["qc_flag"] != "seed"] + assert len(pool_rows) > 0 + valid = pool_rows["predicted_label"] + allowed_labels = set(seeds["structure_label"].tolist()) | {ABSTAIN_LABEL} + assert set(valid.tolist()).issubset(allowed_labels) + assert (pool_rows["predicted_prob"] >= 0.0).all() + assert (pool_rows["predicted_prob"] <= 1.0 + 1e-6).all() + assert (pool_rows["propagation_kind"] == "same_slide").all() + + report = json.loads(report_path.read_text(encoding="utf-8")) + assert report["schema_version"].startswith("stgpt.region_auto_annotation") + assert report["n_seed_regions"] == len(seeds) + assert report["n_pool_regions"] == len(pool_rows) + assert report["primary_classifier"] in {"structure_head", "prototype_knn"} + assert "label_vocab" in report + assert "seed_cross_validation" in report + assert isinstance(report["warnings"], list) + + +def test_annotate_regions_unknown_region_id_raises(tmp_path: Path) -> None: + cfg, checkpoint, seeds_path = _train_and_seed(tmp_path) + bad_seeds = tmp_path / "bad_seeds.csv" + pd.DataFrame({ + "region_id": ["definitely_not_a_real_region"], + "structure_label": ["structure_0"], + "confidence": [1.0], + }).to_csv(bad_seeds, index=False) + with pytest.raises(ValueError, match="unknown region_ids"): + annotate_regions( + config=cfg, + checkpoint=checkpoint, + seed_labels=bad_seeds, + output_dir=tmp_path / "bad", + batch_size=4, + device="cpu", + ) + + +def test_annotate_regions_runtime_alias_matches(tmp_path: Path) -> None: + cfg, checkpoint, seeds_path = _train_and_seed(tmp_path) + direct = annotate_regions( + config=cfg, + checkpoint=checkpoint, + seed_labels=seeds_path, + output_dir=tmp_path / "direct", + batch_size=4, + device="cpu", + ) + runtime_result = runtime_annotate_regions( + config=cfg, + checkpoint=checkpoint, + seed_labels=seeds_path, + output_dir=tmp_path / "runtime", + batch_size=4, + device="cpu", + ) + assert direct.keys() == runtime_result.keys() + direct_pred = pd.read_parquet(direct["predictions"]).sort_values("region_id").reset_index(drop=True) + runtime_pred = pd.read_parquet(runtime_result["predictions"]).sort_values("region_id").reset_index(drop=True) + pd.testing.assert_frame_equal(direct_pred, runtime_pred) + + +def test_annotate_regions_cli(tmp_path: Path) -> None: + cfg, checkpoint, seeds_path = _train_and_seed(tmp_path) + config_path = tmp_path / "config.yaml" + config_path.write_text( + f""" +case_name: {cfg.case_name} +data: + mode: synthetic + output_dir: {cfg.data.output_dir} + n_cells: {cfg.data.n_cells} + n_genes: {cfg.data.n_genes} + n_structures: {cfg.data.n_structures} + image_size: {cfg.data.image_size} + seed: {cfg.data.seed} +model: + d_model: {cfg.model.d_model} + n_heads: {cfg.model.n_heads} + n_layers: {cfg.model.n_layers} + max_genes: {cfg.model.max_genes} + image_size: {cfg.model.image_size} + n_expression_bins: {cfg.model.n_expression_bins} + n_prototypes: {cfg.model.n_prototypes} +training: + batch_size: {cfg.training.batch_size} + max_steps: {cfg.training.max_steps} + output_dir: {cfg.training.output_dir} + device: cpu + seed: {cfg.training.seed} +split: + seed: {cfg.split.seed} +""", + encoding="utf-8", + ) + runner = CliRunner() + result = runner.invoke( + app, + [ + "annotate-regions", + "--config", str(config_path), + "--checkpoint", str(checkpoint), + "--seed-labels", str(seeds_path), + "--output", str(tmp_path / "cli"), + "--classifier", "both", + "--batch-size", "4", + "--device", "cpu", + ], + ) + assert result.exit_code == 0, result.output + payload = json.loads(result.output) + assert Path(payload["predictions"]).exists() + assert Path(payload["report"]).exists()