sleap-roots-predict

A lightweight CLI and library that uses sleap-nn for prediction and produces artifacts in the format expected by sleap-roots pipelines. Intended to interoperate with the sleap-roots model registry.

Features

• SLEAP-NN Integration: Direct inference using SLEAP neural network models
• Flexible Video Processing: Create sleap_io.Video objects from image sequences for prediction
• Timelapse Experiment Support: Process plate-based imaging systems with automated metadata extraction
• Metadata Extraction: Parse datetime, plate numbers, and experimental conditions from standardized filenames
• Dual Output Formats: Generate either Video objects for direct prediction or compressed H5 files for storage
• Comprehensive Validation: Check image directories for consistency, datetime formats, and suffix patterns
• Batch Processing: Handle multi-plate experiments with CSV-based metadata
• GPU Acceleration: Automatic device selection (CUDA, MPS, or CPU) for optimal performance
• JSON Export: Save experiment results and metadata in JSON format for downstream analysis

Installation

We recommend doing this in an isolated environment.

Install using uv for faster, more reliable package management.

Create an isolated environment using uv:

# Create and activate a new virtual environment
uv venv
source .venv/bin/activate  # On Windows: .venv\Scripts\activate

Then install the package:

# CPU-only installation (all platforms)
uv pip install sleap-roots-predict[cpu]

# Windows with CUDA support
uv pip install sleap-roots-predict[windows_cuda]

# Linux with CUDA support
uv pip install sleap-roots-predict[linux_cuda]

# macOS (Apple Silicon or Intel)
uv pip install sleap-roots-predict[macos]

Otherwise just use pip.

Development Setup

# Install with development dependencies (replace with your platform's extra)
uv sync --extra dev --extra cpu         # For CPU-only
uv sync --extra dev --extra windows_cuda  # For Windows with CUDA
uv sync --extra dev --extra linux_cuda    # For Linux with CUDA
uv sync --extra dev --extra macos         # For macOS

# Run tests
uv run pytest

# Format code
uv run black sleap_roots_predict tests scripts

# Lint code
uv run ruff check sleap_roots_predict/ scripts/

# Check spelling
uv run codespell

Usage

Quick Start - High-Level API

from sleap_roots_predict import (
    process_timelapse_experiment,
    make_predictor,
    predict_on_video
)

# Process an entire timelapse experiment (video / H5 / metadata).
# NOTE: prediction within this flow is currently deferred — model_paths/device
# are accepted but ignored. Use the prediction API below to run inference.
results = process_timelapse_experiment(
    base_dir="path/to/experiment",
    metadata_csv="path/to/metadata.csv",
    experiment_name="my_experiment",
    output_dir="path/to/output",
    results_json="results.json"  # Save results as JSON
)

# Run prediction directly
predictor = make_predictor(
    model_paths=["path/to/model"],  # one dir per model (e.g. per root type)
    peak_threshold=0.2,
    batch_size=4,
    device="auto"  # Automatically selects GPU (CUDA/MPS) if available
)

# Create a Video object and run predictions
from sleap_roots_predict.video_utils import make_video_from_images
image_files = sorted(Path("path/to/images").glob("*.tif"))
video = make_video_from_images(image_files, greyscale=False)

predictions = predict_on_video(
    predictor,
    video,
    save_path="predictions.slp"  # Optional: save predictions
)

Advanced Usage - Utility Functions

For more control, you can import utility functions directly from their modules:

# Import utility functions from their modules
from sleap_roots_predict.plates_timelapse_experiment import (
    process_timelapse_image_directory,
    check_timelapse_image_directory,
    find_image_directories,
    extract_timelapse_metadata_from_filename,
    create_timelapse_metadata_dataframe
)

from sleap_roots_predict.video_utils import (
    make_video_from_images,
    load_images,
    convert_to_greyscale,
    save_array_as_h5,
    natural_sort
)

# Process a single directory
video, csv_path = process_timelapse_image_directory(
    source_dir="path/to/plate_001",
    experiment_name="exp1",
    treatment="control",
    num_plants=3,
    save_h5=False,  # Returns Video object
    output_dir="output/"
)

