diff --git a/CHANGELOG.md b/CHANGELOG.md index 84df743..6e46f85 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,7 @@ * Replaced OpenCV-generated synthetic test videos with committed stub videos from DANDI (H.264, MJPEG, mp4v) at 160x120 resolution. [PR #24](https://github.com/catalystneuro/nwb-video-widgets/pull/24) * Added comprehensive unit tests for PoseEstimation widgets covering discovery, widget creation, lazy loading, keypoint colors, error handling, and video mapping. [PR #16](https://github.com/catalystneuro/nwb-video-widgets/pull/16) +* Added support for Pose Estimation Objects anywhere in the NWB file [PR #17](https://github.com/catalystneuro/nwb-video-widgets/pull/17) # v0.1.5 (2026-02-03) diff --git a/src/nwb_video_widgets/_utils.py b/src/nwb_video_widgets/_utils.py index b4b5d71..b2a51a5 100644 --- a/src/nwb_video_widgets/_utils.py +++ b/src/nwb_video_widgets/_utils.py @@ -404,6 +404,10 @@ def start_video_server(directory: Path) -> int: def discover_pose_estimation_cameras(nwbfile: NWBFile) -> dict: """Discover all PoseEstimation containers in an NWB file. + Searches all objects in the file regardless of where they are stored, + so PoseEstimation data in any processing module (e.g. 'pose_estimation', + 'behavior') is found. + Parameters ---------- nwbfile : NWBFile @@ -412,20 +416,13 @@ def discover_pose_estimation_cameras(nwbfile: NWBFile) -> dict: Returns ------- dict - Mapping of camera names to PoseEstimation objects from - processing['pose_estimation']. + Mapping of camera names to PoseEstimation objects. """ - if "pose_estimation" not in nwbfile.processing: - return {} - - pose_module = nwbfile.processing["pose_estimation"] - - # Get only PoseEstimation objects (not Skeletons or other types) cameras = {} - for name, obj in pose_module.data_interfaces.items(): - if type(obj).__name__ == "PoseEstimation": - cameras[name] = obj - + for obj in nwbfile.objects.values(): + if obj.neurodata_type == "PoseEstimation": + assert obj.name not in cameras, f"Duplicate PoseEstimation name found: {obj.name}" + cameras[obj.name] = obj return cameras diff --git a/src/nwb_video_widgets/dandi_pose_widget.py b/src/nwb_video_widgets/dandi_pose_widget.py index 2534ecc..e1e02ea 100644 --- a/src/nwb_video_widgets/dandi_pose_widget.py +++ b/src/nwb_video_widgets/dandi_pose_widget.py @@ -29,7 +29,7 @@ class NWBDANDIPoseEstimationWidget(anywidget.AnyWidget): Overlays DeepLabCut keypoints on streaming video with support for camera selection via a settings panel. - This widget discovers PoseEstimation containers in processing['pose_estimation'] + This widget discovers PoseEstimation containers anywhere in the NWB file and resolves video paths to S3 URLs via the DANDI API. An interactive settings panel allows users to select which camera to display. @@ -156,13 +156,10 @@ def __init__( colormap_name = "tab10" custom_colors = keypoint_colors - # Get pose estimation container - if "pose_estimation" not in nwbfile.processing: - raise ValueError("NWB file does not contain pose_estimation processing module") - pose_estimation = nwbfile.processing["pose_estimation"] - - # Get all PoseEstimation containers (excludes Skeletons and other metadata) + # Get all PoseEstimation containers (location-agnostic) pose_containers = discover_pose_estimation_cameras(nwbfile) + if not pose_containers: + raise ValueError("NWB file does not contain any PoseEstimation objects") available_cameras = list(pose_containers.keys()) # Get camera info for settings panel display @@ -185,7 +182,7 @@ def __init__( selected_camera = "" # Store references for lazy loading (not synced to JS) - self._pose_estimation = pose_estimation + self._pose_containers = pose_containers self._cmap = plt.get_cmap(colormap_name) self._custom_colors = custom_colors @@ -216,7 +213,7 @@ def _on_camera_selected(self, change): try: # Load pose data for this camera camera_data = self._load_camera_pose_data( - self._pose_estimation, camera_name, self._cmap, self._custom_colors + self._pose_containers, camera_name, self._cmap, self._custom_colors ) # Update all_camera_data (must create new dict for traitlets to detect change) @@ -296,7 +293,7 @@ def _get_video_urls_from_dandi( return video_urls @staticmethod - def _load_camera_pose_data(pose_estimation, camera_name: str, cmap, custom_colors: dict) -> dict: + def _load_camera_pose_data(pose_containers: dict, camera_name: str, cmap, custom_colors: dict) -> dict: """Load pose data for a single camera. Returns a dict with: @@ -304,7 +301,7 @@ def _load_camera_pose_data(pose_estimation, camera_name: str, cmap, custom_color - pose_coordinates: {name: [[x, y], ...]} as JSON-serializable lists - timestamps: [t0, t1, ...] as JSON-serializable list """ - camera_pose = pose_estimation[camera_name] + camera_pose = pose_containers[camera_name] keypoint_names = list(camera_pose.pose_estimation_series.keys()) n_kp = len(keypoint_names) diff --git a/src/nwb_video_widgets/local_pose_widget.py b/src/nwb_video_widgets/local_pose_widget.py index 38d3b4a..8da51bc 100644 --- a/src/nwb_video_widgets/local_pose_widget.py +++ b/src/nwb_video_widgets/local_pose_widget.py @@ -26,7 +26,7 @@ class NWBLocalPoseEstimationWidget(anywidget.AnyWidget): Overlays DeepLabCut keypoints on streaming video with support for camera selection via a settings panel. - This widget discovers PoseEstimation containers in processing['pose_estimation'] + This widget discovers PoseEstimation containers anywhere in the NWB file and resolves video paths relative to the NWB file location. An interactive settings panel allows users to select which camera to display. @@ -37,8 +37,7 @@ class NWBLocalPoseEstimationWidget(anywidget.AnyWidget): Parameters ---------- nwbfile : pynwb.NWBFile - NWB file containing pose estimation in processing['pose_estimation']. - Must have been loaded from disk. + NWB file containing pose estimation. Must have been loaded from disk. video_nwbfile : pynwb.NWBFile, optional NWB file containing video ImageSeries in acquisition. If not provided, videos are assumed to be in `nwbfile`. Use this when videos are in a @@ -125,13 +124,10 @@ def __init__( colormap_name = "tab10" custom_colors = keypoint_colors - # Get pose estimation container - if "pose_estimation" not in nwbfile.processing: - raise ValueError("NWB file does not contain pose_estimation processing module") - pose_estimation = nwbfile.processing["pose_estimation"] - - # Get all PoseEstimation containers (excludes Skeletons and other metadata) + # Get all PoseEstimation containers (location-agnostic) pose_containers = discover_pose_estimation_cameras(nwbfile) + if not pose_containers: + raise ValueError("NWB file does not contain any PoseEstimation objects") available_cameras = list(pose_containers.keys()) # Get camera info for settings panel display @@ -154,7 +150,7 @@ def __init__( selected_camera = "" # Store references for lazy loading (not synced to JS) - self._pose_estimation = pose_estimation + self._pose_containers = pose_containers self._cmap = plt.get_cmap(colormap_name) self._custom_colors = custom_colors @@ -185,7 +181,7 @@ def _on_camera_selected(self, change): try: # Load pose data for this camera camera_data = self._load_camera_pose_data( - self._pose_estimation, camera_name, self._cmap, self._custom_colors + self._pose_containers, camera_name, self._cmap, self._custom_colors ) # Update all_camera_data (must create new dict for traitlets to detect change) @@ -279,7 +275,7 @@ def _get_video_urls_from_local(nwbfile: NWBFile) -> dict[str, str]: return video_urls @staticmethod - def _load_camera_pose_data(pose_estimation, camera_name: str, cmap, custom_colors: dict) -> dict: + def _load_camera_pose_data(pose_containers: dict, camera_name: str, cmap, custom_colors: dict) -> dict: """Load pose data for a single camera. Returns a dict with: @@ -287,7 +283,7 @@ def _load_camera_pose_data(pose_estimation, camera_name: str, cmap, custom_color - pose_coordinates: {name: [[x, y], ...]} as JSON-serializable lists - timestamps: [t0, t1, ...] as JSON-serializable list """ - camera_pose = pose_estimation[camera_name] + camera_pose = pose_containers[camera_name] keypoint_names = list(camera_pose.pose_estimation_series.keys()) n_kp = len(keypoint_names) diff --git a/tests/conftest.py b/tests/conftest.py index f2d7c51..e69f959 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -101,6 +101,40 @@ def nwbfile_with_multiple_cameras_pose(tmp_path): return read_nwb(nwb_path) +@pytest.fixture +def nwbfile_with_behavior_module_pose(tmp_path): + """Create an NWB file with pose estimation stored in the 'behavior' processing module.""" + nwbfile = create_nwbfile_with_pose_estimation( + camera_names=["LeftCamera"], + keypoint_names=["Nose", "LeftEar", "RightEar"], + num_frames=30, + processing_module_name="behavior", + ) + nwb_path = tmp_path / "test_behavior_pose.nwb" + + with NWBHDF5IO(nwb_path, "w") as io: + io.write(nwbfile) + + return read_nwb(nwb_path) + + +@pytest.fixture +def nwbfile_with_custom_module_pose(tmp_path): + """Create an NWB file with pose estimation stored in a custom-named processing module.""" + nwbfile = create_nwbfile_with_pose_estimation( + camera_names=["LeftCamera", "RightCamera"], + keypoint_names=["Nose", "LeftEar"], + num_frames=30, + processing_module_name="my_custom_module", + ) + nwb_path = tmp_path / "test_custom_pose.nwb" + + with NWBHDF5IO(nwb_path, "w") as io: + io.write(nwbfile) + + return read_nwb(nwb_path) + + @pytest.fixture def nwbfile_with_videos_and_pose(tmp_path): """Create an NWB file with both videos and pose estimation.""" diff --git a/tests/fixtures/synthetic_nwb.py b/tests/fixtures/synthetic_nwb.py index 212fcde..01504d5 100644 --- a/tests/fixtures/synthetic_nwb.py +++ b/tests/fixtures/synthetic_nwb.py @@ -60,6 +60,7 @@ def create_nwbfile_with_pose_estimation( timestamps: np.ndarray | None = None, video_width: int = 160, video_height: int = 120, + processing_module_name: str = "pose_estimation", ) -> NWBFile: """Create an NWBFile with PoseEstimation data. @@ -77,6 +78,8 @@ def create_nwbfile_with_pose_estimation( Width of the source video in pixels, by default 160 video_height : int, optional Height of the source video in pixels, by default 120 + processing_module_name : str, optional + Name of the processing module to store pose data in, by default "pose_estimation" Returns ------- @@ -85,9 +88,9 @@ def create_nwbfile_with_pose_estimation( """ nwbfile = mock_NWBFile() - # Create pose_estimation processing module + # Create pose estimation processing module pose_module = ProcessingModule( - name="pose_estimation", + name=processing_module_name, description="Pose estimation data from DeepLabCut or similar", ) nwbfile.add_processing_module(pose_module) diff --git a/tests/test_local_pose_widget.py b/tests/test_local_pose_widget.py index f239ac8..bfaf121 100644 --- a/tests/test_local_pose_widget.py +++ b/tests/test_local_pose_widget.py @@ -137,8 +137,8 @@ class TestErrorHandling: """Tests for error handling.""" def test_raises_for_missing_pose_module(self, nwbfile_with_single_video): - """Test that error is raised when pose_estimation module is missing.""" - with pytest.raises(ValueError, match="pose_estimation processing module"): + """Test that error is raised when no PoseEstimation objects are found.""" + with pytest.raises(ValueError, match="NWB file does not contain any PoseEstimation objects"): NWBLocalPoseEstimationWidget(nwbfile_with_single_video) def test_raises_for_in_memory_nwbfile(self): @@ -173,6 +173,37 @@ def test_raises_for_in_memory_nwbfile(self): NWBLocalPoseEstimationWidget(nwbfile) +class TestProcessingModuleLocation: + """Tests that pose estimation is discovered regardless of processing module name.""" + + def test_pose_in_pose_estimation_module(self, nwbfile_with_single_camera_pose): + """Standard case: pose in 'pose_estimation' processing module.""" + widget = NWBLocalPoseEstimationWidget(nwbfile_with_single_camera_pose) + assert "LeftCamera" in widget.available_cameras + + def test_pose_in_behavior_module(self, nwbfile_with_behavior_module_pose): + """Pose stored in 'behavior' module is discovered correctly.""" + widget = NWBLocalPoseEstimationWidget(nwbfile_with_behavior_module_pose) + assert len(widget.available_cameras) == 1 + assert "LeftCamera" in widget.available_cameras + + def test_pose_data_loads_from_behavior_module(self, nwbfile_with_behavior_module_pose): + """Pose data in 'behavior' module loads correctly on camera selection.""" + widget = NWBLocalPoseEstimationWidget(nwbfile_with_behavior_module_pose) + widget.selected_camera = "LeftCamera" + assert "LeftCamera" in widget.all_camera_data + camera_data = widget.all_camera_data["LeftCamera"] + assert "keypoint_metadata" in camera_data + assert "Nose" in camera_data["keypoint_metadata"] + + def test_pose_in_custom_module(self, nwbfile_with_custom_module_pose): + """Pose stored in a custom-named module is discovered correctly.""" + widget = NWBLocalPoseEstimationWidget(nwbfile_with_custom_module_pose) + assert len(widget.available_cameras) == 2 + assert "LeftCamera" in widget.available_cameras + assert "RightCamera" in widget.available_cameras + + class TestVideoNameMapping: """Tests for video name to URL mapping."""