Python 3.12 GitHub tests

README.md

@@ -41,6 +41,11 @@ labels = sio.load_file("predictions.slp")
 sio.save_file(labels, "predictions.nwb")
 # Or:
 # labels.save("predictions.nwb")
+
+# Save to an NWB file and convert SLEAP training data to NWB training data:
+frame_inds = [i for i in range(20)]
+sio.save_file(labels, "predictions.nwb", as_training=True, frame_inds=frame_inds)
+# This will save the first 20 frames of the video as individual images
 ```
 
 ### Convert labels to raw arrays

environment.yml

@@ -13,7 +13,7 @@ dependencies:
   - pandas
   - simplejson
   - h5py >=3.8.0
-  - hdmf
+  - hdmf == 3.13.0
   - numpy <2.0.0
   - opencv
   - pynwb
@@ -24,7 +24,7 @@ dependencies:
   - pydocstyle
   - toml
   - twine
-  - build
+  - python-build
   - pip
   - pip:
     - "--editable=.[dev]"

profiling.ipynb

@@ -0,0 +1,131 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import cProfile\n",
+    "from sleap_io.io.main import load_slp\n",
+    "from sleap_io.io.nwb import NWBFile, append_nwb_data, datetime"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "labels = load_slp(\"tests/data/slp/centered_pair_predictions.slp\")\n",
+    "labels.clean(tracks=True)\n",
+    "nwbfile = NWBFile(\n",
+    "    session_description=\"Testing session for nwb\",\n",
+    "    identifier=\"identifier\",\n",
+    "    session_start_time=datetime.datetime.now(datetime.timezone.utc),\n",
+    ")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/keya/mambaforge3/envs/io_dev/lib/python3.12/site-packages/pynwb/base.py:195: UserWarning: Timeseries has a rate of 0.0 Hz, but the length of the data is greater than 1.\n",
+      "  warn('Timeseries has a rate of 0.0 Hz, but the length of the data is greater than 1.')\n",
+      "/Users/keya/mambaforge3/envs/io_dev/lib/python3.12/site-packages/pynwb/base.py:195: UserWarning: Timeseries has a rate of 0.0 Hz, but the length of the data is greater than 1.\n",
+      "  warn('Timeseries has a rate of 0.0 Hz, but the length of the data is greater than 1.')\n",
+      "/Users/keya/mambaforge3/envs/io_dev/lib/python3.12/site-packages/pynwb/base.py:195: UserWarning: Timeseries has a rate of 0.0 Hz, but the length of the data is greater than 1.\n",
+      "  warn('Timeseries has a rate of 0.0 Hz, but the length of the data is greater than 1.')\n",
+      "/Users/keya/mambaforge3/envs/io_dev/lib/python3.12/site-packages/pynwb/base.py:195: UserWarning: Timeseries has a rate of 0.0 Hz, but the length of the data is greater than 1.\n",
+      "  warn('Timeseries has a rate of 0.0 Hz, but the length of the data is greater than 1.')\n",
+      "/Users/keya/mambaforge3/envs/io_dev/lib/python3.12/site-packages/pynwb/base.py:195: UserWarning: Timeseries has a rate of 0.0 Hz, but the length of the data is greater than 1.\n",
+      "  warn('Timeseries has a rate of 0.0 Hz, but the length of the data is greater than 1.')\n",
+      "/Users/keya/mambaforge3/envs/io_dev/lib/python3.12/site-packages/pynwb/base.py:195: UserWarning: Timeseries has a rate of 0.0 Hz, but the length of the data is greater than 1.\n",
+      "  warn('Timeseries has a rate of 0.0 Hz, but the length of the data is greater than 1.')\n",
+      "/Users/keya/mambaforge3/envs/io_dev/lib/python3.12/site-packages/pynwb/base.py:195: UserWarning: Timeseries has a rate of 0.0 Hz, but the length of the data is greater than 1.\n",
+      "  warn('Timeseries has a rate of 0.0 Hz, but the length of the data is greater than 1.')\n",
+      "/Users/keya/mambaforge3/envs/io_dev/lib/python3.12/site-packages/pynwb/base.py:195: UserWarning: Timeseries has a rate of 0.0 Hz, but the length of the data is greater than 1.\n",
+      "  warn('Timeseries has a rate of 0.0 Hz, but the length of the data is greater than 1.')\n"
+     ]
+    }
+   ],
+   "source": [
+    "cProfile.run(\"nwbfile = append_nwb_data(labels, nwbfile)\", \"profile_output2\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Fri Aug 30 12:14:46 2024    profile_output2\n",
+      "\n",
+      "         318923324 function calls (305248509 primitive calls) in 88.785 seconds\n",
+      "\n",
+      "   Ordered by: internal time\n",
+      "   List reduced from 1027 to 10 due to restriction <10>\n",
+      "\n",
+      "   ncalls  tottime  percall  cumtime  percall filename:lineno(function)\n",
+      "  6211428   25.942    0.000   59.150    0.000 /Users/keya/mambaforge3/envs/io_dev/lib/python3.12/site-packages/hdmf/utils.py:191(__parse_args)\n",
+      "75422705/75420717    7.006    0.000    8.989    0.000 {built-in method builtins.isinstance}\n",
+      "16770924/10891398    6.833    0.000   13.012    0.000 /Users/keya/mambaforge3/envs/io_dev/lib/python3.12/site-packages/hdmf/utils.py:70(check_type)\n",
+      "  6211428    6.045    0.000   67.803    0.000 /Users/keya/mambaforge3/envs/io_dev/lib/python3.12/site-packages/hdmf/utils.py:634(_check_args)\n",
+      "  6224242    5.523    0.000    8.837    0.000 /Users/keya/mambaforge3/envs/io_dev/lib/python3.12/site-packages/hdmf/utils.py:764(getargs)\n",
+      " 37500806    3.995    0.000    3.995    0.000 {method 'get' of 'dict' objects}\n",
+      "6204020/705    3.993    0.000   87.528    0.124 /Users/keya/mambaforge3/envs/io_dev/lib/python3.12/site-packages/hdmf/utils.py:666(func_call)\n",
+      "42094046/42060076    3.263    0.000    3.271    0.000 {built-in method builtins.len}\n",
+      "  6211428    2.933    0.000    9.012    0.000 /Users/keya/mambaforge3/envs/io_dev/lib/python3.12/collections/__init__.py:595(__init__)\n",
+      "  1489008    2.401    0.000   60.155    0.000 /Users/keya/mambaforge3/envs/io_dev/lib/python3.12/site-packages/hdmf/build/manager.py:743(register_container_type)\n",
+      "\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "<pstats.Stats at 0x15c6672f0>"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "import pstats\n",
+    "p = pstats.Stats(\"profile_output2\")\n",
+    "p.sort_stats(\"tottime\").print_stats(10)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "io_dev",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.4"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

pyproject.toml

@@ -24,7 +24,7 @@ dependencies = [
     "attrs",
     "h5py>=3.8.0",
     "pynwb",
-    "ndx-pose",
+    "ndx-pose @ git+https://github.com/rly/ndx-pose@a847ad4be75e60ef9e413b8cbfc99c616fc9fd05",
     "pandas",
     "simplejson",
     "imageio",
@@ -45,7 +45,7 @@ dev = [
     "pydocstyle",
     "toml",
     "twine",
-    "build",
+    "python-build",
     "mkdocs-material[imaging]",
     "mkdocs-jupyter",
     "mkdocstrings[python]>=0.18",

sleap_io/io/main.py

@@ -6,6 +6,8 @@
 from typing import Optional, Union
 from pathlib import Path
 
+from pynwb import NWBHDF5IO
+
 
 def load_slp(filename: str) -> Labels:
     """Load a SLEAP dataset.
@@ -59,21 +61,45 @@
     return nwb.read_nwb(filename)
 
 
-def save_nwb(labels: Labels, filename: str, append: bool = True):
+def save_nwb(
+    labels: Labels,
+    filename: str,
+    as_training: bool = False,
+    append: bool = True,
+    frame_inds: Optional[list[int]] = None,
+    frame_path: Optional[str] = None,
+):
     """Save a SLEAP dataset to NWB format.
 
     Args:
         labels: A SLEAP `Labels` object (see `load_slp`).
         filename: Path to NWB file to save to. Must end in `.nwb`.
+        as_training: If `True`, save the dataset as a training dataset.
         append: If `True` (the default), append to existing NWB file. File will be
             created if it does not exist.
+        frame_inds: Optional list of frame indices to save. If None, all frames
+            will be saved.
+        frame_path: The path to save the frames. If None, the path is the video
+            filename without the extension.
 
-    See also: nwb.write_nwb, nwb.append_nwb
+    See also: nwb.write_nwb, nwb.append_nwb, nwb.append_nwb_training
     """
     if append and Path(filename).exists():
-        nwb.append_nwb(labels, filename)
+        nwb.append_nwb(
+            labels,
+            filename,
+            as_training=as_training,
+            frame_inds=frame_inds,
+            frame_path=frame_path,
+        )
     else:
-        nwb.write_nwb(labels, filename)
+        nwb.write_nwb(
+            labels,
+            filename,
+            as_training=as_training,
+            frame_inds=frame_inds,
+            frame_path=frame_path,
+        )
 
 
 def load_labelstudio(
@@ -190,6 +216,8 @@
         return load_jabs(filename, **kwargs)
     elif format == "video":
         return load_video(filename, **kwargs)
+    else:
+        raise ValueError(f"Unknown format '{format}' for filename: '{filename}'.")
 
 
 def save_file(
@@ -219,8 +247,10 @@
 
     if format == "slp":
         save_slp(labels, filename, **kwargs)
-    elif format == "nwb":
-        save_nwb(labels, filename, **kwargs)
+    elif format in ("nwb", "nwb_predictions"):
+        save_nwb(labels, filename, False)
+    elif format == "nwb_training":
+        save_nwb(labels, filename, True, frame_inds=kwargs.pop("frame_inds", None))
     elif format == "labelstudio":
         save_labelstudio(labels, filename, **kwargs)
     elif format == "jabs":