diff --git a/CLAUDE.md b/CLAUDE.md index 0ad0f877..a263f0f0 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -151,6 +151,7 @@ cache invalidation logic (version mismatch, pose hash, distance scale). - `__init__.py` files: unused imports (F401) are allowed - **Naming**: `PascalCase` classes, `snake_case` functions/methods/files, `UPPER_SNAKE_CASE` constants, `_` prefix for private members +- Prefer American English spelling (e.g., "behavior" not "behaviour"). ## Coding Standards diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index a7acea3b..fa06363f 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -109,7 +109,7 @@ For detailed information about JABS architecture, development setup, and impleme - **General questions**: Open an issue on GitHub - **Security issues**: Email jabs@jax.org (do not open public issues) -- **Development questions**: See [DEVELOPMENT.md](docs/DEVELOPMENT.md) or contact jabs@jax.org +- **Development questions**: See [DEVELOPMENT.md](docs/development/DEVELOPMENT.md) or contact jabs@jax.org ## Code of Conduct diff --git a/docs/development/jabs-nwb-format.md b/docs/development/jabs-nwb-format.md new file mode 100644 index 00000000..0f151e29 --- /dev/null +++ b/docs/development/jabs-nwb-format.md @@ -0,0 +1,598 @@ +# JABS NWB Format + +This document describes the NWB files produced by JABS. It covers the two output modes +(combined and per-identity), the full file layout, how animal pose, static objects, and +dynamic objects are stored, the `jabs_metadata` scratch field, and why the +`ndx-multisubjects` extension is not currently used. + +JABS NWB files use the [ndx-pose 0.2](https://github.com/rly/ndx-pose) extension for +all pose and object data. + +--- + +## Output modes + +JABS can write NWB in two modes, selectable at export time. + +### Combined file (default) + +All identities from a single recording session are written into one NWB file. This is +the simplest output and potentially the most compatible with third-party NWB tooling. + +``` +session.nwb + └── all identities, all objects +``` + +**When to use:** sharing data with collaborators, archiving, downstream analysis that +needs all animals in one place. + +### Per-identity files + +One NWB file is written per animal. The output path is used as a naming template; the +combined file is never created. Static and dynamic objects are written to every +per-identity file identically (they are session-level, not animal-level data). + +``` +session_subject_0.nwb ← identity 0 + all objects +session_subject_1.nwb ← identity 1 + all objects +session_subject_2.nwb ← identity 2 + all objects +``` + +Identity names in the filenames come from `external_ids` in the pose file (sanitized +for HDF5 compatibility) or fall back to `subject_0`, `subject_1`, … when no external +IDs are present. + +**When to use:** downstream workflows that require one NWB file per animal (e.g. tools +that expect a single `Subject` in each file). + +#### Reading per-identity files + +The JABS reader re-assembles per-identity files transparently. Point it at **any one** +sibling file; it detects the `per_identity_files` flag in `jabs_metadata`, globs for +siblings matching `{base_stem}_*.nwb` in the same directory, filters to those that +share the same `total_identities` count, sorts by `source_identity_index`, and +concatenates them into a single `PoseData` with all identities in their original order. + +``` +# read any sibling — result is identical +pose_data = load("session_subject_0.nwb", PoseData) +pose_data = load("session_subject_2.nwb", PoseData) +``` + +Validation ensures the expected number of sibling files are present before merging; a +`ValueError` is raised if any file is missing. + +--- + +## Full NWB layout + +The layout below shows a combined file containing two animal identities, two static +objects (`corners`, `lixit`), and one dynamic object (`fecal_boli`). + +``` +NWBFile +├── processing/ +│ └── behavior/ [ProcessingModule] +│ ├── Skeletons/ [Skeletons container] +│ │ ├── subject/ Skeleton — animal keypoints + edges +│ │ ├── corners/ Skeleton — static object (4 nodes) +│ │ ├── lixit/ Skeleton — static object (1 or 3 nodes) +│ │ └── fecal_boli/ Skeleton — dynamic object (max_count nodes) +│ │ +│ ├── subject_0/ [PoseEstimation] animal identity 0 +│ │ ├── nose/ [PoseEstimationSeries] num_frames timestamps +│ │ ├── left_ear/ +│ │ └── ... +│ │ +│ ├── subject_1/ [PoseEstimation] animal identity 1 +│ │ ├── nose/ +│ │ └── ... +│ │ +│ ├── corners/ [PoseEstimation] static object +│ │ ├── corners_0/ [PoseEstimationSeries] 1 timestamp +│ │ ├── corners_1/ +│ │ ├── corners_2/ +│ │ └── corners_3/ +│ │ +│ ├── lixit/ [PoseEstimation] static object +│ │ └── lixit_0/ [PoseEstimationSeries] 1 timestamp +│ │ +│ ├── fecal_boli/ [PoseEstimation] dynamic object +│ │ ├── fecal_boli_0/ [PoseEstimationSeries] n_predictions timestamps +│ │ ├── fecal_boli_1/ +│ │ └── ... +│ │ +│ ├── jabs_identity_mask [TimeSeries] uint8 identity presence mask +│ ├── jabs_bounding_boxes_subject_0 [TimeSeries] optional, one per identity +│ └── jabs_bounding_boxes_subject_1 [TimeSeries] optional, one per identity +│ +└── scratch/ + └── jabs_metadata/ [ScratchData] JSON string (see below) +``` + +In a per-identity file the layout is identical, except only one animal identity +container is present and `jabs_identity_mask` / `jabs_bounding_boxes_` cover +that identity only. + +--- + +## Animal pose + +Each animal identity is a `PoseEstimation` container in `processing/behavior`. The +container name is the sanitized external ID from the pose file, or `subject_{i}` when +no external IDs are available. + +A single `Skeleton` named `subject` (or overridden via `skeleton_name`) is shared by +all animal identities and stored in the `Skeletons` container. + +### PoseEstimationSeries fields (per keypoint) + +| Field | Value | +|-------------------------|-----------------------------------------------------------------------------------| +| `name` | Keypoint name (e.g. `"nose"`, `"left_ear"`) | +| `data` | shape `(num_frames, 2)` — `(x, y)` coordinates in pixels | +| `rate` | Frames per second (float) | +| `unit` | `"pixels"` | +| `reference_frame` | `"Top-left corner of video frame, x increases rightward, y increases downward"` | +| `confidence` | shape `(num_frames,)` — `0.0` = missing keypoint, `> 0.0` = valid | +| `confidence_definition` | `"0.0=invalid/missing keypoint, >0.0=valid keypoint"` | + +### Identity mask + +`jabs_identity_mask` is a `TimeSeries` that records whether each identity is present in +each frame. + +| Mode | Shape stored in file | Shape returned by reader | +|------------------|-------------------------------|----------------------------------| +| Combined | `(num_frames, num_identities)` | `(num_identities, num_frames)` | +| Per-identity | `(num_frames,)` | `(1, num_frames)` | + +### Bounding boxes (optional) + +When the pose file contains bounding box data, one `TimeSeries` per identity is written +with the name `jabs_bounding_boxes_{identity_name}`. This naming makes the containers +self-describing — no external index mapping is required. + +| Property | Value | +|----------------------|-----------------------------------------------------------| +| Name | `jabs_bounding_boxes_{identity_name}` (one per identity) | +| Shape stored in file | `(num_frames, 2, 2)` | +| Shape returned by reader | `(num_identities, num_frames, 2, 2)` (all stacked) | +| Both modes | Same per-identity shape in combined and per-identity files | + +Format: `[[upper_left_x, upper_left_y], [lower_right_x, lower_right_y]]` in pixels. + +The reader looks for keys `jabs_bounding_boxes_{name}` for each name in `identity_names` +(from `jabs_metadata`). If all are present, they are stacked in identity order to form +the returned array. If any are missing, `bounding_boxes` is `None`. + +--- + +## Static objects + +Static objects are fixed-position spatial landmarks that do not move during a session. +They are read from `static_objects/` in JABS pose HDF5 files (v5+). + +Common static objects: + +| Object | Shape | Description | +|---------------|----------------------|---------------------------------------------------| +| `corners` | `(4, 2)` | Four corners of the arena | +| `lixit` | `(1, 2)` or `(3, 2)` | Water spout — single tip, or tip + left + right | +| `food_hopper` | `(4, 2)` | Four corners of the food hopper opening | + +### NWB representation + +Each static object is a `PoseEstimation` container with a **single timestamp +(`t = 0.0 s`)**, one `PoseEstimationSeries` per keypoint, and a dedicated `Skeleton` +in the `Skeletons` container. Nodes are named `{object_name}_{i}` (zero-indexed). + +**PoseEstimationSeries fields:** + +| Field | Value | +|-------------------------|-----------------------------------------------------------------------------------| +| `name` | `{object_name}_{i}` | +| `data` | shape `(1, 2)` — the `(x, y)` coordinate | +| `timestamps` | `[0.0]` | +| `unit` | `"pixels"` | +| `reference_frame` | `"Top-left corner of video frame, x increases rightward, y increases downward"` | +| `confidence` | `[1.0]` | +| `confidence_definition` | `"Static landmark; confidence is always 1.0"` | + +*JABS pose files carry no confidence values for static objects; `1.0` is a placeholder. +Consumers should ignore the confidence field for static objects.* + +**PoseEstimation fields:** + +| Field | Value | +|-------------------|---------------------------------------------------------| +| `name` | `{object_name}` (e.g. `"corners"`) | +| `description` | `"Static object: {object_name}"` | +| `source_software` | `"JABS"` | +| `skeleton` | The matching `Skeleton` from the `Skeletons` container | + +### Example — `corners` (4 keypoints) + +``` +Skeletons/ + corners/ + nodes: ["corners_0", "corners_1", "corners_2", "corners_3"] + +processing/behavior/ + corners/ PoseEstimation + corners_0/ PoseEstimationSeries + data: [[10.0, 20.0]] shape (1, 2) + timestamps: [0.0] + confidence: [1.0] + corners_1/ + data: [[300.0, 20.0]] + timestamps: [0.0] + confidence: [1.0] + corners_2/ + data: [[10.0, 300.0]] + timestamps: [0.0] + confidence: [1.0] + corners_3/ + data: [[300.0, 300.0]] + timestamps: [0.0] + confidence: [1.0] +``` + +### Example — `lixit` (3-keypoint variant) + +``` +Skeletons/ + lixit/ + nodes: ["lixit_0", "lixit_1", "lixit_2"] + +processing/behavior/ + lixit/ PoseEstimation + lixit_0/ tip + data: [[62.0, 166.0]] + timestamps: [0.0] + confidence: [1.0] + lixit_1/ left side + data: [[65.0, 160.0]] + timestamps: [0.0] + confidence: [1.0] + lixit_2/ right side + data: [[60.0, 172.0]] + timestamps: [0.0] + confidence: [1.0] +``` + +--- + +## Dynamic objects + +Dynamic objects are objects whose position or count may change over time. Unlike animal +pose, predictions are not made every frame — only a sparse subset of frames is sampled. +Dynamic objects are introduced in JABS pose format v7. + +The HDF5 pose v7 format stores dynamic objects under `dynamic_objects/[name]/`: + +| Dataset | Shape | Description | +|------------------|----------------------------------------------------------------------------------------------------|-----------------------------------------------------| +| `points` | `(n_predictions, max_count, 2)` for single-keypoint; `(n_predictions, max_count, n_keypoints, 2)` for multi-keypoint | Keypoint coordinates; axis order set by `axis_order` attribute | +| `counts` | `(n_predictions,)` | Number of valid object instances at each prediction | +| `sample_indices` | `(n_predictions,)` | Frame indices at which predictions were made | + +The `points` dataset carries an optional HDF5 attribute: + +| Attribute | Values | Default | Meaning | +|--------------|-------------------|---------|--------------------------------------------------------------------------| +| `axis_order` | `"xy"` or `"yx"` | `"yx"` | Coordinate ordering in the file. JABS always normalizes to `(x, y)` on read. | + +> **Note:** The `"yx"` default matches the fecal boli network, which was trained with +> HRNet and stores coordinates in row-major (y, x) order. + +JABS normalizes all dynamic object `points` arrays to 4-D +`(n_predictions, max_count, n_keypoints, 2)` internally. Single-keypoint objects stored +as 3-D in HDF5 are expanded to `n_keypoints=1` on read. + +### NWB representation + +Each dynamic object is a `PoseEstimation` container with `n_predictions` **irregular +timestamps**. One `PoseEstimationSeries` is written per **instance slot × keypoint** +combination, with a dedicated `Skeleton` in the `Skeletons` container. + +#### Timestamps + +Frame indices are converted to seconds for the NWB time axis: + +``` +timestamps[p] = sample_indices[p] / fps + +# Recover on read: +sample_indices[p] = round(timestamps[p] * fps) +``` + +#### Instance slot validity via confidence + +Not all `max_count` slots are occupied at every prediction. Occupancy is encoded in +`confidence`, using the same field used for keypoint validity in animal pose: + +``` +confidence[p] = 1.0 if counts[p] > slot_index + = 0.0 otherwise +``` + +The same confidence value applies to every keypoint within a slot. `counts` can be +recovered on read by summing slots where `confidence > 0` at each prediction timestamp. +Coordinate values in empty slots are meaningless padding and must not be used by +consumers. + +#### Node naming convention + +| Condition | Node name pattern | Example | +|--------------------|----------------------|--------------------------------------| +| `n_keypoints == 1` | `{name}_{slot}` | `fecal_boli_0` | +| `n_keypoints > 1` | `{name}_{slot}_{kp}` | `door_0_0`, `door_0_1`, `door_1_0` | + +**PoseEstimationSeries fields:** + +| Field | Value | +|-------------------------|------------------------------------------------------------------------------------| +| `name` | `{name}_{slot}` or `{name}_{slot}_{kp}` | +| `data` | shape `(n_predictions, 2)` — `(x, y)` coordinates | +| `timestamps` | `sample_indices / fps` | +| `unit` | `"pixels"` | +| `reference_frame` | `"Top-left corner of video frame, x increases rightward, y increases downward"` | +| `confidence` | `1.0` if slot occupied, `0.0` otherwise | +| `confidence_definition` | `"1.0=valid object instance in this slot, 0.0=slot unoccupied at this prediction"` | + +**PoseEstimation fields:** + +| Field | Value | +|-------------------|---------------------------------------------------------| +| `name` | `{object_name}` (e.g. `"fecal_boli"`) | +| `description` | `"Dynamic object: {object_name}"` | +| `source_software` | `"JABS"` | +| `skeleton` | The matching `Skeleton` from the `Skeletons` container | + +### Example — `fecal_boli` (single keypoint per instance, up to 3 instances) + +50 predictions were made; up to 3 fecal boli are visible at once. + +``` +Skeletons/ + fecal_boli/ + nodes: ["fecal_boli_0", "fecal_boli_1", "fecal_boli_2"] + +processing/behavior/ + fecal_boli/ PoseEstimation + fecal_boli_0/ PoseEstimationSeries — slot 0 + data: shape (50, 2) + timestamps: [t_0, t_1, ..., t_49] # sample_indices / fps + confidence: [1.0, 1.0, 0.0, ...] # 1.0 where counts > 0 + fecal_boli_1/ slot 1 + data: shape (50, 2) + timestamps: [t_0, t_1, ..., t_49] + confidence: [1.0, 0.0, 0.0, ...] # 1.0 where counts > 1 + fecal_boli_2/ slot 2 + data: shape (50, 2) + timestamps: [t_0, t_1, ..., t_49] + confidence: [0.0, 0.0, 0.0, ...] # 1.0 where counts > 2 +``` + +At prediction `p=0`, `counts[0]=2`: slots 0 and 1 are valid, slot 2 is padding. +At prediction `p=1`, `counts[1]=1`: only slot 0 is valid. + +### Example — multi-keypoint dynamic object (2 keypoints per instance, up to 2 instances) + +A hypothetical `door` object with 2 keypoints per instance (left edge, right edge), +maximum 2 doors in the arena, 30 predictions. + +``` +Skeletons/ + door/ + nodes: ["door_0_0", "door_0_1", "door_1_0", "door_1_1"] + # slot 0 slot 0 slot 1 slot 1 + # kp 0 kp 1 kp 0 kp 1 + +processing/behavior/ + door/ PoseEstimation + door_0_0/ slot 0, keypoint 0 — left edge of door 0 + data: shape (30, 2) + timestamps: [t_0, ..., t_29] + confidence: [1.0, 1.0, ...] # 1.0 where counts > 0 + door_0_1/ slot 0, keypoint 1 — right edge of door 0 + data: shape (30, 2) + timestamps: [t_0, ..., t_29] + confidence: [1.0, 1.0, ...] # same as door_0_0 — slot-level validity + door_1_0/ slot 1, keypoint 0 — left edge of door 1 + data: shape (30, 2) + timestamps: [t_0, ..., t_29] + confidence: [0.0, 1.0, ...] # 1.0 where counts > 1 + door_1_1/ slot 1, keypoint 1 — right edge of door 1 + data: shape (30, 2) + timestamps: [t_0, ..., t_29] + confidence: [0.0, 1.0, ...] # same as door_1_0 +``` + +--- + +## `jabs_metadata` scratch field + +Every JABS NWB file contains a `ScratchData` object named `jabs_metadata` in the NWB +`scratch` space. Its `data` field is a JSON string carrying all JABS-specific metadata +needed for a lossless round-trip. Standard NWB fields alone are insufficient because +pynwb returns `PoseEstimationSeries` in alphabetical order from HDF5, which would +otherwise scramble the keypoint ordering. + +### Keys + +| Key | Type | Present | Description | +|-------------------------|-------------------------|------------------------------|-------------| +| `format_version` | `int` | Always | JABS NWB format version. Currently `1`. | +| `identity_names` | `list[str]` | Always | Ordered list of `PoseEstimation` container names that are animal identities. Defines identity order on read. | +| `num_identities` | `int` | Always | Total number of animal identities in the recording session. In per-identity mode this equals `total_identities`; the file itself contains only one identity. | +| `body_parts` | `list[str]` | Always | Ordered list of keypoint names for animal skeletons. Preserves original write order, since HDF5 returns groups alphabetically. | +| `cm_per_pixel` | `float \| null` | Always | Pixel-to-centimetre scale factor. `null` if not available in the source pose file. | +| `external_ids` | `list[str] \| null` | Always | Original external identity names from the pose file (e.g. mouse cage IDs). `null` if the pose file had no external IDs. | +| `subjects` | `dict[str, dict] \| null` | Always | Per-identity subject metadata keyed by identity name. `null` if no subject metadata is available. Inner dict may contain `subject_id`, `sex`, `genotype`, `strain`, `age`, `weight`, `species`, and `description`. Values are `null` when not available. | +| `metadata` | `dict` | Always | Provenance metadata from the source pose file. Includes `source_file`, `pose_format_version`, and optionally `source_file_hash`. | +| `static_object_names` | `list[str]` | When static objects present | Names of all `PoseEstimation` containers that are static objects. | +| `dynamic_object_names` | `list[str]` | When dynamic objects present | Names of all `PoseEstimation` containers that are dynamic objects. | +| `dynamic_object_shapes` | `dict[str, [int, int]]` | When dynamic objects present | Maps each dynamic object name to `[max_count, n_keypoints]`. Required to reconstruct the 4-D `points` array `(n_predictions, max_count, n_keypoints, 2)` from the flat series list on read. | +| `per_identity_files` | `bool` | Per-identity mode only | `true` if this file is one of a set of per-identity NWB files. | +| `source_identity_index` | `int` | Per-identity mode only | Zero-based index of the identity in this file within the original multi-identity dataset. Used to restore original identity order when merging siblings. | +| `total_identities` | `int` | Per-identity mode only | Total number of identity files in the set. Used to validate that all siblings are present before merging. | + +### Example — combined file with two identities, static objects, and dynamic objects + +```json +{ + "format_version": 1, + "identity_names": ["subject_0", "subject_1"], + "num_identities": 2, + "body_parts": ["nose", "left_ear", "right_ear", "base_neck", "left_front_paw", + "right_front_paw", "center_spine", "left_rear_paw", "right_rear_paw", + "base_tail", "mid_tail", "tip_tail"], + "cm_per_pixel": 0.043, + "external_ids": null, + "subjects": { + "subject_0": { + "subject_id": "M123", + "sex": "M", + "genotype": "WT", + "strain": "C57BL/6J", + "age": "P70D", + "weight": null, + "species": "Mus musculus", + "description": null + }, + "subject_1": { + "subject_id": "M124", + "sex": "F", + "genotype": "Shank3+/-", + "strain": "C57BL/6J", + "age": "P72D", + "weight": null, + "species": "Mus musculus", + "description": null + } + }, + "metadata": { + "source_file": "/data/session_pose_est_v7.h5", + "pose_format_version": 7, + "source_file_hash": "a3f1c8..." + }, + "static_object_names": ["corners", "lixit"], + "dynamic_object_names": ["fecal_boli"], + "dynamic_object_shapes": { + "fecal_boli": [3, 1] + } +} +``` + +### Example — per-identity file (identity 1 of 3) + +```json +{ + "format_version": 1, + "identity_names": ["subject_1"], + "num_identities": 3, + "body_parts": ["nose", "left_ear", "..."], + "cm_per_pixel": 0.043, + "external_ids": null, + "subjects": { + "subject_0": { "subject_id": "M123", "sex": "M", "genotype": "WT", "..." : "..." }, + "subject_1": { "subject_id": "M124", "sex": "F", "genotype": "Shank3+/-", "..." : "..." }, + "subject_2": { "subject_id": "M125", "sex": "M", "genotype": "WT", "..." : "..." } + }, + "metadata": { "source_file": "...", "pose_format_version": 7 }, + "static_object_names": ["corners", "lixit"], + "dynamic_object_names": ["fecal_boli"], + "dynamic_object_shapes": { "fecal_boli": [3, 1] }, + "per_identity_files": true, + "source_identity_index": 1, + "total_identities": 3 +} +``` + +> **Note:** Per-identity files store the full `subjects` dict for all identities, not +> just the one identity in that file. This makes each file self-contained and means the +> reader can recover complete subject metadata from any sibling. + +--- + +## Read-path disambiguation + +All `PoseEstimation` containers in `processing/behavior` are classified using three +explicit lists from `jabs_metadata`. Each container name appears in exactly one list: + +``` +all PoseEstimation containers in behavior + │ + ├── name in identity_names → animal identity + ├── name in static_object_names → static object + └── name in dynamic_object_names → dynamic object (use dynamic_object_shapes to reconstruct) +``` + +Using explicit lists rather than inference rules (e.g. "everything else is a static +object") ensures the classification remains correct if new container types are added in +future format versions. + +--- + +## Coordinate system + +All coordinates in JABS NWB files use the following convention: + +| Property | Value | +|----------|------------------------------------------| +| Origin | Top-left corner of the video frame | +| x axis | Increases rightward (column direction) | +| y axis | Increases downward (row direction) | +| Units | Pixels | + +This applies to animal keypoints, static object points, and dynamic object points. +Dynamic object coordinates are stored as `(y, x)` in the HDF5 pose file by default (the +`axis_order` attribute controls this); JABS flips them to `(x, y)` on read before +writing to NWB. NWB files always contain `(x, y)` order. + +--- + +## Why ndx-multisubjects is not used + +[ndx-multisubjects](https://github.com/nehatk17/ndx-multisubjects) is a NWB extension +that adds multi-subject support to NWB files through three new types: + +- **`SubjectsTable`** — a `DynamicTable` with one row per animal, storing standard + subject fields (`subject_id`, `sex`, `genotype`, `strain`, `age`, `weight`, etc.) +- **`NdxMultiSubjectsNWBFile`** — a subclass of `NWBFile` that embeds the + `SubjectsTable` in general metadata +- **`SelectSubjectsContainer`** — an `NWBDataInterface` that links data to a subject + subset via a `DynamicTableRegion` + +Standard NWB only supports a single `Subject` on `NWBFile.subject`, so this extension +addresses a real gap for multi-animal recordings. JABS evaluated it and chose not to +adopt it for the following reasons: + +**1. `NdxMultiSubjectsNWBFile` is a non-standard `NWBFile` subclass.** +Any tool that opens a JABS NWB file without the extension installed will either fail +or silently lose the subjects table. The core value of NWB is that files are readable +by the broader ecosystem using only pynwb. This extension undermines that guarantee. + +**2. `SelectSubjectsContainer` does not compose cleanly with ndx-pose.** +The extension's model for associating data with subjects requires wrapping data +containers inside `SelectSubjectsContainer`. JABS `PoseEstimation` containers live +directly in `processing/behavior` and are already named by identity — adding a wrapper +layer would significantly restructure the layout without a commensurate benefit. + +**3. The extension is still Beta.** +ndx-multisubjects is published on PyPI (v0.1.1, November 2025) and sleap-io merged +support for it in December 2025, so the ecosystem is beginning to form. However, at +Beta (0.1.x) the API may still change, and broad adoption across NWB tooling has not +yet occurred. Since JABS NWB support is itself under active development, taking a +dependency on an immature extension adds unnecessary coupling at this stage. + +**What JABS does instead.** +Per-animal biological metadata (`subject_id`, `sex`, `genotype`, `strain`, `age`, +`weight`, `species`, `description`) can be stored in the `subjects` key of `jabs_metadata`. +This keeps the file readable by any standard NWB tool while preserving the metadata in +a structured, machine-readable form. If ndx-multisubjects stabilises and achieves +broader adoption, migrating to it would be straightforward since all the underlying +data is already present. \ No newline at end of file diff --git a/packages/jabs-core/src/jabs/core/abstract/pose_est.py b/packages/jabs-core/src/jabs/core/abstract/pose_est.py index d0d0e7bf..29cd30d3 100644 --- a/packages/jabs-core/src/jabs/core/abstract/pose_est.py +++ b/packages/jabs-core/src/jabs/core/abstract/pose_est.py @@ -8,6 +8,7 @@ import numpy as np from shapely.geometry import MultiPoint +from jabs.core.types import DynamicObjectData from jabs.core.utils import hash_file MINIMUM_CONFIDENCE = 0.3 @@ -132,7 +133,8 @@ def __init__(self, file_path: Path, cache_dir: Path | None = None, fps: int = 30 self._hash = hash_file(file_path) self._fps = fps - self._static_objects = {} + self._static_objects: dict = {} + self._dynamic_objects: dict[str, DynamicObjectData] = {} # check cache version, if it doesn't match, clear the cache file for this pose file if self._cache_dir is not None and not self.check_cache_version(): @@ -270,6 +272,28 @@ def static_objects(self): """get static objects from the pose file""" return self._static_objects + @property + def dynamic_objects(self) -> dict[str, DynamicObjectData]: + """Get dynamic objects from the pose file. + + Returns: + Mapping of object name to jabs.core.types.DynamicObjectData. + Empty for pose versions that do not support dynamic objects (v2-v6). + """ + return self._dynamic_objects + + def get_dynamic_object(self, name: str) -> DynamicObjectData | None: + """Get a named dynamic object from the pose file. + + Args: + name: Name of the dynamic object (e.g. "fecal_boli"). + + Returns: + DynamicObjectData for the requested object, or None if not present. + Always None for pose versions that do not support dynamic objects (v2-v6). + """ + return self._dynamic_objects.get(name) + def get_identity_convex_hulls(self, identity): """get a list of length #frames containing convex hulls for the given identity. @@ -419,3 +443,11 @@ def _cache_file_path(self) -> Path | None: return None filename = self._path.name.replace(".h5", "_cache.h5") return self._cache_dir / filename + + def get_bounding_boxes(self, identity: int) -> np.ndarray | None: + """Get bounding box array for an identity index. + + Default implementation returns None, indicating no bounding box data. + Bounding boxes are not available in pose file versions <8, this is included for interface consistency. + """ + return None diff --git a/packages/jabs-core/src/jabs/core/types/__init__.py b/packages/jabs-core/src/jabs/core/types/__init__.py index 7917c6c1..e991fe63 100644 --- a/packages/jabs-core/src/jabs/core/types/__init__.py +++ b/packages/jabs-core/src/jabs/core/types/__init__.py @@ -3,7 +3,7 @@ from .inference import AggregationSpec, InferenceSampling from .keypoints import FrameKeypoints, FrameKeypointsData, KeypointAnnotation from .model import ModelInfo -from .pose import PoseData +from .pose import DynamicObjectData, PoseData from .prediction import BehaviorPrediction, ClassifierMetadata from .results import InferenceRunMetadata, KeypointInferenceResult from .video import VideoInfo @@ -12,6 +12,7 @@ "AggregationSpec", "BehaviorPrediction", "ClassifierMetadata", + "DynamicObjectData", "FrameKeypoints", "FrameKeypointsData", "InferenceRunMetadata", diff --git a/packages/jabs-core/src/jabs/core/types/pose.py b/packages/jabs-core/src/jabs/core/types/pose.py index 3cb31240..e65db761 100644 --- a/packages/jabs-core/src/jabs/core/types/pose.py +++ b/packages/jabs-core/src/jabs/core/types/pose.py @@ -2,6 +2,33 @@ from typing import Any import numpy as np +import numpy.typing as npt + + +@dataclass(frozen=True) +class DynamicObjectData: + """Container for a single dynamic object read from a pose file. + + Dynamic objects are objects that may change position or count over time, but + are not predicted every frame. Only frames listed in sample_indices have + valid predictions. Coordinates are always stored in (x, y) order. + + points is always 4-D regardless of how many keypoints each object + instance has. Single-keypoint objects (e.g. fecal boli) have + n_keypoints=1 after normalization on read. + + Attributes: + points: Detected keypoint coordinates in (x, y) order, shape + (n_predictions, max_count, n_keypoints, 2). + counts: Number of valid detected objects for each prediction, shape + (n_predictions,). + sample_indices: Frame indices at which each prediction was made, shape + (n_predictions,). + """ + + points: npt.NDArray[np.float64] + counts: npt.NDArray[np.int64] + sample_indices: npt.NDArray[np.int64] @dataclass(frozen=True) @@ -20,8 +47,13 @@ class PoseData: Format is [[upper_left_x, upper_left_y], [lower_right_x, lower_right_y]]. segmentation_data: Optional segmentation masks or data. static_objects: Dictionary of static objects (e.g., 'lixit') and their positions. + dynamic_objects: Dictionary of dynamic objects (e.g., 'fecal_boli') and their data. external_ids: Optional list of external identifiers for each identity. Maps an identity index to an external ID string. + subjects: Optional per-animal biological metadata, keyed by identity name + (matching the values in external_ids). Each value is a free-form + dict; standard keys are subject_id, sex, genotype, + strain, age, weight, species, description. metadata: Dictionary for any additional provenance or experimental metadata. """ @@ -35,7 +67,9 @@ class PoseData: bounding_boxes: np.ndarray | None = None segmentation_data: np.ndarray | None = None static_objects: dict[str, np.ndarray] = field(default_factory=dict) + dynamic_objects: dict[str, DynamicObjectData] = field(default_factory=dict) external_ids: list[str] | None = None + subjects: dict[str, dict] | None = None metadata: dict[str, Any] = field(default_factory=dict) def __post_init__(self): diff --git a/packages/jabs-io/pyproject.toml b/packages/jabs-io/pyproject.toml index cee2250f..9dbfc871 100644 --- a/packages/jabs-io/pyproject.toml +++ b/packages/jabs-io/pyproject.toml @@ -14,7 +14,6 @@ authors = [ dependencies = [ "jabs-core", - "ndx-pose>=0.2.2", "numpy>=2.0.0,<3.0.0", ] diff --git a/packages/jabs-io/src/jabs/io/internal/pose/__init__.py b/packages/jabs-io/src/jabs/io/internal/pose/__init__.py index 33c2c49c..7f79ac79 100644 --- a/packages/jabs-io/src/jabs/io/internal/pose/__init__.py +++ b/packages/jabs-io/src/jabs/io/internal/pose/__init__.py @@ -1,4 +1,4 @@ -"""Pose estimation NWB adapters.""" +"""Pose estimation adapters (NWB requires the [nwb] extra).""" from jabs.io.internal.pose.nwb import PoseNWBAdapter diff --git a/packages/jabs-io/src/jabs/io/internal/pose/nwb.py b/packages/jabs-io/src/jabs/io/internal/pose/nwb.py index a03250d8..fa265455 100644 --- a/packages/jabs-io/src/jabs/io/internal/pose/nwb.py +++ b/packages/jabs-io/src/jabs/io/internal/pose/nwb.py @@ -1,5 +1,7 @@ """NWB adapter for PoseData using ndx-pose.""" +from __future__ import annotations + import datetime import json import logging @@ -9,12 +11,18 @@ import numpy as np import numpy.typing as npt -from ndx_pose import PoseEstimation, PoseEstimationSeries, Skeleton, Skeletons -from pynwb import NWBHDF5IO, NWBFile, TimeSeries -from pynwb.core import ScratchData + +try: + from ndx_pose import PoseEstimation, PoseEstimationSeries, Skeleton, Skeletons + from pynwb import NWBHDF5IO, NWBFile, TimeSeries + from pynwb.core import ScratchData + + _NWB_AVAILABLE = True +except ImportError: + _NWB_AVAILABLE = False from jabs.core.enums import StorageFormat -from jabs.core.types import PoseData +from jabs.core.types import DynamicObjectData, PoseData from jabs.io.base import Adapter from jabs.io.registry import register_adapter @@ -23,18 +31,39 @@ _JABS_NWB_FORMAT_VERSION = 1 _JABS_METADATA_KEY = "jabs_metadata" _IDENTITY_MASK_KEY = "jabs_identity_mask" -_BOUNDING_BOXES_KEY = "jabs_bounding_boxes" +_BOUNDING_BOXES_PREFIX = "jabs_bounding_boxes" _PROCESSING_MODULE_NAME = "behavior" _PROCESSING_MODULE_DESC = "JABS pose estimation data" _SKELETON_NAME = "subject" _REFERENCE_FRAME = "Top-left corner of video frame, x increases rightward, y increases downward" _CONFIDENCE_DEFINITION = "0.0=invalid/missing keypoint, >0.0=valid keypoint" +_DYNAMIC_CONFIDENCE_DEFINITION = ( + "1.0=valid object instance in this slot, 0.0=slot unoccupied at this prediction" +) + + +def _bounding_box_key(identity_name: str) -> str: + """Return the TimeSeries name for bounding boxes of a given identity.""" + return f"{_BOUNDING_BOXES_PREFIX}_{identity_name}" @register_adapter(StorageFormat.NWB, PoseData, priority=10) class PoseNWBAdapter(Adapter): """NWB adapter for PoseData.""" + def __init__(self) -> None: + """Initialize the adapter, raising ImportError if NWB deps are not installed.""" + self._require_nwb() + + @staticmethod + def _require_nwb() -> None: + """Raise a clear ImportError if pynwb / ndx-pose are not installed.""" + if not _NWB_AVAILABLE: + raise ImportError( + "pynwb and ndx-pose are required for NWB format support. " + "Install with: pip install 'jabs-io[nwb]'" + ) + @classmethod def can_handle(cls, data_type): # noqa: D102 return data_type is PoseData @@ -79,8 +108,8 @@ def write(self, data: PoseData, path: str | Path, **kwargs) -> None: / ← one PoseEstimationSeries per keypoint / ← one PoseEstimation per static object _0/ ← one PoseEstimationSeries per point - jabs_identity_mask ← TimeSeries, uint8 presence mask - jabs_bounding_boxes ← TimeSeries, optional (num_frames, …, 2, 2) + jabs_identity_mask ← TimeSeries, uint8 presence mask + jabs_bounding_boxes_ ← TimeSeries per identity, optional (num_frames, 2, 2) scratch/ jabs_metadata ← JSON: format_version, cm_per_pixel, identity_names, body_parts, metadata, … @@ -171,7 +200,10 @@ def _write_single_file(self, data: PoseData, path: Path, **kwargs) -> None: nwbfile = self._make_nwb_file(**kwargs) skeleton = self._make_skeleton(data.body_parts, data.edges, **kwargs) static_skeletons = self._build_static_skeletons(data.static_objects) - skeletons = Skeletons(skeletons=[skeleton, *static_skeletons.values()]) + dynamic_skeletons = self._build_dynamic_skeletons(data.dynamic_objects) + skeletons = Skeletons( + skeletons=[skeleton, *static_skeletons.values(), *dynamic_skeletons.values()] + ) behavior = nwbfile.create_processing_module( # type: ignore[attr-defined] name=_PROCESSING_MODULE_NAME, @@ -197,6 +229,13 @@ def _write_single_file(self, data: PoseData, path: Path, **kwargs) -> None: ) ) + for obj_name, obj_skeleton in dynamic_skeletons.items(): + behavior.add( + self._build_dynamic_object_pose_estimation( + obj_name, data.dynamic_objects[obj_name], data.fps, obj_skeleton + ) + ) + behavior.add( TimeSeries( name=_IDENTITY_MASK_KEY, @@ -207,16 +246,15 @@ def _write_single_file(self, data: PoseData, path: Path, **kwargs) -> None: ) if data.bounding_boxes is not None: - behavior.add( - TimeSeries( - name=_BOUNDING_BOXES_KEY, - data=np.transpose( - data.bounding_boxes, (1, 0, 2, 3) - ), # (num_frames, num_identities, 2, 2) - unit="pixels", - rate=float(data.fps), + for i, name in enumerate(identity_names): + behavior.add( + TimeSeries( + name=_bounding_box_key(name), + data=data.bounding_boxes[i], # (num_frames, 2, 2) + unit="pixels", + rate=float(data.fps), + ) ) - ) jabs_meta = self._build_jabs_metadata(data, identity_names) nwbfile.add_scratch( @@ -242,10 +280,13 @@ def _write_per_identity(self, data: PoseData, path: Path, **kwargs) -> None: nwbfile = self._make_nwb_file(**kwargs) skeleton = self._make_skeleton(data.body_parts, data.edges, **kwargs) - # Rebuild static skeletons each iteration: HDMF objects can only - # belong to one container, so they cannot be shared across files. + # Rebuild static/dynamic skeletons each iteration: HDMF objects can + # only belong to one container, so they cannot be shared across files. static_skeletons = self._build_static_skeletons(data.static_objects) - skeletons = Skeletons(skeletons=[skeleton, *static_skeletons.values()]) + dynamic_skeletons = self._build_dynamic_skeletons(data.dynamic_objects) + skeletons = Skeletons( + skeletons=[skeleton, *static_skeletons.values(), *dynamic_skeletons.values()] + ) behavior = nwbfile.create_processing_module( # type: ignore[attr-defined] name=_PROCESSING_MODULE_NAME, @@ -270,6 +311,13 @@ def _write_per_identity(self, data: PoseData, path: Path, **kwargs) -> None: ) ) + for obj_name, obj_skeleton in dynamic_skeletons.items(): + behavior.add( + self._build_dynamic_object_pose_estimation( + obj_name, data.dynamic_objects[obj_name], data.fps, obj_skeleton + ) + ) + behavior.add( TimeSeries( name=_IDENTITY_MASK_KEY, @@ -282,8 +330,8 @@ def _write_per_identity(self, data: PoseData, path: Path, **kwargs) -> None: if data.bounding_boxes is not None: behavior.add( TimeSeries( - name=_BOUNDING_BOXES_KEY, - data=data.bounding_boxes[i], + name=_bounding_box_key(identity_name), + data=data.bounding_boxes[i], # (num_frames, 2, 2) unit="pixels", rate=float(data.fps), ) @@ -324,6 +372,8 @@ def _read_single(self, path: Path) -> tuple[PoseData, dict]: # Discover PoseEstimation containers (skip Skeletons, TimeSeries) identity_names = jabs_meta.get("identity_names", []) identity_names_set = set(identity_names) + dynamic_object_names = jabs_meta.get("dynamic_object_names", []) + dynamic_object_shapes = jabs_meta.get("dynamic_object_shapes", {}) pe_containers = { name: obj for name, obj in behavior.data_interfaces.items() @@ -400,33 +450,34 @@ def _read_single(self, path: Path) -> tuple[PoseData, dict]: else: identity_mask = identity_mask.T # (num_identities, num_frames) - # Bounding boxes - # Per-identity files store (num_frames, 2, 2); single files store (num_frames, num_identities, 2, 2). - # Both need to be returned as (num_identities, num_frames, 2, 2). + # Bounding boxes — one TimeSeries per identity, each (num_frames, 2, 2). + # Stack to (num_identities, num_frames, 2, 2). Present only if all identity + # containers exist (bounding boxes are all-or-nothing in PoseData). bounding_boxes = None - if _BOUNDING_BOXES_KEY in behavior.data_interfaces: - bb_ts = behavior[_BOUNDING_BOXES_KEY] - bounding_boxes = np.array(bb_ts.data[:]) - if bounding_boxes.ndim == 3: - bounding_boxes = bounding_boxes[np.newaxis, :] # (1, num_frames, 2, 2) - else: - bounding_boxes = np.transpose( - bounding_boxes, (1, 0, 2, 3) - ) # (num_identities, num_frames, 2, 2) + bb_keys = [_bounding_box_key(name) for name in ordered_names] + if all(k in behavior.data_interfaces for k in bb_keys): + bounding_boxes = np.stack( + [np.array(behavior[k].data[:]) for k in bb_keys], axis=0 + ) # (num_identities, num_frames, 2, 2) # Recover JABS-specific fields cm_per_pixel = jabs_meta.get("cm_per_pixel") - # Per-identity files store the full external_ids list, but this - # intermediate PoseData holds only one identity. Pass None here and - # recover the full list from jabs_meta in _read_merged. - external_ids = ( - None if jabs_meta.get("per_identity_files") else jabs_meta.get("external_ids") - ) + # Per-identity files store the full external_ids and subjects dicts, but + # this intermediate PoseData holds only one identity. Pass None here and + # recover the full values from jabs_meta in _read_merged. + per_id = jabs_meta.get("per_identity_files", False) + external_ids = None if per_id else jabs_meta.get("external_ids") + subjects = None if per_id else jabs_meta.get("subjects") metadata = jabs_meta.get("metadata", {}) + # Read dynamic objects from NWB-native PoseEstimation containers. + dynamic_objects = self._read_dynamic_objects( + pe_containers, dynamic_object_names, dynamic_object_shapes, int(fps_value) + ) + # Read static objects from NWB-native PoseEstimation containers. - # Any PoseEstimation not in identity_names_set is a static object. - static_objects = self._read_static_objects(pe_containers, identity_names_set) + static_object_names = jabs_meta.get("static_object_names", []) + static_objects = self._read_static_objects(pe_containers, static_object_names) pose_data = PoseData( points=points, @@ -438,7 +489,9 @@ def _read_single(self, path: Path) -> tuple[PoseData, dict]: cm_per_pixel=cm_per_pixel, bounding_boxes=bounding_boxes, static_objects=static_objects, + dynamic_objects=dynamic_objects, external_ids=external_ids, + subjects=subjects, metadata=metadata, ) return pose_data, jabs_meta @@ -493,9 +546,11 @@ def _read_merged(self, path: Path, jabs_meta: dict) -> PoseData: if all(pd.bounding_boxes is not None for pd in pose_datas): bounding_boxes = np.concatenate([pd.bounding_boxes for pd in pose_datas], axis=0) - # Recover external_ids from jabs_meta of the first file; each per-identity - # file stores the full original list, so any file's meta will do. - external_ids = parts[0][2].get("external_ids") + # Recover external_ids and subjects from jabs_meta of the first file; + # each per-identity file stores the full original values, so any file's meta will do. + first_meta = parts[0][2] + external_ids = first_meta.get("external_ids") + subjects = first_meta.get("subjects") return PoseData( points=points, @@ -507,7 +562,9 @@ def _read_merged(self, path: Path, jabs_meta: dict) -> PoseData: cm_per_pixel=ref.cm_per_pixel, bounding_boxes=bounding_boxes, static_objects=ref.static_objects, + dynamic_objects=ref.dynamic_objects, external_ids=external_ids, + subjects=subjects, metadata=ref.metadata, ) @@ -638,34 +695,156 @@ def _build_static_object_pose_estimation( source_software="JABS", ) + @staticmethod + def _build_dynamic_skeletons( + dynamic_objects: dict[str, DynamicObjectData], + ) -> dict[str, Skeleton]: # type: ignore[valid-type] + """Build a Skeleton for each dynamic object. + + For single-keypoint objects, nodes are named ``{name}_{slot}``. + For multi-keypoint objects, nodes are named ``{name}_{slot}_{kp}``. + + Args: + dynamic_objects: Mapping of object name to DynamicObjectData. + + Returns: + Ordered dict mapping object name to its Skeleton. + """ + skeletons: dict[str, Skeleton] = {} # type: ignore[valid-type] + for name, dyn_obj in dynamic_objects.items(): + max_count = dyn_obj.points.shape[1] + n_keypoints = dyn_obj.points.shape[2] + if n_keypoints == 1: + nodes = [f"{name}_{slot}" for slot in range(max_count)] + else: + nodes = [ + f"{name}_{slot}_{kp}" for slot in range(max_count) for kp in range(n_keypoints) + ] + skeletons[name] = Skeleton(name=name, nodes=nodes) + return skeletons + + @staticmethod + def _build_dynamic_object_pose_estimation( + name: str, + dyn_obj: DynamicObjectData, + fps: int, + skeleton: Skeleton, # type: ignore[valid-type] + ) -> PoseEstimation: + """Build a PoseEstimation container for one dynamic object. + + Each slot/keypoint combination is stored as a PoseEstimationSeries with + irregular timestamps corresponding to the sample indices. Confidence of + 1.0 indicates the slot is occupied at that prediction; 0.0 means empty. + + Args: + name: Name of the dynamic object (e.g. "fecal_boli"). + dyn_obj: DynamicObjectData with shape (n_predictions, max_count, n_keypoints, 2). + fps: Frames per second used to convert sample indices to timestamps. + skeleton: Skeleton built by :meth:`_build_dynamic_skeletons`. + + Returns: + PoseEstimation container ready to add to the behavior module. + """ + n_predictions, max_count, n_keypoints, _ = dyn_obj.points.shape + timestamps = (dyn_obj.sample_indices / fps).tolist() + series_list = [] + for slot in range(max_count): + slot_confidence = (dyn_obj.counts > slot).astype(np.float64) + for kp in range(n_keypoints): + series_name = f"{name}_{slot}" if n_keypoints == 1 else f"{name}_{slot}_{kp}" + series_list.append( + PoseEstimationSeries( + name=series_name, + data=dyn_obj.points[:, slot, kp, :].astype(np.float64), + confidence=slot_confidence, + confidence_definition=_DYNAMIC_CONFIDENCE_DEFINITION, + timestamps=timestamps, + unit="pixels", + reference_frame=_REFERENCE_FRAME, + ) + ) + return PoseEstimation( + name=name, + pose_estimation_series=series_list, + description=f"Dynamic object: {name}", + skeleton=skeleton, + source_software="JABS", + ) + + @staticmethod + def _read_dynamic_objects( + pe_containers: dict[str, PoseEstimation], + dynamic_object_names: list[str], + dynamic_object_shapes: dict[str, list[int]], + fps: int, + ) -> dict[str, DynamicObjectData]: + """Reconstruct dynamic_objects from PoseEstimation containers. + + Args: + pe_containers: All PoseEstimation containers from the behavior module. + dynamic_object_names: Names of dynamic objects recorded in jabs_metadata. + dynamic_object_shapes: Mapping of name to [max_count, n_keypoints]. + fps: Frames per second used to convert timestamps back to sample indices. + + Returns: + Mapping of object name to DynamicObjectData. + """ + dynamic_objects: dict[str, DynamicObjectData] = {} + for name in dynamic_object_names: + if name not in pe_containers: + logger.warning("Dynamic object %r not found in behavior module", name) + continue + pe = pe_containers[name] + max_count, n_keypoints = dynamic_object_shapes[name] + first_series = next(iter(pe.pose_estimation_series.values())) + timestamps = np.array(first_series.timestamps[:]) + sample_indices = np.round(timestamps * fps).astype(np.int64) + n_predictions = len(sample_indices) + points = np.empty((n_predictions, max_count, n_keypoints, 2), dtype=np.float64) + counts = np.zeros(n_predictions, dtype=np.int64) + for slot in range(max_count): + slot_confidence = None + for kp in range(n_keypoints): + series_name = f"{name}_{slot}" if n_keypoints == 1 else f"{name}_{slot}_{kp}" + series = pe.pose_estimation_series[series_name] + points[:, slot, kp, :] = np.array(series.data[:]) + if slot_confidence is None: + slot_confidence = np.array(series.confidence[:]) + if slot_confidence is not None: + counts += (slot_confidence > 0).astype(np.int64) + dynamic_objects[name] = DynamicObjectData( + points=points, counts=counts, sample_indices=sample_indices + ) + return dynamic_objects + @staticmethod def _read_static_objects( pe_containers: dict[str, PoseEstimation], - identity_names_set: set[str], + static_names: list[str], ) -> dict[str, npt.NDArray[np.float64]]: - """Reconstruct static_objects from non-identity PoseEstimation containers. + """Reconstruct static_objects from PoseEstimation containers. Args: pe_containers: All PoseEstimation containers from the behavior module. - identity_names_set: Set of container names that are animal identities. + static_names: Ordered list of static object container names from + ``jabs_metadata``. Returns: Mapping of static object name to (N, 2) float64 array. """ static_objects: dict[str, npt.NDArray[np.float64]] = {} - for name, pe in pe_containers.items(): - if name in identity_names_set: + for name in static_names: + if name not in pe_containers: + logger.warning("Static object %r listed in metadata but not found in file", name) continue - # Sort series by numeric suffix to restore original point order. series_items = sorted( - pe.pose_estimation_series.items(), + pe_containers[name].pose_estimation_series.items(), key=lambda kv: int(kv[0].rsplit("_", 1)[-1]), ) - pts = np.array( + static_objects[name] = np.array( [np.array(series.data[0]) for _, series in series_items], dtype=np.float64, ) # shape (N, 2) - static_objects[name] = pts return static_objects @staticmethod @@ -681,8 +860,17 @@ def _build_jabs_metadata( "identity_names": identity_names, "num_identities": data.points.shape[0], "body_parts": data.body_parts, + "subjects": data.subjects, "metadata": data.metadata, } + if data.static_objects: + meta["static_object_names"] = list(data.static_objects.keys()) + if data.dynamic_objects: + meta["dynamic_object_names"] = list(data.dynamic_objects.keys()) + meta["dynamic_object_shapes"] = { + name: [dyn.points.shape[1], dyn.points.shape[2]] + for name, dyn in data.dynamic_objects.items() + } meta.update(extra) return meta diff --git a/packages/jabs-io/tests/internal/pose/test_nwb.py b/packages/jabs-io/tests/internal/pose/test_nwb.py index 1176e353..bb176001 100644 --- a/packages/jabs-io/tests/internal/pose/test_nwb.py +++ b/packages/jabs-io/tests/internal/pose/test_nwb.py @@ -4,7 +4,7 @@ import pytest from ndx_pose import PoseEstimation -from jabs.core.types import PoseData +from jabs.core.types import DynamicObjectData, PoseData from jabs.io.internal.pose.nwb import PoseNWBAdapter @@ -24,6 +24,8 @@ def _make_pose_data( with_bounding_boxes=False, with_static_objects=True, with_metadata=True, + with_subjects=False, + with_dynamic_objects=False, edges=None, ): rng = np.random.default_rng(42) @@ -40,6 +42,24 @@ def _make_pose_data( if edges is None: edges = [(0, 1), (1, 2)] + names = external_ids or [f"subject_{i}" for i in range(num_identities)] + subjects = {name: {"sex": "M", "genotype": "WT"} for name in names} if with_subjects else None + + if with_dynamic_objects: + dyn_rng = np.random.default_rng(42) + points_dyn = dyn_rng.random((5, 2, 1, 2)) * 100 + counts_dyn = np.array([2, 1, 0, 2, 1], dtype=np.int64) + sample_indices_dyn = np.array([10, 25, 40, 60, 80], dtype=np.int64) + dynamic_objects: dict[str, DynamicObjectData] = { + "fecal_boli": DynamicObjectData( + points=points_dyn, + counts=counts_dyn, + sample_indices=sample_indices_dyn, + ) + } + else: + dynamic_objects = {} + return PoseData( points=points, point_mask=point_mask, @@ -50,7 +70,9 @@ def _make_pose_data( cm_per_pixel=cm_per_pixel, bounding_boxes=bounding_boxes, static_objects=static_objects, + dynamic_objects=dynamic_objects, external_ids=external_ids, + subjects=subjects, metadata=metadata, ) @@ -65,10 +87,20 @@ def _assert_pose_data_equal(a: PoseData, b: PoseData): assert a.fps == b.fps assert a.cm_per_pixel == b.cm_per_pixel assert a.external_ids == b.external_ids + assert a.subjects == b.subjects assert a.metadata == b.metadata assert a.static_objects.keys() == b.static_objects.keys() for key in a.static_objects: np.testing.assert_allclose(a.static_objects[key], b.static_objects[key]) + assert set(a.dynamic_objects.keys()) == set(b.dynamic_objects.keys()) + for key in a.dynamic_objects: + np.testing.assert_allclose( + a.dynamic_objects[key].points, b.dynamic_objects[key].points, atol=1e-10 + ) + np.testing.assert_array_equal(a.dynamic_objects[key].counts, b.dynamic_objects[key].counts) + np.testing.assert_array_equal( + a.dynamic_objects[key].sample_indices, b.dynamic_objects[key].sample_indices + ) if a.bounding_boxes is None: assert b.bounding_boxes is None else: @@ -91,6 +123,17 @@ def test_roundtrip_single_file(tmp_path, adapter): _assert_pose_data_equal(data, loaded) +def test_roundtrip_with_subjects(tmp_path, adapter): + """subjects metadata survives roundtrip and is caught by _assert_pose_data_equal.""" + path = tmp_path / "pose_subjects.nwb" + data = _make_pose_data(external_ids=["mouse_a", "mouse_b"], with_subjects=True) + + adapter.write(data, path) + loaded = adapter.read(path) + + _assert_pose_data_equal(data, loaded) + + def test_roundtrip_with_bounding_boxes(tmp_path, adapter): """Bounding boxes survive roundtrip.""" path = tmp_path / "pose_bb.nwb" @@ -250,13 +293,13 @@ def test_static_objects_nwb_structure(tmp_path, adapter): assert "lixit_0" in lixit_pe.pose_estimation_series -def test_static_objects_not_in_jabs_metadata_json(tmp_path, adapter): - """Static objects are no longer stored in the JSON scratch blob.""" +def test_static_object_names_in_jabs_metadata_json(tmp_path, adapter): + """static_object_names is written to jabs_metadata when static objects are present.""" import json from pynwb import NWBHDF5IO - path = tmp_path / "pose_no_json_static.nwb" + path = tmp_path / "pose_json_static.nwb" data = _make_pose_data(with_static_objects=True) adapter.write(data, path) @@ -264,7 +307,8 @@ def test_static_objects_not_in_jabs_metadata_json(tmp_path, adapter): with NWBHDF5IO(str(path), "r") as io: nwb = io.read() jabs_meta = json.loads(str(nwb.scratch["jabs_metadata"].data)) - assert "static_objects" not in jabs_meta + assert "static_object_names" in jabs_meta + assert jabs_meta["static_object_names"] == ["lixit"] def test_empty_static_objects(tmp_path, adapter): @@ -306,6 +350,98 @@ def test_static_objects_1d_skipped_with_warning(tmp_path, adapter, caplog): assert "bad_obj" not in loaded.static_objects +def test_subjects_roundtrip(tmp_path, adapter): + """Per-animal subjects metadata survives a single-file roundtrip.""" + import json + + from pynwb import NWBHDF5IO + + path = tmp_path / "pose_subjects.nwb" + subjects = { + "mouse_a": {"sex": "M", "genotype": "WT", "strain": "C57BL/6", "age": "P60"}, + "mouse_b": {"sex": "F", "genotype": "KO", "strain": "C57BL/6", "age": "P60"}, + } + data = _make_pose_data(external_ids=["mouse_a", "mouse_b"]) + data = data.__class__( + **{ + **{f: getattr(data, f) for f in data.__dataclass_fields__}, + "subjects": subjects, + } + ) + + adapter.write(data, path) + loaded = adapter.read(path) + + assert loaded.subjects == subjects + + # Verify it's in the jabs_metadata JSON + with NWBHDF5IO(str(path), "r") as io: + nwb = io.read() + meta = json.loads(str(nwb.scratch["jabs_metadata"].data)) + assert meta["subjects"] == subjects + + +def test_subjects_none_roundtrip(tmp_path, adapter): + """subjects=None reads back as None.""" + path = tmp_path / "pose_no_subjects.nwb" + data = _make_pose_data() + + adapter.write(data, path) + loaded = adapter.read(path) + + assert loaded.subjects is None + + +def test_subjects_per_identity_roundtrip(tmp_path, adapter): + """Per-animal subjects metadata survives a per-identity file roundtrip.""" + path = tmp_path / "pose.nwb" + subjects = { + "mouse_a": {"sex": "M", "genotype": "WT"}, + "mouse_b": {"sex": "F", "genotype": "KO"}, + } + data = _make_pose_data(external_ids=["mouse_a", "mouse_b"]) + data = data.__class__( + **{ + **{f: getattr(data, f) for f in data.__dataclass_fields__}, + "subjects": subjects, + } + ) + + adapter.write(data, path, per_identity_files=True) + loaded = adapter.read(tmp_path / "pose_mouse_a.nwb") + + assert loaded.subjects == subjects + + +def test_bounding_boxes_per_identity_containers(tmp_path, adapter): + """Bounding boxes are stored as one TimeSeries per identity, not a single combined array.""" + from pynwb import NWBHDF5IO + + path = tmp_path / "pose_bb_struct.nwb" + data = _make_pose_data( + num_identities=2, + num_frames=10, + external_ids=["mouse_a", "mouse_b"], + with_bounding_boxes=True, + ) + + adapter.write(data, path) + + with NWBHDF5IO(str(path), "r") as io: + nwb = io.read() + behavior = nwb.processing["behavior"] + + assert "jabs_bounding_boxes_mouse_a" in behavior.data_interfaces + assert "jabs_bounding_boxes_mouse_b" in behavior.data_interfaces + # The old combined key must not be present + assert "jabs_bounding_boxes" not in behavior.data_interfaces + + bb_a = behavior["jabs_bounding_boxes_mouse_a"] + bb_b = behavior["jabs_bounding_boxes_mouse_b"] + assert np.array(bb_a.data[:]).shape == (10, 2, 2) + assert np.array(bb_b.data[:]).shape == (10, 2, 2) + + def test_edges_roundtrip(tmp_path, adapter): """Skeleton edges survive roundtrip.""" path = tmp_path / "pose_edges.nwb" @@ -485,3 +621,121 @@ def test_write_raises_on_collision_after_sanitization(tmp_path, adapter): with pytest.raises(ValueError, match="not unique after sanitization"): adapter.write(data, path) + + +# --------------------------------------------------------------------------- +# Dynamic objects +# --------------------------------------------------------------------------- + + +def test_dynamic_objects_roundtrip(tmp_path, adapter): + """Dynamic objects survive a single-file roundtrip.""" + path = tmp_path / "pose_dyn.nwb" + data = _make_pose_data(with_dynamic_objects=True) + + adapter.write(data, path) + loaded = adapter.read(path) + + _assert_pose_data_equal(data, loaded) + + +def test_dynamic_objects_per_identity_roundtrip(tmp_path, adapter): + """Dynamic objects survive a per-identity file roundtrip.""" + path = tmp_path / "pose.nwb" + data = _make_pose_data(external_ids=["mouse_a", "mouse_b"], with_dynamic_objects=True) + + adapter.write(data, path, per_identity_files=True) + loaded = adapter.read(tmp_path / "pose_mouse_a.nwb") + + _assert_pose_data_equal(data, loaded) + + +def test_dynamic_objects_nwb_structure(tmp_path, adapter): + """Dynamic objects are stored as PoseEstimation + Skeleton in the behavior module.""" + from pynwb import NWBHDF5IO + + path = tmp_path / "pose_dyn_struct.nwb" + data = _make_pose_data(with_dynamic_objects=True) + + adapter.write(data, path) + + with NWBHDF5IO(str(path), "r") as io: + nwb = io.read() + behavior = nwb.processing["behavior"] + + # fecal_boli should be a PoseEstimation container + assert "fecal_boli" in behavior.data_interfaces + assert isinstance(behavior.data_interfaces["fecal_boli"], PoseEstimation) + + # The Skeletons container should include a fecal_boli skeleton + skeletons_obj = behavior.data_interfaces["Skeletons"] + assert "fecal_boli" in skeletons_obj.skeletons + + # Single-keypoint, 2 slots → series named fecal_boli_0 and fecal_boli_1 + fb_pe = behavior.data_interfaces["fecal_boli"] + assert "fecal_boli_0" in fb_pe.pose_estimation_series + assert "fecal_boli_1" in fb_pe.pose_estimation_series + + +def test_dynamic_objects_empty(tmp_path, adapter): + """Empty dynamic_objects writes no extra containers and reads back as empty dict.""" + import json + + from pynwb import NWBHDF5IO + + path = tmp_path / "pose_no_dyn.nwb" + data = _make_pose_data(with_dynamic_objects=False) + + adapter.write(data, path) + loaded = adapter.read(path) + + assert loaded.dynamic_objects == {} + + with NWBHDF5IO(str(path), "r") as io: + nwb = io.read() + meta = json.loads(str(nwb.scratch["jabs_metadata"].data)) + assert "dynamic_object_names" not in meta + assert "dynamic_object_shapes" not in meta + + +def test_dynamic_objects_multi_keypoint(tmp_path, adapter): + """Multi-keypoint dynamic objects use {name}_{slot}_{kp} node naming and roundtrip.""" + path = tmp_path / "pose_dyn_multi.nwb" + rng = np.random.default_rng(7) + # foo: 2 instances, 2 keypoints each → shape (4, 2, 2, 2) + points_foo = rng.random((4, 2, 2, 2)) * 200 + counts_foo = np.array([2, 1, 2, 0], dtype=np.int64) + sample_indices_foo = np.array([5, 15, 30, 50], dtype=np.int64) + dyn_obj = DynamicObjectData( + points=points_foo, counts=counts_foo, sample_indices=sample_indices_foo + ) + base = _make_pose_data(with_dynamic_objects=False) + data = PoseData( + points=base.points, + point_mask=base.point_mask, + identity_mask=base.identity_mask, + body_parts=base.body_parts, + edges=base.edges, + fps=base.fps, + cm_per_pixel=base.cm_per_pixel, + dynamic_objects={"foo": dyn_obj}, + metadata=base.metadata, + ) + + adapter.write(data, path) + loaded = adapter.read(path) + + assert "foo" in loaded.dynamic_objects + foo = loaded.dynamic_objects["foo"] + np.testing.assert_allclose(foo.points, points_foo, atol=1e-10) + np.testing.assert_array_equal(foo.counts, counts_foo) + np.testing.assert_array_equal(foo.sample_indices, sample_indices_foo) + + from pynwb import NWBHDF5IO + + with NWBHDF5IO(str(path), "r") as io: + nwb = io.read() + foo_pe = nwb.processing["behavior"]["foo"] + # 2 slots * 2 keypoints -> 4 series + expected_names = {"foo_0_0", "foo_0_1", "foo_1_0", "foo_1_1"} + assert set(foo_pe.pose_estimation_series.keys()) == expected_names diff --git a/pyproject.toml b/pyproject.toml index 6c36708d..ffcd5875 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,6 +10,7 @@ authors = [ { name = "Glen Beane" }, { name = "Brian Geuther" }, { name = "Keith Sheppard" }, + { name = "Alexander Berger-Liedtka" }, ] dependencies = [ diff --git a/ruff.toml b/ruff.toml index f7a355f6..f9c1daa0 100644 --- a/ruff.toml +++ b/ruff.toml @@ -32,3 +32,7 @@ convention = "google" [lint.per-file-ignores] "__init__.py" = ["F401"] # Unused imports in __init__ files +"src/jabs/scripts/cli/cli.py" = [ + "D301", # Click uses \b in docstrings to control help formatting + "D412", # Click Examples sections look better with a blank line before \b blocks +] diff --git a/src/jabs/pose_estimation/pose_est_v7.py b/src/jabs/pose_estimation/pose_est_v7.py index 584714ab..2227d55b 100644 --- a/src/jabs/pose_estimation/pose_est_v7.py +++ b/src/jabs/pose_estimation/pose_est_v7.py @@ -1,12 +1,117 @@ +"""Pose estimation handler for version 7 pose files with dynamic object support.""" + +import logging +from pathlib import Path + +import h5py +import numpy as np + +from jabs.core.types import DynamicObjectData + from .pose_est_v6 import PoseEstimationV6 +logger = logging.getLogger(__name__) + +# Attribute name on dynamic_objects/[name]/points datasets that specifies +# the axis ordering of the stored coordinates. Valid values: "xy" or "yx". +# When this attribute is absent the reader defaults to "yx" (matching the +# fecal_boli network which was trained with HRNet and stores points as y, x). +_AXIS_ORDER_ATTR = "axis_order" +_DEFAULT_AXIS_ORDER = "yx" + class PoseEstimationV7(PoseEstimationV6): - """Pose estimation version 7 + """Pose estimation handler for version 7 pose files with dynamic object support. + + Extends PoseEstimationV6 to add reading and management of dynamic object data + (e.g. fecal boli positions) from pose v7 HDF5 files. Dynamic objects differ + from static objects in that their position or count may change over time, and + predictions are made only at a subset of frames recorded in sample_indices. + + The dynamic_objects/[name]/points dataset may carry an axis_order + attribute specifying whether coordinates are stored as "xy" or "yx". + If the attribute is absent the reader defaults to "yx" and flips the last + axis so that all data is returned in (x, y) order. - Currently handled the same as v6 because we're not using the v7 dynamic_objects dataset yet. + Args: + file_path: Path to the pose HDF5 file. + cache_dir: Optional cache directory for intermediate data. + fps: Frames per second for the video. """ + def __init__(self, file_path: Path, cache_dir: Path | None = None, fps: int = 30) -> None: + """Initialize PoseEstimationV7 from an HDF5 pose file. + + Args: + file_path: Path to the pose v7 HDF5 file. + cache_dir: Optional cache directory for intermediate data. + fps: Frames per second for the video. + """ + super().__init__(file_path, cache_dir, fps) + + with h5py.File(self._path, "r") as pose_h5: + if "dynamic_objects" not in pose_h5: + return + + for obj_name in pose_h5["dynamic_objects"]: + obj_group = pose_h5[f"dynamic_objects/{obj_name}"] + + required = {"points", "counts", "sample_indices"} + if not required.issubset(obj_group.keys()): + logger.warning( + "Dynamic object %r is missing required datasets %r; skipping.", + obj_name, + required - set(obj_group.keys()), + ) + continue + + points_ds = obj_group["points"] + axis_order: str = points_ds.attrs.get(_AXIS_ORDER_ATTR, _DEFAULT_AXIS_ORDER) + + points = points_ds[:].astype(np.float64) + counts = obj_group["counts"][:].astype(np.int64) + sample_indices = obj_group["sample_indices"][:].astype(np.int64) + + if points.ndim not in (3, 4): + logger.warning( + "Dynamic object %r has unexpected points shape %s; skipping.", + obj_name, + points.shape, + ) + continue + + if axis_order == "yx": + points = np.flip(points, axis=-1).copy() + elif axis_order == "xy": + pass + else: + logger.warning( + "Dynamic object %r has unknown axis_order=%r; defaulting to 'yx'.", + obj_name, + axis_order, + ) + points = np.flip(points, axis=-1).copy() + + # Normalize to 4-D (n_predictions, max_count, n_keypoints, 2). + # Some older single-keypoint objects (e.g. fecal_boli) are stored as 3-D + # in the HDF5 file; expand to (n_predictions, max_count, 1, 2). + if points.ndim == 3: + points = points[:, :, np.newaxis, :] + + self._dynamic_objects[obj_name] = DynamicObjectData( + points=points, + counts=counts, + sample_indices=sample_indices, + ) + logger.debug( + "Loaded dynamic object %r: %d predictions, max_count=%d, n_keypoints=%d, axis_order=%r", + obj_name, + len(sample_indices), + points.shape[1], + points.shape[2], + axis_order, + ) + @property def format_major_version(self) -> int: """Returns the major version of the pose file format.""" diff --git a/src/jabs/scripts/cli/cli.py b/src/jabs/scripts/cli/cli.py index ec0bc1fa..80c8b601 100644 --- a/src/jabs/scripts/cli/cli.py +++ b/src/jabs/scripts/cli/cli.py @@ -7,6 +7,7 @@ functionality from existing JABS scripts such as `jabs-init` or `jabs-merge`). """ +import json from pathlib import Path import click @@ -16,6 +17,7 @@ from jabs.core.enums import ClassifierType, CrossValidationGroupingStrategy from jabs.project import Project, export_training_data, get_videos_to_prune +from .convert_to_nwb import run_conversion from .cross_validation import run_cross_validation # find out which classifiers are supported in this environment @@ -328,6 +330,114 @@ def cross_validation( raise click.ClickException(str(e)) from e +@cli.command(name="convert-to-nwb") +@click.argument( + "input_path", + type=click.Path(exists=True, dir_okay=False, path_type=Path), +) +@click.argument( + "output", + type=click.Path(dir_okay=False, writable=True, path_type=Path), +) +@click.option( + "--per-identity", + is_flag=True, + default=False, + help=( + "Write one NWB file per identity instead of a single combined file. " + "OUTPUT is used as a naming template; files are written as " + "{output_stem}_{identity_name}.nwb alongside it." + ), +) +@click.option( + "--session-description", + type=str, + default=None, + help="NWB session description string. Defaults to 'JABS PoseEstimation Data'.", +) +@click.option( + "--subjects", + "subjects_path", + type=click.Path(exists=True, dir_okay=False, path_type=Path), + default=None, + help=( + "Path to a JSON file containing per-animal biological metadata. " + "Keys are identity names: use the external IDs from the pose file " + "if present (e.g. 'mouse_a'), or 'subject_0', 'subject_1', … if the " + "pose file has no external IDs. " + "Standard fields: subject_id, sex, genotype, strain, age, weight, " + "species, description." + ), +) +@click.pass_context +def convert_to_nwb( + ctx: click.Context, + input_path: Path, + output: Path, + per_identity: bool, + session_description: str | None, + subjects_path: Path | None, +) -> None: + """Convert a JABS pose estimation file to NWB format. + + INPUT_PATH is a JABS pose HDF5 file (any version, v2-v8). The format + version is inferred automatically from the filename (e.g. _pose_est_v6.h5). + + OUTPUT is the destination NWB file. In --per-identity mode, OUTPUT is a + naming template and is not created directly; instead one file per identity + is written as {output_stem}_{identity_name}.nwb in the same directory. + + Examples: + + \b + # Single file, all identities + jabs-cli convert-to-nwb session_pose_est_v6.h5 session.nwb + + \b + # One NWB file per identity + jabs-cli convert-to-nwb session_pose_est_v6.h5 session.nwb --per-identity + + \b + # Include per-animal metadata + jabs-cli convert-to-nwb session_pose_est_v6.h5 session.nwb --subjects subjects.json + """ + if ctx.obj["VERBOSE"]: + click.echo(f"Input: {input_path}") + click.echo(f"Output: {output}") + click.echo(f"Per-identity: {per_identity}") + if subjects_path: + click.echo(f"Subjects: {subjects_path}") + + subjects: dict[str, dict] | None = None + if subjects_path is not None: + try: + subjects = json.loads(subjects_path.read_text()) + except Exception as e: + raise click.ClickException(f"Failed to read subjects file: {e}") from e + if not isinstance(subjects, dict): + raise click.ClickException( + f"Subjects file must contain a JSON object, got {type(subjects).__name__}" + ) + + console = Console() + with console.status(f"Converting {input_path.name} → NWB ...", spinner="dots"): + try: + run_conversion( + input_path=input_path, + output_path=output, + per_identity=per_identity, + session_description=session_description, + subjects=subjects, + ) + except Exception as e: + raise click.ClickException(str(e)) from e + + if per_identity: + click.echo(f"Wrote per-identity NWB files to {output.parent}") + else: + click.echo(f"Wrote {output}") + + def main(): """Entry point for the JABS CLI.""" cli(obj={}) diff --git a/src/jabs/scripts/cli/convert_to_nwb.py b/src/jabs/scripts/cli/convert_to_nwb.py new file mode 100644 index 00000000..4552b11b --- /dev/null +++ b/src/jabs/scripts/cli/convert_to_nwb.py @@ -0,0 +1,143 @@ +"""Convert a JABS pose estimation file (any version) to NWB format.""" + +import logging +from pathlib import Path + +import numpy as np + +from jabs.core.abstract.pose_est import PoseEstimation +from jabs.core.types.pose import PoseData +from jabs.io import save +from jabs.pose_estimation import open_pose_file + +logger = logging.getLogger(__name__) + + +def _segments_to_edges(segments) -> list[tuple[int, int]]: + """Convert connected segment paths to a list of (src, dst) edge pairs. + + A segment like (0, 3, 6, 9) produces edges (0,3), (3,6), (6,9). + + Args: + segments: Iterable of sequences of keypoint indices. + + Returns: + List of (src, dst) index tuples. + """ + edges = [] + for segment in segments: + for i in range(len(segment) - 1): + edges.append((int(segment[i]), int(segment[i + 1]))) + return edges + + +def pose_to_pose_data( + pose: PoseEstimation, + subjects: dict[str, dict] | None = None, +) -> PoseData: + """Convert any PoseEstimation object to a PoseData dataclass. + + Handles all supported JABS pose versions (v2-v8). + + Args: + pose: A loaded PoseEstimation object (any version). + subjects: Optional per-animal biological metadata, keyed by identity + name (matching external_identities values). Passed through + directly to PoseData.subjects. + + Returns: + A PoseData instance ready for NWB export. + """ + all_points = [] + all_point_masks = [] + for identity in pose.identities: + points, mask = pose.get_identity_poses(identity) + all_points.append(points) + all_point_masks.append(mask) + + points_array = np.stack(all_points, axis=0) + point_mask_array = np.stack(all_point_masks, axis=0) + + identity_mask_array = np.stack( + [pose.identity_mask(identity) for identity in pose.identities], + axis=0, + ) + + body_parts = [kpt.name for kpt in PoseEstimation.KeypointIndex] + edges = _segments_to_edges(pose.get_connected_segments()) + + cm_per_pixel = getattr(pose, "cm_per_pixel", None) + static_objects = getattr(pose, "static_objects", {}) + external_ids = getattr(pose, "external_identities", None) + + per_identity_boxes = [pose.get_bounding_boxes(i) for i in pose.identities] + bounding_boxes: np.ndarray | None = None + if all(b is not None for b in per_identity_boxes): + bounding_boxes = np.stack(per_identity_boxes, axis=0) # (num_identities, num_frames, 2, 2) + + file_hash = getattr(pose, "hash", None) + metadata: dict = { + "source_file": str(pose.pose_file), + "pose_format_version": pose.format_major_version, + } + if file_hash is not None: + metadata["source_file_hash"] = file_hash + + return PoseData( + points=points_array, + point_mask=point_mask_array, + identity_mask=identity_mask_array, + body_parts=body_parts, + edges=edges, + fps=pose.fps, + cm_per_pixel=cm_per_pixel, + bounding_boxes=bounding_boxes, + static_objects=static_objects, + external_ids=external_ids, + subjects=subjects, + metadata=metadata, + ) + + +def run_conversion( + input_path: Path, + output_path: Path, + per_identity: bool = False, + session_description: str | None = None, + subjects: dict[str, dict] | None = None, +) -> None: + """Convert a JABS pose HDF5 file to NWB and write to disk. + + The pose format version is inferred from the filename (e.g. + "_pose_est_v6.h5" → v6). Supported versions: v2-v8. + + Args: + input_path: Path to the input JABS pose HDF5 file. + output_path: Destination path for the NWB file. In per-identity mode + this is used as a naming template; actual files are written + alongside it as "{stem}_{identity_name}.nwb". + per_identity: If True, write one NWB file per identity. + session_description: Optional NWB session description string. + subjects: Optional per-animal biological metadata dict, keyed by + identity name. See PoseData.subjects for the expected + structure. + + Raises: + ValueError: If the input file is not a recognized JABS pose file. + FileNotFoundError: If the input file does not exist. + """ + logger.info("Loading %s", input_path) + pose = open_pose_file(input_path) + identity_word = "identity" if pose.num_identities == 1 else "identities" + logger.info( + "%d %s, %d frames, %d fps", pose.num_identities, identity_word, pose.num_frames, pose.fps + ) + + pose_data = pose_to_pose_data(pose, subjects=subjects) + + write_kwargs: dict = {"per_identity_files": per_identity} + if session_description is not None: + write_kwargs["session_description"] = session_description + + logger.info("Writing NWB to %s", output_path) + save(pose_data, output_path, **write_kwargs) diff --git a/src/jabs/ui/player_widget/overlays/bounding_box.py b/src/jabs/ui/player_widget/overlays/bounding_box.py index 1d67edaf..61a34bb7 100644 --- a/src/jabs/ui/player_widget/overlays/bounding_box.py +++ b/src/jabs/ui/player_widget/overlays/bounding_box.py @@ -66,6 +66,8 @@ def paint(self, painter: QtGui.QPainter, crop_rect: QtCore.QRect) -> None: # we use methods available only in PoseEstimationV8 and later # cast the base PoseEstimation to PoseEstimationV8 to avoid IDE warnings about undefined methods pose = cast(PoseEstimationV8, self.parent.pose) + + # no need to proceed if there are no bounding boxes to draw if not pose.has_bounding_boxes: return diff --git a/uv.lock b/uv.lock index 5ac4d04a..9ce7cf73 100644 --- a/uv.lock +++ b/uv.lock @@ -1033,7 +1033,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/38/3f/9859f655d11901e7b2996c6e3d33e0caa9a1d4572c3bc61ed0faa64b2f4c/greenlet-3.3.2-cp310-cp310-macosx_11_0_universal2.whl", hash = "sha256:9bc885b89709d901859cf95179ec9f6bb67a3d2bb1f0e88456461bd4b7f8fd0d", size = 277747, upload-time = "2026-02-20T20:16:21.325Z" }, { url = "https://files.pythonhosted.org/packages/fb/07/cb284a8b5c6498dbd7cba35d31380bb123d7dceaa7907f606c8ff5993cbf/greenlet-3.3.2-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b568183cf65b94919be4438dc28416b234b678c608cafac8874dfeeb2a9bbe13", size = 579202, upload-time = "2026-02-20T20:47:28.955Z" }, { url = "https://files.pythonhosted.org/packages/ed/45/67922992b3a152f726163b19f890a85129a992f39607a2a53155de3448b8/greenlet-3.3.2-cp310-cp310-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:527fec58dc9f90efd594b9b700662ed3fb2493c2122067ac9c740d98080a620e", size = 590620, upload-time = "2026-02-20T20:55:55.581Z" }, - { url = "https://files.pythonhosted.org/packages/03/5f/6e2a7d80c353587751ef3d44bb947f0565ec008a2e0927821c007e96d3a7/greenlet-3.3.2-cp310-cp310-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:508c7f01f1791fbc8e011bd508f6794cb95397fdb198a46cb6635eb5b78d85a7", size = 602132, upload-time = "2026-02-20T21:02:43.261Z" }, { url = "https://files.pythonhosted.org/packages/ad/55/9f1ebb5a825215fadcc0f7d5073f6e79e3007e3282b14b22d6aba7ca6cb8/greenlet-3.3.2-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ad0c8917dd42a819fe77e6bdfcb84e3379c0de956469301d9fd36427a1ca501f", size = 591729, upload-time = "2026-02-20T20:20:58.395Z" }, { url = "https://files.pythonhosted.org/packages/24/b4/21f5455773d37f94b866eb3cf5caed88d6cea6dd2c6e1f9c34f463cba3ec/greenlet-3.3.2-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:97245cc10e5515dbc8c3104b2928f7f02b6813002770cfaffaf9a6e0fc2b94ef", size = 1551946, upload-time = "2026-02-20T20:49:31.102Z" }, { url = "https://files.pythonhosted.org/packages/00/68/91f061a926abead128fe1a87f0b453ccf07368666bd59ffa46016627a930/greenlet-3.3.2-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:8c1fdd7d1b309ff0da81d60a9688a8bd044ac4e18b250320a96fc68d31c209ca", size = 1618494, upload-time = "2026-02-20T20:21:06.541Z" }, @@ -1041,7 +1040,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f3/47/16400cb42d18d7a6bb46f0626852c1718612e35dcb0dffa16bbaffdf5dd2/greenlet-3.3.2-cp311-cp311-macosx_11_0_universal2.whl", hash = "sha256:c56692189a7d1c7606cb794be0a8381470d95c57ce5be03fb3d0ef57c7853b86", size = 278890, upload-time = "2026-02-20T20:19:39.263Z" }, { url = "https://files.pythonhosted.org/packages/a3/90/42762b77a5b6aa96cd8c0e80612663d39211e8ae8a6cd47c7f1249a66262/greenlet-3.3.2-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1ebd458fa8285960f382841da585e02201b53a5ec2bac6b156fc623b5ce4499f", size = 581120, upload-time = "2026-02-20T20:47:30.161Z" }, { url = "https://files.pythonhosted.org/packages/bf/6f/f3d64f4fa0a9c7b5c5b3c810ff1df614540d5aa7d519261b53fba55d4df9/greenlet-3.3.2-cp311-cp311-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a443358b33c4ec7b05b79a7c8b466f5d275025e750298be7340f8fc63dff2a55", size = 594363, upload-time = "2026-02-20T20:55:56.965Z" }, - { url = "https://files.pythonhosted.org/packages/9c/8b/1430a04657735a3f23116c2e0d5eb10220928846e4537a938a41b350bed6/greenlet-3.3.2-cp311-cp311-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:4375a58e49522698d3e70cc0b801c19433021b5c37686f7ce9c65b0d5c8677d2", size = 605046, upload-time = "2026-02-20T21:02:45.234Z" }, { url = "https://files.pythonhosted.org/packages/72/83/3e06a52aca8128bdd4dcd67e932b809e76a96ab8c232a8b025b2850264c5/greenlet-3.3.2-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8e2cd90d413acbf5e77ae41e5d3c9b3ac1d011a756d7284d7f3f2b806bbd6358", size = 594156, upload-time = "2026-02-20T20:20:59.955Z" }, { url = "https://files.pythonhosted.org/packages/70/79/0de5e62b873e08fe3cef7dbe84e5c4bc0e8ed0c7ff131bccb8405cd107c8/greenlet-3.3.2-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:442b6057453c8cb29b4fb36a2ac689382fc71112273726e2423f7f17dc73bf99", size = 1554649, upload-time = "2026-02-20T20:49:32.293Z" }, { url = "https://files.pythonhosted.org/packages/5a/00/32d30dee8389dc36d42170a9c66217757289e2afb0de59a3565260f38373/greenlet-3.3.2-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:45abe8eb6339518180d5a7fa47fa01945414d7cca5ecb745346fc6a87d2750be", size = 1619472, upload-time = "2026-02-20T20:21:07.966Z" }, @@ -1050,7 +1048,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ea/ab/1608e5a7578e62113506740b88066bf09888322a311cff602105e619bd87/greenlet-3.3.2-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:ac8d61d4343b799d1e526db579833d72f23759c71e07181c2d2944e429eb09cd", size = 280358, upload-time = "2026-02-20T20:17:43.971Z" }, { url = "https://files.pythonhosted.org/packages/a5/23/0eae412a4ade4e6623ff7626e38998cb9b11e9ff1ebacaa021e4e108ec15/greenlet-3.3.2-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3ceec72030dae6ac0c8ed7591b96b70410a8be370b6a477b1dbc072856ad02bd", size = 601217, upload-time = "2026-02-20T20:47:31.462Z" }, { url = "https://files.pythonhosted.org/packages/f8/16/5b1678a9c07098ecb9ab2dd159fafaf12e963293e61ee8d10ecb55273e5e/greenlet-3.3.2-cp312-cp312-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a2a5be83a45ce6188c045bcc44b0ee037d6a518978de9a5d97438548b953a1ac", size = 611792, upload-time = "2026-02-20T20:55:58.423Z" }, - { url = "https://files.pythonhosted.org/packages/5c/c5/cc09412a29e43406eba18d61c70baa936e299bc27e074e2be3806ed29098/greenlet-3.3.2-cp312-cp312-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:ae9e21c84035c490506c17002f5c8ab25f980205c3e61ddb3a2a2a2e6c411fcb", size = 626250, upload-time = "2026-02-20T21:02:46.596Z" }, { url = "https://files.pythonhosted.org/packages/50/1f/5155f55bd71cabd03765a4aac9ac446be129895271f73872c36ebd4b04b6/greenlet-3.3.2-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:43e99d1749147ac21dde49b99c9abffcbc1e2d55c67501465ef0930d6e78e070", size = 613875, upload-time = "2026-02-20T20:21:01.102Z" }, { url = "https://files.pythonhosted.org/packages/fc/dd/845f249c3fcd69e32df80cdab059b4be8b766ef5830a3d0aa9d6cad55beb/greenlet-3.3.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:4c956a19350e2c37f2c48b336a3afb4bff120b36076d9d7fb68cb44e05d95b79", size = 1571467, upload-time = "2026-02-20T20:49:33.495Z" }, { url = "https://files.pythonhosted.org/packages/2a/50/2649fe21fcc2b56659a452868e695634722a6655ba245d9f77f5656010bf/greenlet-3.3.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:6c6f8ba97d17a1e7d664151284cb3315fc5f8353e75221ed4324f84eb162b395", size = 1640001, upload-time = "2026-02-20T20:21:09.154Z" }, @@ -1059,7 +1056,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ac/48/f8b875fa7dea7dd9b33245e37f065af59df6a25af2f9561efa8d822fde51/greenlet-3.3.2-cp313-cp313-macosx_11_0_universal2.whl", hash = "sha256:aa6ac98bdfd716a749b84d4034486863fd81c3abde9aa3cf8eff9127981a4ae4", size = 279120, upload-time = "2026-02-20T20:19:01.9Z" }, { url = "https://files.pythonhosted.org/packages/49/8d/9771d03e7a8b1ee456511961e1b97a6d77ae1dea4a34a5b98eee706689d3/greenlet-3.3.2-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ab0c7e7901a00bc0a7284907273dc165b32e0d109a6713babd04471327ff7986", size = 603238, upload-time = "2026-02-20T20:47:32.873Z" }, { url = "https://files.pythonhosted.org/packages/59/0e/4223c2bbb63cd5c97f28ffb2a8aee71bdfb30b323c35d409450f51b91e3e/greenlet-3.3.2-cp313-cp313-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:d248d8c23c67d2291ffd47af766e2a3aa9fa1c6703155c099feb11f526c63a92", size = 614219, upload-time = "2026-02-20T20:55:59.817Z" }, - { url = "https://files.pythonhosted.org/packages/94/2b/4d012a69759ac9d77210b8bfb128bc621125f5b20fc398bce3940d036b1c/greenlet-3.3.2-cp313-cp313-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:ccd21bb86944ca9be6d967cf7691e658e43417782bce90b5d2faeda0ff78a7dd", size = 628268, upload-time = "2026-02-20T21:02:48.024Z" }, { url = "https://files.pythonhosted.org/packages/7a/34/259b28ea7a2a0c904b11cd36c79b8cef8019b26ee5dbe24e73b469dea347/greenlet-3.3.2-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b6997d360a4e6a4e936c0f9625b1c20416b8a0ea18a8e19cabbefc712e7397ab", size = 616774, upload-time = "2026-02-20T20:21:02.454Z" }, { url = "https://files.pythonhosted.org/packages/0a/03/996c2d1689d486a6e199cb0f1cf9e4aa940c500e01bdf201299d7d61fa69/greenlet-3.3.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:64970c33a50551c7c50491671265d8954046cb6e8e2999aacdd60e439b70418a", size = 1571277, upload-time = "2026-02-20T20:49:34.795Z" }, { url = "https://files.pythonhosted.org/packages/d9/c4/2570fc07f34a39f2caf0bf9f24b0a1a0a47bc2e8e465b2c2424821389dfc/greenlet-3.3.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:1a9172f5bf6bd88e6ba5a84e0a68afeac9dc7b6b412b245dd64f52d83c81e55b", size = 1640455, upload-time = "2026-02-20T20:21:10.261Z" }, @@ -1068,7 +1064,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3f/ae/8bffcbd373b57a5992cd077cbe8858fff39110480a9d50697091faea6f39/greenlet-3.3.2-cp314-cp314-macosx_11_0_universal2.whl", hash = "sha256:8d1658d7291f9859beed69a776c10822a0a799bc4bfe1bd4272bb60e62507dab", size = 279650, upload-time = "2026-02-20T20:18:00.783Z" }, { url = "https://files.pythonhosted.org/packages/d1/c0/45f93f348fa49abf32ac8439938726c480bd96b2a3c6f4d949ec0124b69f/greenlet-3.3.2-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:18cb1b7337bca281915b3c5d5ae19f4e76d35e1df80f4ad3c1a7be91fadf1082", size = 650295, upload-time = "2026-02-20T20:47:34.036Z" }, { url = "https://files.pythonhosted.org/packages/b3/de/dd7589b3f2b8372069ab3e4763ea5329940fc7ad9dcd3e272a37516d7c9b/greenlet-3.3.2-cp314-cp314-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c2e47408e8ce1c6f1ceea0dffcdf6ebb85cc09e55c7af407c99f1112016e45e9", size = 662163, upload-time = "2026-02-20T20:56:01.295Z" }, - { url = "https://files.pythonhosted.org/packages/cd/ac/85804f74f1ccea31ba518dcc8ee6f14c79f73fe36fa1beba38930806df09/greenlet-3.3.2-cp314-cp314-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:e3cb43ce200f59483eb82949bf1835a99cf43d7571e900d7c8d5c62cdf25d2f9", size = 675371, upload-time = "2026-02-20T21:02:49.664Z" }, { url = "https://files.pythonhosted.org/packages/d2/d8/09bfa816572a4d83bccd6750df1926f79158b1c36c5f73786e26dbe4ee38/greenlet-3.3.2-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:63d10328839d1973e5ba35e98cccbca71b232b14051fd957b6f8b6e8e80d0506", size = 664160, upload-time = "2026-02-20T20:21:04.015Z" }, { url = "https://files.pythonhosted.org/packages/48/cf/56832f0c8255d27f6c35d41b5ec91168d74ec721d85f01a12131eec6b93c/greenlet-3.3.2-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:8e4ab3cfb02993c8cc248ea73d7dae6cec0253e9afa311c9b37e603ca9fad2ce", size = 1619181, upload-time = "2026-02-20T20:49:36.052Z" }, { url = "https://files.pythonhosted.org/packages/0a/23/b90b60a4aabb4cec0796e55f25ffbfb579a907c3898cd2905c8918acaa16/greenlet-3.3.2-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:94ad81f0fd3c0c0681a018a976e5c2bd2ca2d9d94895f23e7bb1af4e8af4e2d5", size = 1687713, upload-time = "2026-02-20T20:21:11.684Z" }, @@ -1077,7 +1072,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/98/6d/8f2ef704e614bcf58ed43cfb8d87afa1c285e98194ab2cfad351bf04f81e/greenlet-3.3.2-cp314-cp314t-macosx_11_0_universal2.whl", hash = "sha256:e26e72bec7ab387ac80caa7496e0f908ff954f31065b0ffc1f8ecb1338b11b54", size = 286617, upload-time = "2026-02-20T20:19:29.856Z" }, { url = "https://files.pythonhosted.org/packages/5e/0d/93894161d307c6ea237a43988f27eba0947b360b99ac5239ad3fe09f0b47/greenlet-3.3.2-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8b466dff7a4ffda6ca975979bab80bdadde979e29fc947ac3be4451428d8b0e4", size = 655189, upload-time = "2026-02-20T20:47:35.742Z" }, { url = "https://files.pythonhosted.org/packages/f5/2c/d2d506ebd8abcb57386ec4f7ba20f4030cbe56eae541bc6fd6ef399c0b41/greenlet-3.3.2-cp314-cp314t-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:b8bddc5b73c9720bea487b3bffdb1840fe4e3656fba3bd40aa1489e9f37877ff", size = 658225, upload-time = "2026-02-20T20:56:02.527Z" }, - { url = "https://files.pythonhosted.org/packages/d1/67/8197b7e7e602150938049d8e7f30de1660cfb87e4c8ee349b42b67bdb2e1/greenlet-3.3.2-cp314-cp314t-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:59b3e2c40f6706b05a9cd299c836c6aa2378cabe25d021acd80f13abf81181cf", size = 666581, upload-time = "2026-02-20T21:02:51.526Z" }, { url = "https://files.pythonhosted.org/packages/8e/30/3a09155fbf728673a1dea713572d2d31159f824a37c22da82127056c44e4/greenlet-3.3.2-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b26b0f4428b871a751968285a1ac9648944cea09807177ac639b030bddebcea4", size = 657907, upload-time = "2026-02-20T20:21:05.259Z" }, { url = "https://files.pythonhosted.org/packages/f3/fd/d05a4b7acd0154ed758797f0a43b4c0962a843bedfe980115e842c5b2d08/greenlet-3.3.2-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:1fb39a11ee2e4d94be9a76671482be9398560955c9e568550de0224e41104727", size = 1618857, upload-time = "2026-02-20T20:49:37.309Z" }, { url = "https://files.pythonhosted.org/packages/6f/e1/50ee92a5db521de8f35075b5eff060dd43d39ebd46c2181a2042f7070385/greenlet-3.3.2-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:20154044d9085151bc309e7689d6f7ba10027f8f5a8c0676ad398b951913d89e", size = 1680010, upload-time = "2026-02-20T20:21:13.427Z" },