Skip to content

SLEAP-NN v0.3.1

Choose a tag to compare

@gitttt-1234 gitttt-1234 released this 28 Jul 06:05
· 9 commits to main since this release
880f40c

sleap-nn v0.3.1 Release Notes

Summary

SLEAP-NN v0.3.1 is a focused follow-up to v0.3.0, adding sliding-window tiling for high-res / small-object frames, a new whole-frame semantic_segmentation model type, pretrained HuggingFace backbones, and several centroid-training correctness fixes (single-source centroid targets, training directly from UserCentroid annotations, including pure-centroid frames). It also restores CPU/GPU overlap in sleap-nn predict (dropped silently in an earlier inference-pipeline refactor) and fixes a repo-wide sweep of predict-pipeline and training/evaluation correctness bugs (postprocess-override isolation, checkpoint-monitor metrics, OKS eval config forwarding, a ground-truth mutation bug, and more), and bumps the sleap-io dependency to >=0.9.2,<0.10.0 (re-ID/Category/Event annotations, large-project save hardening, O(N) merges). 23 PRs since v0.3.0.

⚠️ Breaking changes: whole-frame segmentation models (semantic_segmentation, bottomup_segmentation) trained before #693 have a baked-in mask/image misalignment and should be retrained; default mask-evaluation metrics change (predicted-instance-linked GT masks are now excluded); frame-caching failures during training now hard-fail instead of warning-and-continuing; and centroid training's default target-source resolution changed. Read Breaking Changes before upgrading.

