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/sleap_io/io/main.py b/sleap_io/io/main.py
@@ -47,34 +47,50 @@ 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.
 
  Returns:
  The dataset as a `Labels` object.
  """
- return nwb.read_nwb(filename)
+ if as_training is None:
+ return 
+
+ 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.
-
- Returns:
- The dataset as a `Labels` object.
- """
- return nwb.read_nwb(filename)
- if as_training is None:
- return 
-
- if as_training:
- return nwb.read_nwb_training(filename)
- else:
- return nwb.read_nwb(filename)
+ 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.
-
- Returns:
- The dataset as a `Labels` object.
- """
- return nwb.read_nwb(filename)
- if as_training is None:
- return 
-
- if as_training:
- return nwb.read_nwb_training(filename)
- else:
- return nwb.read_nwb(filename)
+ if as_training:
+ return nwb.read_nwb_training(filename)
+ else:
+ 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)
- else:
- nwb.write_nwb(labels, filename)
+ if as_training:
+ pose_training = nwb.labels_to_pose_training(labels, **kwargs)
+ if append and Path(filename).exists():
+ nwb.append_nwb_training(pose_training, filename, **kwargs)
+ else:
+ nwb.write_nwb_training(pose_training, filename, **kwargs)
 
+ else:
+ 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)
- else:
- nwb.write_nwb(labels, filename)
- if as_training:
- pose_training = nwb.labels_to_pose_training(labels, **kwargs)
- if append and Path(filename).exists():
- nwb.append_nwb_training(pose_training, filename, **kwargs)
- else:
- nwb.write_nwb_training(pose_training, filename, **kwargs)
-
- else:
- 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
+ """
+ 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)
- else:
- nwb.write_nwb(labels, filename)
- if as_training:
- pose_training = nwb.labels_to_pose_training(labels, **kwargs)
- if append and Path(filename).exists():
- nwb.append_nwb_training(pose_training, filename, **kwargs)
- else:
- nwb.write_nwb_training(pose_training, filename, **kwargs)
-
- else:
- 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
+ """
+ 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(
  filename: str, skeleton: Optional[Union[Skeleton, list[str]]] = None
@@ -190,6 +206,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,7 +237,7 @@ def save_file(
 
  if format == "slp":
  save_slp(labels, filename, **kwargs)
- elif format == "nwb":
+ elif format in ("nwb", "nwb_training", "nwb_predictions"):
  save_nwb(labels, filename, **kwargs)
  elif format == "labelstudio":
  save_labelstudio(labels, filename, **kwargs)