talmolab · keyaloding · Jun 28, 2024 · Jul 12, 2024 · Jul 16, 2024 · Jul 16, 2024
diff --git a/pyproject.toml b/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",

diff --git a/save_test.ipynb b/save_test.ipynb
@@ -0,0 +1,82 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "minimal_instance.pkg.nwb_images/img_0.png\n"
+     ]
+    },
+    {
+     "ename": "ValueError",
+     "evalue": "Can't write images with one color channel.",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[2], line 16\u001b[0m\n\u001b[1;32m     14\u001b[0m     img_path \u001b[38;5;241m=\u001b[39m save_path \u001b[38;5;241m/\u001b[39m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mimg_\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mi\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m.png\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m     15\u001b[0m     \u001b[38;5;28mprint\u001b[39m(img_path)\n\u001b[0;32m---> 16\u001b[0m     \u001b[43miio\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mimwrite\u001b[49m\u001b[43m(\u001b[49m\u001b[43mimg_path\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlf\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mimage\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     17\u001b[0m     img_paths\u001b[38;5;241m.\u001b[39mappend(img_path)\n\u001b[1;32m     18\u001b[0m \u001b[38;5;28mprint\u001b[39m(img_paths)\n",
+      "File \u001b[0;32m~/mambaforge3/envs/io_dev/lib/python3.12/site-packages/imageio/v3.py:147\u001b[0m, in \u001b[0;36mimwrite\u001b[0;34m(uri, image, plugin, extension, format_hint, **kwargs)\u001b[0m\n\u001b[1;32m    105\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Write an ndimage to the given URI.\u001b[39;00m\n\u001b[1;32m    106\u001b[0m \n\u001b[1;32m    107\u001b[0m \u001b[38;5;124;03mThe exact behavior depends on the file type and plugin used. To learn about\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    136\u001b[0m \n\u001b[1;32m    137\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    139\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m imopen(\n\u001b[1;32m    140\u001b[0m     uri,\n\u001b[1;32m    141\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mw\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    145\u001b[0m     extension\u001b[38;5;241m=\u001b[39mextension,\n\u001b[1;32m    146\u001b[0m ) \u001b[38;5;28;01mas\u001b[39;00m img_file:\n\u001b[0;32m--> 147\u001b[0m     encoded \u001b[38;5;241m=\u001b[39m \u001b[43mimg_file\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mwrite\u001b[49m\u001b[43m(\u001b[49m\u001b[43mimage\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    149\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m encoded\n",
+      "File \u001b[0;32m~/mambaforge3/envs/io_dev/lib/python3.12/site-packages/imageio/plugins/pillow.py:433\u001b[0m, in \u001b[0;36mPillowPlugin.write\u001b[0;34m(self, ndimage, mode, format, is_batch, **kwargs)\u001b[0m\n\u001b[1;32m    431\u001b[0m     is_batch \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[1;32m    432\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m ndimage\u001b[38;5;241m.\u001b[39mndim \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m3\u001b[39m \u001b[38;5;129;01mand\u001b[39;00m ndimage\u001b[38;5;241m.\u001b[39mshape[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m] \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m1\u001b[39m:\n\u001b[0;32m--> 433\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mCan\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mt write images with one color channel.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    434\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m ndimage\u001b[38;5;241m.\u001b[39mndim \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m3\u001b[39m \u001b[38;5;129;01mand\u001b[39;00m ndimage\u001b[38;5;241m.\u001b[39mshape[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m] \u001b[38;5;129;01min\u001b[39;00m [\u001b[38;5;241m2\u001b[39m, \u001b[38;5;241m3\u001b[39m, \u001b[38;5;241m4\u001b[39m]:\n\u001b[1;32m    435\u001b[0m     \u001b[38;5;66;03m# Note: this makes a channel-last assumption\u001b[39;00m\n\u001b[1;32m    436\u001b[0m     is_batch \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n",
+      "\u001b[0;31mValueError\u001b[0m: Can't write images with one color channel."
+     ]
+    }
+   ],
+   "source": [
+    "import imageio.v3 as iio\n",
+    "from pathlib import Path\n",
+    "import sleap_io as sio\n",
+    "\n",
+    "save_path = Path(\"minimal_instance.pkg.nwb_images\")\n",
+    "try:\n",
+    "    save_path.mkdir(parents=True, exist_ok=True)\n",
+    "except Exception as e:\n",
+    "    print(f\"An error {e} occurred. The directory could not be created.\")\n",
+    "img_paths = []\n",
+    "\n",
+    "labels_original = sio.load_slp(\"tests/data/slp/minimal_instance.pkg.slp\")\n",
+    "for i, lf in enumerate(labels_original):\n",
+    "    img_path = save_path / f\"img_{i}.png\"\n",
+    "    print(img_path)\n",
+    "    iio.imwrite(img_path, lf.image)\n",
+    "    img_paths.append(img_path)\n",
+    "print(img_paths)"
+   ]
+  }
+ ],
+ "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
+}
diff --git a/sleap_io/io/main.py b/sleap_io/io/main.py
@@ -47,33 +47,48 @@ def save_slp(
     return slp.write_labels(filename, labels, embed=embed)
 
 
-def load_nwb(filename: str) -> Labels:
+def load_nwb(filename: str, as_training: Optional[bool] = None) -> Labels:
     """Load an NWB dataset as a SLEAP `Labels` object.
 
     Args:
         filename: Path to a NWB file (`.nwb`).
+        as_training: If `True`, load the dataset as a training dataset.
-def load_nwb(filename: str, as_training: Optional[bool] = None) -> Labels:
-    """Load an NWB dataset as a SLEAP `Labels` object.
-
-    Args:
-        filename: Path to a NWB file (`.nwb`).
-        as_training: If `True`, load the dataset as a training dataset.
+def load_nwb(filename: str, as_training: Optional[bool] = None) -> Labels:
+    """Load an NWB dataset as a SLEAP `Labels` object.
+
+    Args:
+        filename: Path to a NWB file (`.nwb`).
+        as_training: If `True`, load the dataset as a training dataset.
+
+    if as_training:
+        return nwb.read_nwb_training(filename)
+    else:
+        return nwb.read_nwb(filename)
-def load_nwb(filename: str, as_training: Optional[bool] = None) -> Labels:
-    """Load an NWB dataset as a SLEAP `Labels` object.
-
-    Args:
-        filename: Path to a NWB file (`.nwb`).
-        as_training: If `True`, load the dataset as a training dataset.
+def load_nwb(filename: str, as_training: Optional[bool] = None) -> Labels:
+    """Load an NWB dataset as a SLEAP `Labels` object.
+
+    Args:
+        filename: Path to a NWB file (`.nwb`).
+        as_training: If `True`, load the dataset as a training dataset.
+
+    if as_training:
+        return nwb.read_nwb_training(filename)
+    else:
+        return nwb.read_nwb(filename)
 
     Returns:
         The dataset as a `Labels` object.
     """
     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 = None,
+    append: bool = True,
+    **kwargs,
+):
     """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.
 
     See also: nwb.write_nwb, nwb.append_nwb
     """
-    if append and Path(filename).exists():
-        nwb.append_nwb(labels, filename)
+    if as_training:
+        if append and Path(filename).exists():
+            nwb.append_nwb_training(labels, filename)
+        else:
+            nwb.write_nwb(labels, filename, None, None, True)
+
     else:
-        nwb.write_nwb(labels, filename)
+        if append and Path(filename).exists():
+            nwb.append_nwb(labels, filename, **kwargs)
+        else:
+            nwb.write_nwb(labels, filename)
-def save_nwb(
-    labels: Labels,
-    filename: str,
-    as_training: bool = None,
-    append: bool = True,
-    **kwargs,
-):
-    """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.
-
-    See also: nwb.write_nwb, nwb.append_nwb
-    """
-    if append and Path(filename).exists():
-        nwb.append_nwb(labels, filename)
-    if as_training:
-        if append and Path(filename).exists():
-            nwb.append_nwb_training(labels, filename)
-        else:
-            nwb.write_nwb(labels, filename, None, None, True)
-
-    else:
-        nwb.write_nwb(labels, filename)
-        if append and Path(filename).exists():
-            nwb.append_nwb(labels, filename, **kwargs)
-        else:
-            nwb.write_nwb(labels, filename)
+    func = nwb.labels_to_pose_training if as_training else lambda l, **kw: l
+    action = nwb.append_nwb_training if as_training else nwb.append_nwb
+    write = nwb.write_nwb_training if as_training else nwb.write_nwb
+
+    data = func(labels, **kwargs)
+    if append and Path(filename).exists():
+        action(data, filename, **kwargs)
+    else:
+        write(data, filename, **kwargs)
-def save_nwb(
-    labels: Labels,
-    filename: str,
-    as_training: bool = None,
-    append: bool = True,
-    **kwargs,
-):
-    """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.
-
-    See also: nwb.write_nwb, nwb.append_nwb
-    """
-    if append and Path(filename).exists():
-        nwb.append_nwb(labels, filename)
-    if as_training:
-        if append and Path(filename).exists():
-            nwb.append_nwb_training(labels, filename)
-        else:
-            nwb.write_nwb(labels, filename, None, None, True)
-
-    else:
-        nwb.write_nwb(labels, filename)
-        if append and Path(filename).exists():
-            nwb.append_nwb(labels, filename, **kwargs)
-        else:
-            nwb.write_nwb(labels, filename)
+    func = nwb.labels_to_pose_training if as_training else lambda l, **kw: l
+    action = nwb.append_nwb_training if as_training else nwb.append_nwb
+    write = nwb.write_nwb_training if as_training else nwb.write_nwb
+
+    data = func(labels, **kwargs)
+    if append and Path(filename).exists():
+        action(data, filename, **kwargs)
+    else:
+        write(data, filename, **kwargs)
 
 
 def load_labelstudio(
@@ -190,6 +205,8 @@ def load_file(
         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 +236,10 @@ def save_file(
 
     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, as_training=False, **kwargs)
+    elif format == "nwb_training":
+        save_nwb(labels, filename, as_training=True, **kwargs)
     elif format == "labelstudio":
         save_labelstudio(labels, filename, **kwargs)
     elif format == "jabs":