# Validate directories before processing
check_results = check_timelapse_image_directory(
    image_dir="path/to/plate_001",
    expected_suffix_pattern=r'^\d{3}$',  # e.g., "001", "002"
    min_images=5,
    max_images=1000,
    check_datetime=True,
    check_suffix_consistency=True
)

# Load and process images
image_paths = ["img1.tif", "img2.tif", "img3.tif"]
volume, filenames = load_images(image_paths, greyscale=True)
print(f"Loaded volume shape: {volume.shape}")  # (frames, height, width, channels)

# Convert RGB to greyscale with proper weights
grey_image = convert_to_greyscale(
    rgb_image,
    method="weights"  # Uses standard RGB weights (0.299, 0.587, 0.114)
)

# Save processed data as H5
save_array_as_h5(
    volume,
    output_path="processed_data.h5",
    compression="gzip",
    compression_opts=4
)

# Natural sorting for filenames with numbers
files = ["img_2.tif", "img_10.tif", "img_1.tif"]
sorted_files = natural_sort(files)
# Result: ["img_1.tif", "img_2.tif", "img_10.tif"]

Configuration

Model fetching from the wandb registry is configured through environment variables. The only required variable is WANDB_API_KEY — with it set, the warm worker fetches production models from the live registry out-of-the-box; everything else has a sensible default. Copy .env.example to .env and fill in your key to get started.

Variable	Required?	Default
WANDB_API_KEY	yes	— (a missing key fails loud on first use)
SRP_WANDB_ENTITY	no	eberrigan-salk-institute-for-biological-studies
SRP_WANDB_MODEL_REGISTRY	no	sleap-roots-models (the live production registry)
SRP_WANDB_MODEL_ALIAS	no	production
SRP_MODEL_CACHE_DIR	no	falls back to WANDB_CACHE_DIR, then wandb's default
SRP_DEVICE	no	auto-detect (cuda / mps / cpu)
SRP_PREDICT_CODE_SHA	no	recorded in each manifest's predict_code_sha; baked into the image at build time (build-arg → ENV), so operators don't set it
SRP_PREDICT_CONTAINER_DIGEST	no	recorded in each manifest's predict_container_digest (fail-soft to "")
ARGO_WORKFLOW_NAME	no — set by the cluster predictor template; leave unset locally	unset → only run_manifest.json is read (absent = unscoped discovery); set → this run's run_manifest.<ARGO_WORKFLOW_NAME>.json (or, during the rollout, the legacy run_manifest.json) is required, and its absence exits 1

Running the predict container

The service image runs the warm-batch predict CLI over a directory of staged scans:

docker run --rm -e WANDB_API_KEY=$WANDB_API_KEY \
  -v /path/to/scans:/in -v /path/to/out:/out \
  ghcr.io/talmolab/sleap-roots-predict:<tag> /in /out

Each scan is a directory of image frames with a co-located {scan_key}.scan_metadata.json sidecar (carrying the resolved {species, mode, age} params). This run's manifest, staged by an upstream pipeline stage — the per-run run_manifest.<ARGO_WORKFLOW_NAME>.json under Argo (the legacy run_manifest.json is still accepted during the rollout), or run_manifest.json with no run identity — scopes discovery to exactly its scan_keys: a leftover sidecar from a prior run is silently excluded rather than reprocessed. Under Argo a missing manifest fails the batch rather than processing every scan in the shared directory. An empty (zero-scan) input directory raises rather than silently succeeding, so a misconfigured stage-in mount never looks like a completed batch. The container loads models once, predicts every scan, and writes per scan out/{scan_key}/{scan_key}.predictions.json + named per-root .slp + a copy of the sidecar; the run manifest is also forwarded, under the name it was read, to the top level of out/, so the downstream trait-extraction stage — whose input directory is this output directory — stays scoped to the same run. All writes are atomic (a per-writer temp file + rename) so no reader ever observes a partially-written file. Skip-if-done compares a recomputed idempotency key against the prior run's own artifacts (no new storage needed) and skips only on an exact match, (re)predicting otherwise. A SIGTERM (Argo preemption) stops the batch at the next scan boundary rather than mid-scan. See the predict-container OpenSpec spec (openspec/specs/predict-container/spec.md) for the exact exit-code contract (0 success / 3 partial / default 1 staging-error-or-crash / 143 SIGTERM-terminated; 2 is reserved for a CLI usage error). GPU is used when available (nvidia.com/gpu); it also runs CPU-only. The same entrypoint is available as a library: from sleap_roots_predict import run_batch, or python -m sleap_roots_predict <in> <out>.