Highlights:

  • Sliding-window tiling (data_config.preprocessing.tiling) for 4K++ frames with small objects — cut into overlapping tiles, run per-tile at native resolution, stitch via Gaussian-weighted merge. Opt-in; supported for single_instance and bottomup_segmentation. (#687)
  • New semantic_segmentation model type — whole-frame binary fg/bg mask, no instance grouping, with a matching-free --match_method semantic eval mode. (#688)
  • Pretrained HuggingFace backbones (sleap-nn[backbones]) — use any AutoBackbone (ConvNeXtV2, ResNet, Swinv2, DINOv2/v3, ...) as a model's encoder, frozen or fine-tuned. (#681)
  • Centroid training correctness: centroid_source config ("user"/"computed"/None) fixes a mixed-annotation footgun where the centroid head could train against two different centroid definitions in one run (#704); CentroidDataset can now train directly on UserCentroid annotations, including pure-centroid frames with no pose instance (#702, #703).
  • clDice mask metric for connectivity-aware segmentation quality, plus a configurable best-checkpoint monitor metric so segmentation runs can checkpoint on full-resolution quality instead of coarse val/loss. (#682, #691)
  • ⚠️ sleap-io >=0.9.2,<0.10.0 — audited past the previous <0.9.0 ceiling (re-ID, Category, Event annotations; large-project save hardening; O(N) merges). See sleap-io's v0.9.0/v0.9.1/v0.9.2 release notes.

Installation

# Install / upgrade the CLI tool (auto-selects the right torch backend)
uv tool install sleap-nn --torch-backend auto --upgrade

# Verify
sleap-nn --version
# Expected output: sleap-nn 0.3.1

Breaking Changes

⚠️ Whole-frame segmentation masks were misaligned with the image grid (#693)

semantic_segmentation and bottomup_segmentation models trained before this fix had their GT masks resized straight to the padded frame size instead of threaded through the same size-match/scale/stride-pad chain as the image — offsetting the foreground target by roughly half the padding (worse toward the bottom-right) whenever height/width weren't multiples of max_stride. This is a training-target bug, not an inference-geometry change; affected models should be retrained.

⚠️ Mask evaluation now excludes predicted-instance-linked GT masks by default (#694)

run_evaluation(match_method="mask") previously counted every mask in frame.masks as ground truth, even ones linked to a PredictedInstance, which imposed an artificial, unreachable recall/F1 ceiling on labels files containing stray predicted instances. exclude_predicted_instance_masks (driven by the existing user_labels_only argument, default True) now drops those from the GT side only — expect higher default recall/F1 on affected files. Predicted-side masks and match_method="semantic" are unaffected; pass user_labels_only=False to restore the old counting.

⚠️ Frame-caching failures now raise instead of warn (#701)

Training runs that previously continued past frame-caching errors (logged as warnings, then crashed later mid-epoch with a confusing FileNotFoundError/KeyError) now hard-fail immediately with a RuntimeError naming the bad frame(s) and a re-encode hint. A corrupt-video or disk-full condition will now stop training up front rather than silently producing an incomplete cache.

⚠️ Centroid target source is now resolved once per dataset, not per frame (#704)

model_config.head_configs.centroid.confmaps.centroid_source ("user", "computed", or None) replaces #702's per-frame fallback: the source is now inferred once for the whole dataset and applied consistently across train/val, and CentroidDataset drops frames that can't supply the chosen target (e.g. pose-only frames in "user" mode). This changes the effective training set for datasets that mix annotation styles. Leaving it unset infers a source with a loud warning recommending it be set explicitly.


New Features

Tiling, segmentation & pretrained backbones

  • Sliding-window tiling (#687) for high-res / small-object frames (4K microscopy, plant-root scans, ...). Tile geometry auto-sizes from labels, is written into the model config, and is parity-checked at inference; unsupported configs (pretrained-encoder backbones, ClassVectorsHead) get a clear error instead of a silent no-op. sleap-nn export now warns when exporting a tiled model to ONNX/TensorRT, since tiled export isn't supported yet and the exported model would otherwise silently run whole-frame instead of tiled.
  • semantic_segmentation model type (#688) — a lone SegmentationHead on the whole frame predicting one binary fg/bg mask, no instance grouping; matching-free --match_method semantic eval (whole-frame IoU/clDice/boundary-IoU).
  • Pretrained HuggingFace backbones (#681) via the optional sleap-nn[backbones] extra — any AutoBackbone as encoder for pose, centroid, or segmentation models, frozen or fine-tuned (model_config.backbone_config.pretrained).
  • clDice mask metric (#682) — centerline-Dice score (mean_cldice / eval/val/mask_mean_cldice) for connectivity-aware quality on thin/tubular structures where mask IoU is misleading.

Centroid training

  • Train from UserCentroid annotations (#702, #703) — CentroidDataset trains directly on first-class sio.UserCentroid annotations when present (enabling active-learning workflows with a user-seeded, non-node-tied centroid), falling back to keypoint-derived centroids otherwise. #703 is the follow-up that makes the pure-centroid seeding case (a frame with a UserCentroid but no pose instance at all — the Phase-1 active-learning workflow) actually reach the dataset: the train/val split previously filtered to has_user_instances before CentroidDataset ever saw the frame, silently producing 0 training frames for a labels file that was otherwise entirely valid.
  • Single-source centroid targets (#704) — see Breaking Changes.

Training diagnostics & robustness

  • Configurable best-checkpoint metric (#691) — ModelCkptConfig.monitor/mode (e.g. eval/val/fg_mean_cldice, mode: max) fixes segmentation best.ckpt selection, which previously used the coarse val/loss instead of full-resolution quality metrics.
  • Confmap fg/bg MSE diagnostic (#698) — {train,val}/confmap_loss_fg, confmap_loss_bg, confmap_fg_frac logged (no effect on training/checkpoints) to surface foreground-vs-background fit that a single blended loss number obscures.
  • Trainer-accelerator validation (#708) — a saved config's trainer_accelerator (e.g. mps, cuda) is now checked against actual device availability at training setup and falls back to "auto" with a log message instead of crashing deep in Trainer.train(); fixes reloading a config across machines (e.g. Mac → Linux/CUDA).

CLI Updates

No single PR this release is purely a CLI change, but several touch sleap-nn's command-line surface directly:

  • sleap-nn eval --match_method semantic (#688) — new matching-free evaluation mode for semantic_segmentation models (whole-frame foreground IoU/clDice/boundary-IoU, no instance matching).
  • sleap-nn export now warns on a tiled model (#687) — exporting a model trained with tiling enabled to ONNX/TensorRT logs a warning that tiled export isn't supported yet (the export runs whole-frame instead of failing silently).
  • sleap-nn predict warns on ignored frame-filter flags for video input (#712) — --only_predicted_frames / --only_suggested_frames / --exclude_user_labeled / --only_labeled_frames require a .slp source (they filter by annotation status); passing one with a video now logs a warning naming the ignored flags instead of silently running full-video inference.
  • sleap-nn predict --paf_workers warns on unsupported model types (#711) — previously a silent no-op if the model type couldn't use the pipelined CPU-grouping path.
  • --max_instances/-n fixed for centroid-only and top-down models (#709) — see Fixes; this documented predict-time flag was being silently ignored for these model types.
  • sleap-nn track crash fixed when combining --tracking_clean_instance_count with --post_connect_single_breaks (#707) — see Fixes.

Fixes

  • Segmentation ModelCheckpoint no longer crashes when a monitored mask metric is NaN on validation epochs with no matched instances (#692).
  • sleap-nn predict output built from a pre-constructed provider (--only_suggested_frames, --only_labeled_frames, --only_predicted_frames, --exclude_user_labeled, --video_index) no longer crashes on save due to a dropped source video (#700).
  • Fixed an UnboundLocalError in sleap-nn track's legacy pipeline when combining --tracking_clean_instance_count with --post_connect_single_breaks, and a related crash on empty frame lists (#707).
  • --max_instances/-n overrides at predict time are now honored for centroid-only and top-down models (CentroidLayer), matching the fix already applied to bottom-up models; the sleap-nn predict startup log also now reflects actual peak_threshold/max_instances values (#709).
  • Fixed a crash (ValueError: 'centroid' is not in list) when the configured anchor part isn't a node in the pose skeleton (#702).
  • Training a dataset with zero usable training samples (e.g. a .slp with labeled frames but no user-labeled instances/centroids) now fails fast with a clear, actionable error instead of a cryptic IndexError: list index out of range deep in the trainer's first log line (#706).
  • Restored CPU-decode / GPU-inference overlap in sleap-nn predict — an earlier inference-pipeline refactor silently replaced the legacy background-thread frame readers with synchronous ones, so CPU video/.slp decode no longer overlapped the GPU forward pass; most noticeable on long or high-resolution videos where decode time is non-trivial. Also warns (instead of silently doing nothing) when --paf_workers is set on a model type that can't use the pipelined CPU-grouping path (#711).
  • Predictor.predict_streaming() no longer leaks predict-time postprocess overrides (peak_threshold, max_instances, etc.) across two interleaved streaming calls on the same Predictor object (#712).
  • sleap-nn predict now warns (instead of silently doing nothing) when a label-status frame filter (--only_labeled_frames, --only_suggested_frames, --exclude_user_labeled, --only_predicted_frames) is set for a non-.slp (video) input, since those flags have no annotation data to filter on there (#712).
  • Fixed LabelsProvider's frame-status filter priority order (when more than one only_*/exclude_* flag is set) to match the legacy pipeline's precedence (#712).
  • model_ckpt.monitor can now target pose-model (eval/val/mOKS, etc.) or centroid-model (eval/val/centroid_dist_avg, etc.) eval metrics — previously only segmentation eval metrics worked as a checkpoint monitor target; the other two callbacks computed their metrics but never exposed them to ModelCheckpoint, which crashed the first time a run tried to monitor one (#713).
  • Post-training final-split evaluation for pose models now uses the configured oks_stddev/oks_scale (matching the per-epoch training eval) instead of silently falling back to defaults — previously the saved metrics.<split>.npz could disagree with the training curves for any run that customized these (#713).
  • run_evaluation(user_labels_only=True) no longer mutates the caller's ground-truth Labels object in place; a second evaluation on the same Labels (e.g. with user_labels_only=False) no longer silently sees fewer instances than it should (#713).

Dependencies & Build

  • sleap-io >=0.9.2,<0.10.0 (was >=0.8.0,<0.9.0) — audited through v0.9.0–v0.9.2; the two breaking changes there (Identity.color removed, .category promoted from str to a Category object) don't affect any sleap-nn code path.
  • macOS CI no longer hangs to the 45-minute timeout or fails on MPS backend out of memory — an opt-in SLEAP_NN_DISABLE_MPS=1 env var (honored at sleap_nn import) forces CPU device selection; also usable by end users as an escape hatch for flaky MPS drivers (#695).

Upgrade Notes

  • Retrain any semantic_segmentation / bottomup_segmentation model trained before this release (#693 mask/image alignment fix).
  • Expect mask-evaluation recall/F1 to increase by default on labels files containing predicted instances (#694); pass user_labels_only=False to restore old behavior.
  • A training run that hits a frame-caching error will now stop immediately with a RuntimeError instead of continuing (#701) — check for corrupt videos / disk space if you see this.
  • If training a centroid model on a dataset that mixes user-labeled poses and UserCentroid annotations, set model_config.head_configs.centroid.confmaps.centroid_source explicitly ("user" or "computed") rather than leaving it unset (#704).
  • If you pin sleap-io, move to >=0.9.2,<0.10.0.

Changelog

  • #681: Reuse pretrained HuggingFace encoders as backbones (#680) (@talmo)
  • #682: Add clDice metric to eval/wandb + fix bottom-up viz (@talmo)
  • #687: Sliding-window tiling for high-res / small-object frames (@talmo)
  • #688: Add semantic_segmentation model type (whole-frame fg/bg) (@talmo)
  • #691: Configurable best-checkpoint metric + segmentation viz cleanup (#690) (@talmo)
  • #692: Don't crash ModelCheckpoint when a monitored seg metric is NaN (@talmo)
  • #693: Register whole-frame masks to the image grid (thread through shared preprocessing) (@talmo)
  • #694: Exclude PredictedInstance-linked GT masks in mask evaluation (@talmo)
  • #695: Disable MPS on mac CI to stop hangs + MPS-OOM failures (@talmo)
  • #698: Log confmap fg/bg MSE split as a training diagnostic (@talmo)
  • #700: Attach provider source videos so predict output is saveable (#699) (@tom21100227)
  • #701: Raise instead of only logging when frame caching fails (@gitttt-1234)
  • #702: Train centroid model from UserCentroid annotations (@tom21100227)
  • #703: Centroid model trains on pure-centroid frames (no pose instance) (@tom21100227)
  • #704: Single-source centroid targets (no user/computed mix) (@tom21100227)
  • #706: Clear error on empty dataset instead of cryptic IndexError (@tom21100227)
  • #707: Fix unbound corrected_lfs when combining tracking_clean_instance_count with post_connect_single_breaks (@gitttt-1234)
  • #708: Verify trainer_accelerator is available before training (@gitttt-1234)
  • #709: Honor predict-time --max_instances override in CentroidLayer (@alicup29)
  • #710: Bump version to 0.3.1 and sleap-io pin to >=0.9.2,<0.10.0 (@gitttt-1234)
  • #711: Restore CPU decode / GPU inference overlap in predict pipeline (@gitttt-1234)
  • #712: Predict-pipeline correctness fixes (overrides, CLI filters, LabelsProvider priority) (@gitttt-1234)
  • #713: Training/evaluation correctness fixes (callback_metrics, OKS eval config, GT mutation) (@gitttt-1234)

Contributors: @talmo, @tom21100227, @gitttt-1234, @alicup29

Full Changelog: v0.3.0...v0.3.1