make sure validation loader always has at least one sample (closes #189)

lightning_pose/data/utils.py

@@ -369,8 +369,21 @@ def split_sizes_from_probabilities(
     train_number = int(np.floor(train_probability * total_number))
     val_number = int(np.floor(val_probability * total_number))
 
-    # if we lose extra examples by flooring, send these to test_number
-    test_number = total_number - train_number - val_number
+    # if we lose extra examples by flooring, send these to train_number or test_number, depending
+    leftover = total_number - train_number - val_number
+    if leftover < 5:
+        # very few samples, let's bulk up train
+        train_number += leftover
+        test_number = 0
+    else:
+        test_number = leftover
+
+    # make sure that we have at least one validation sample
+    if val_number == 0:
+        train_number -= 1
+        val_number += 1
+        if train_number < 1:
+            raise ValueError("Must have at least two labeled frames, one train and one validation")
 
     # assert that we're using all datapoints
     assert train_number + test_number + val_number == total_number

tests/data/test_utils.py

@@ -78,6 +78,23 @@ def test_split_sizes_from_probabilities():
     out = split_sizes_from_probabilities(101, train_probability=0.7)
     assert out[0] == 70 and out[1] == 15 and out[2] == 16
 
+    # make sure we have at least one example in the validation set
+    total_number = 10
+    train_prob = 0.95
+    val_prob = 0.05
+    out = split_sizes_from_probabilities(
+        total_number, train_probability=train_prob, val_probability=val_prob
+    )
+    assert sum(out) == total_number
+    assert out[0] == 9
+    assert out[1] == 1
+    assert out[2] == 0
+
+    # make sure an error is raised if there are not enough labeled frames
+    total_number = 1
+    with pytest.raises(ValueError):
+        split_sizes_from_probabilities(total_number, train_probability=train_prob)
+
 
 def test_clean_any_nans():