Rolling the image back: delete the forwarded run manifest

A rollback MUST also delete run_manifest.json from the output directory.

The forwarded manifest is sticky. It lives on the shared output mount, so rolling the predict image back to a version that never wrote it does not remove the copy already there. Trait-extraction reads that file to scope itself, so a leftover one pins every subsequent traits run to a frozen scan_keys set: newly staged scans are silently skipped, with no error and a green pipeline.

That is the mirror image of the bug this forwarding fixes (#39) — which over-processed, and was therefore noticeable. Under-processing is not. Predict itself is unaffected either way; it only ever reads the manifest from its input directory.

# on the shared mount, as part of any predict rollback
rm -f <output_dir>/run_manifest.json

Per-run run_manifest.<workflow>.json files need no rollback cleanup — each names one Argo workflow and no later run reads it. Never glob-delete them on the shared mount: a concurrent run's manifest would go too.

Once bloomctl writes per-run manifests, this version is predict's rollback floor. An older image reads only the stale, accumulated legacy run_manifest.json, forwards it, and re-scopes traits to that union — the contamination per-run manifests exist to prevent (talmolab/sleap-roots-pipeline#71). Roll the writer back first.

CI/CD

The project uses GitHub Actions for continuous integration and deployment:

Continuous Integration

On every pull request:

• Linting: black formatting, ruff linting, codespell
• Testing (CPU only — see GPU note below):
  • Ubuntu (latest)
  • Windows (latest)
  • macOS (Apple Silicon) - MPS is available but inference is forced to CPU on the hosted runners

CI does not run CUDA GPU tests (there is no GPU runner). The gpu-marked subset is a required local step in the pre-merge gate — run it on a CUDA/MPS machine: uv sync --extra dev --extra windows_cuda && uv run pytest -m gpu.

Parity Harness

sleap_roots_predict.parity measures sleap-nn's prediction accuracy against classic-SLEAP (the same already-trained weights) on real, human-labeled ground truth, gated behind a parity marker (WANDB_API_KEY + SRP_PARITY_DATA_DIR, a basename-search root for ground truth whose video paths were reorganized; deselected by default/in CI, like gpu/acceptance/wandb — run with uv run pytest -m parity -s). The decided tolerance and the full measured results across the 13 production models registered at the 2026-08-04 measurement (8 physical weight sets) are in docs/superpowers/specs/2026-08-04-define-parity-tolerance-results.json (design writeup: 2026-08-03-define-parity-tolerance-design.md). The 2026-09-29 re-run against the re-seeded (selector-shaped) registry — all 8 physical models within tolerance, unchanged from 2026-08-04 — is in 2026-09-29-parity-reseeded-registry-comparison.md. The 2026-10-04 run of the wheat and sorghum cards staged under the W&B alias candidate (read with SRP_WANDB_MODEL_ALIAS=candidate; the harness lists only that alias) is in 2026-10-04-parity-wheat-sorghum-comparison.md. All three cards passed, wheat crown marginally. That doc also covers frame-bootstrap ranges for the gate value and why distance_p95 is sensitive to the evaluator's greedy instance pairing on crowded crown and lateral frames.

Regenerate that report (needs WANDB_API_KEY and SRP_PARITY_DATA_DIR; lab-only — Windows + a Z: mapped network share): uv run python scripts/run_parity_harness.py --out <path>. --out is required — write to a new dated path; the 2026-08-04 JSON above is a pre-selectors snapshot and must not be overwritten. Pass --share-root if your mapped-share letter/path differs from the lab default. Commit the regenerated JSON as its own standalone commit. The 2026-08-04 baseline was produced on CPU (inferred: a CPU re-run reproduces it bit-for-bit; the report records no device). A GPU run is not bit-comparable to it — sync the cpu extra for a run meant to isolate registry/model changes from numeric noise, or re-run any model that moves on CPU before attributing the move.

Set SRP_PARITY_DATA_DIR=Z:/users/eberrigan/SLEAP (same tree as --share-root's default — build_basename_index walks it recursively, and it contains both the SLEAP_Rice and SLEAP_Soy subtrees the basename-search tier needs). Confirmed against the live share by the 2026-09-29 run (every model's ground truth resolved as on 2026-08-04); update this line if the layout moves.

Build and Publish

On release or manual trigger:

• PyPI Publishing: Automated wheel building and publishing using uv
• Trusted Publishing: Uses PyPI trusted publishing (no API tokens needed)
• TestPyPI Support: Manual workflow dispatch option for test publishing

To publish a new release: 0. For testing, manually trigger the workflow with TestPyPI option enabled

1. Update the semantic version in sleap_roots_predict/__init__.py.
2. Create a new GitHub release with the same semantic version tag (e.g., v0.1.0)
3. The workflow automatically builds and publishes to PyPI

Project Structure

sleap_roots_predict/
├── predict.py                      # SLEAP-NN prediction interface
├── model_selection.py              # Pure model-selection matcher (choose_models)
├── model_registry.py               # Model-card sources (Local + Wandb registry)
├── warm_worker.py                  # WarmModelWorker: resident predictors across scans
├── output_contract.py              # Per-scan output artifacts (.slp + predictions.json)
├── batch.py                        # Warm-batch container runner (run_batch, discover_scans)
├── run_manifest.py                 # Resolve this run's manifest; forward it input_dir -> output_dir
├── parity.py                       # A3-predict parity harness (ground truth + metrics)
├── __main__.py                     # `python -m sleap_roots_predict <in> <out>` CLI
├── video_utils.py                  # Core image processing utilities
├── plates_timelapse_experiment.py  # Timelapse experiment processing
└── __init__.py                     # Package exports and version

tests/
├── test_predict.py                     # Prediction module tests
├── test_param_resolution.py            # choose_models round-trip vs. contracts' resolve_params (offline)
├── test_model_selection.py             # Model-selection matcher tests
├── test_model_registry.py              # Card-source tests (offline + gated wandb)
├── test_warm_worker.py                 # Warm worker tests (real CPU inference)
├── test_output_contract.py             # Output-contract writer/batch tests (real CPU inference)
├── test_batch.py                       # Batch runner / CLI tests (real CPU inference)
├── test_run_manifest.py                # Run-manifest resolution + forward-copy tests (offline)
├── test_parity.py                      # Parity harness tests (offline + gated `parity` marker)
├── test_predict_container_packaging.py # Console-script + docker-workflow guards
├── test_public_api.py                  # Public-surface import test
├── test_video_utils.py                 # Video utilities tests
└── conftest.py                         # Shared test fixtures

scripts/
└── run_parity_harness.py           # Regenerate the parity results JSON (lab-only, credentialed)

.github/workflows/
├── ci.yml                      # Continuous integration
├── docker-build.yml            # Build/push the GHCR service image
└── publish.yml                 # PyPI publishing workflow

Contributing

1. Fork the repository
2. Create a feature branch (git checkout -b feature/amazing-feature)
3. Make your changes and add tests
4. Ensure all tests pass (uv run pytest)
5. Format your code (uv run black .)
6. Check linting (uv run ruff check .)
7. Commit your changes (git commit -m 'Add amazing feature')
8. Push to your branch (git push origin feature/amazing-feature)
9. Open a Pull Request

License

See LICENSE file for details.