Fix split calculation and allow for not embedding

sleap_io/io/slp.py

@@ -1112,7 +1112,7 @@ def write_labels(
     """
     if Path(labels_path).exists():
         Path(labels_path).unlink()
-    if embed is not None:
+    if embed:
         embed_videos(labels_path, labels, embed)
     write_videos(labels_path, labels.videos, restore_source=(embed == "source"))
     write_tracks(labels_path, labels.tracks)

sleap_io/model/labels.py

@@ -622,6 +622,7 @@ def make_training_splits(
         n_test: int | float | None = None,
         save_dir: str | Path | None = None,
         seed: int | None = None,
+        embed: bool = True,
     ) -> tuple[Labels, Labels] | tuple[Labels, Labels, Labels]:
         """Make splits for training with embedded images.
 
@@ -635,6 +636,10 @@ def make_training_splits(
                 test split will not be saved.
             save_dir: If specified, save splits to SLP files with embedded images.
             seed: Optional integer seed to use for reproducibility.
+            embed: If `True` (the default), embed user labeled frame images in the saved
+                files, which is useful for portability but can be slow for large
+                projects. If `False`, labels are saved with references to the source
+                videos files.
 
         Returns:
             A tuple of `labels_train, labels_val` or
@@ -652,6 +657,12 @@ def make_training_splits(
             - `{save_dir}/val.pkg.slp`
             - `{save_dir}/test.pkg.slp` (if `n_test` is specified)
 
+            If `embed` is `False`, the files will be saved without embedded images to:
+
+            - `{save_dir}/train.slp`
+            - `{save_dir}/val.slp`
+            - `{save_dir}/test.slp` (if `n_test` is specified)
+
         See also: `Labels.split`
         """
         # Clean up labels.
@@ -660,16 +671,23 @@ def make_training_splits(
         labels.suggestions = []
         labels.clean()
 
-        # Make splits.
+        # Make train split.
         labels_train, labels_rest = labels.split(n_train, seed=seed)
+
+        # Make test split.
         if n_test is not None:
             if n_test < 1:
                 n_test = (n_test * len(labels)) / len(labels_rest)
             labels_test, labels_rest = labels_rest.split(n=n_test, seed=seed)
+
+        # Make val split.
         if n_val is not None:
             if n_val < 1:
                 n_val = (n_val * len(labels)) / len(labels_rest)
-            labels_val, _ = labels_rest.split(n=n_val, seed=seed)
+            if isinstance(n_val, float) and n_val == 1.0:
+                labels_val = labels_rest
+            else:
+                labels_val, _ = labels_rest.split(n=n_val, seed=seed)
         else:
             labels_val = labels_rest
 
@@ -678,9 +696,14 @@ def make_training_splits(
             save_dir = Path(save_dir)
             save_dir.mkdir(exist_ok=True, parents=True)
 
-            labels_train.save(save_dir / "train.pkg.slp", embed="user")
-            labels_val.save(save_dir / "val.pkg.slp", embed="user")
-            labels_test.save(save_dir / "test.pkg.slp", embed="user")
+            if embed:
+                labels_train.save(save_dir / "train.pkg.slp", embed="user")
+                labels_val.save(save_dir / "val.pkg.slp", embed="user")
+                labels_test.save(save_dir / "test.pkg.slp", embed="user")
+            else:
+                labels_train.save(save_dir / "train.slp", embed=False)
+                labels_val.save(save_dir / "val.slp", embed=False)
+                labels_test.save(save_dir / "test.slp", embed=False)
 
         if n_test is None:
             return labels_train, labels_val

tests/model/test_labels.py

@@ -537,7 +537,7 @@ def test_split(slp_real_data, tmp_path):
     )
 
 
-def test_make_training_splits(slp_real_data, tmp_path):
+def test_make_training_splits(slp_real_data):
     labels = load_slp(slp_real_data)
     assert len(labels.user_labeled_frames) == 5
 
@@ -568,6 +568,11 @@ def test_make_training_splits(slp_real_data, tmp_path):
     assert len(val) == 1
     assert len(test) == 1
 
+    train, val, test = labels.make_training_splits(n_train=0.4, n_val=0.4, n_test=0.2)
+    assert len(train) == 2
+    assert len(val) == 2
+    assert len(test) == 1
+
 
 def test_make_training_splits_save(slp_real_data, tmp_path):
     labels = load_slp(slp_real_data)
@@ -587,3 +592,38 @@ def test_make_training_splits_save(slp_real_data, tmp_path):
     assert train_.provenance["source_labels"] == slp_real_data
     assert val_.provenance["source_labels"] == slp_real_data
     assert test_.provenance["source_labels"] == slp_real_data
+
+
+@pytest.mark.parametrize("embed", [True, False])
+def test_make_training_splits_save(slp_real_data, tmp_path, embed):
+    labels = load_slp(slp_real_data)
+
+    train, val, test = labels.make_training_splits(
+        0.6, 0.2, 0.2, save_dir=tmp_path, embed=embed
+    )
+
+    if embed:
+        train_, val_, test_ = (
+            load_slp(tmp_path / "train.pkg.slp"),
+            load_slp(tmp_path / "val.pkg.slp"),
+            load_slp(tmp_path / "test.pkg.slp"),
+        )
+    else:
+        train_, val_, test_ = (
+            load_slp(tmp_path / "train.slp"),
+            load_slp(tmp_path / "val.slp"),
+            load_slp(tmp_path / "test.slp"),
+        )
+
+    assert len(train_) == len(train)
+    assert len(val_) == len(val)
+    assert len(test_) == len(test)
+
+    if embed:
+        assert train_.provenance["source_labels"] == slp_real_data
+        assert val_.provenance["source_labels"] == slp_real_data
+        assert test_.provenance["source_labels"] == slp_real_data
+    else:
+        assert train_.video.filename == labels.video.filename
+        assert val_.video.filename == labels.video.filename
+        assert test_.video.filename == labels.video.filename