322 sequenced data split

src/animl/export.py

@@ -5,14 +5,15 @@
 
 @ Kyra Swanson 2023
 """
+import math
 import os
 import pandas as pd
 from typing import Optional, Union
 from shutil import copy2
 from random import randrange
 from pathlib import Path
 from tqdm import tqdm
-from sklearn.model_selection import train_test_split
+from sklearn.model_selection import StratifiedGroupKFold, train_test_split
 
 from animl import __version__
 from animl.file_management import build_file_manifest, save_data, save_json, save_yaml
@@ -149,6 +150,74 @@ def update_labels_from_folders(manifest: pd.DataFrame,
 
     return pd.merge(manifest, ground_truth[[unique_name, 'label']], on=unique_name)
 
+def _stratified_grouped_datasplit(manifest: pd.DataFrame,
+                                  groupby_col: str,
+                                  label_col: str = "class",
+                                  val_size: float = 0.1,
+                                  test_size: float = 0.1,
+                                  seed: int = 42):
+    """
+    Returns train_df, val_df, test_df with each group in groupby_col
+    belonging to only one of train, val, or test
+
+    It attempts to stratify by species and meet test/val size requirements
+    as best as possible, but final proportions may deviate from desired.
+    Manual verification and adjustment is recommended afterwards.
+
+    Note:
+        val_size and test_size are rounded to the nearest 5% increment (0.05)
+        to ensure StratifiedGroupKFold uses at most 20 splits
+
+    Args:
+        manifest (pd.DataFrame): DataFrame containing predictions
+        groupby_col (str): column containing group labels (e.g., "sequence")
+        label_col (str): column containing species labels
+        val_size (float): fraction of data to use for validation
+        test_size (float): fraction of data to use for testing
+        seed (int): random seed for reproducibility
+    """
+    # round val_size and test_size to nearest 0.05 and then convert to percentage
+    # val_pct and test_pct must be nonzero (at least 5%)
+    val_pct=int(max(5,round(val_size / 0.05)*5))
+    test_pct=int(max(5,round(test_size / 0.05)*5))
+
+    # the number of splits is determined by the gcd of
+    # the test, val, and train size percentages
+    gcd_val = math.gcd(math.gcd(test_pct, val_pct), 100-test_pct-val_pct)
+    n_splits = 100 // gcd_val
+
+    assert len(manifest[groupby_col].unique())>=n_splits, (
+        f"There must be at least {n_splits} groups for this test size and val size"
+    )
+
+    num_val_folds=int(val_pct*n_splits/100)
+    num_test_folds=int(test_pct*n_splits/100)
+
+    # splits data into n_splits sections with no overlapping groups
+    # and species stratified as best as possible
+    sgkf = StratifiedGroupKFold(n_splits=n_splits,
+                                shuffle=True,
+                                random_state=seed)
+
+    # now we just need to grab the right number of sections to form the val and test df
+    # folds is a list of n_splits folds. every fold isolates one of the n_splits sections
+    # folds[i][1] contains the isolated section indices for that specific fold i.
+    folds = list(sgkf.split(manifest, y=manifest[label_col], groups=manifest[groupby_col]))
+
+    # Grab the relevant folds and retrieve the indices
+    val_folds = folds[:num_val_folds]
+    test_folds = folds[num_val_folds : num_val_folds + num_test_folds]
+    val_indices = [idx for fold in val_folds for idx in fold[1]]
+    test_indices = [idx for fold in test_folds for idx in fold[1]]
+
+    # train indices are whichever indices haven't been selected for val and test
+    train_indices = list(set(manifest.index)-set(test_indices)-set(val_indices))
+
+    train_df=manifest.loc[train_indices]
+    val_df=manifest.loc[val_indices]
+    test_df=manifest.loc[test_indices]
+
+    return train_df, val_df, test_df
 
 def export_train_val_test(manifest: pd.DataFrame,
                           label_col: str = "class",
@@ -157,7 +226,8 @@ def export_train_val_test(manifest: pd.DataFrame,
                           out_dir: Optional[str] = None,
                           val_size: float = 0.1,
                           test_size: float = 0.1,
-                          seed: int = 42):
+                          seed: int = 42,
+                          groupby_col: Optional[str] = None):
     """
     Returns train_df, val_df, test_df with label_col stratified.
     test_size and val_size are fractions of the whole dataset (e.g., 0.2 -> 20%).
@@ -171,36 +241,52 @@ def export_train_val_test(manifest: pd.DataFrame,
         val_size (float): fraction of data to use for validation
         test_size (float): fraction of data to use for testing
         seed (int): random seed for reproducibility
+        groupby_col (Optional[str]): column containing group labels
+            If provided, splits data so test/val/train don't share any groups.
     """
-    assert 0 <= test_size < 1
-    assert 0 <= val_size < 1
+    assert 0 < test_size < 1
+    assert 0 < val_size < 1
     assert test_size + val_size < 1
 
     if label_col not in manifest.columns:
         raise ValueError(f"label_col '{label_col}' not found in dataframe columns")
     if file_col not in manifest.columns:
         raise ValueError(f"file_col '{file_col}' not found in dataframe columns")
 
+    manifest = manifest.reset_index(drop=True)
     # Keep only the highest confidence entry for each file, or one entry per file if no conf_col
     if conf_col not in manifest.columns:
         manifest = manifest.drop_duplicates(subset=[file_col])
     else:
         idx = manifest.groupby(file_col)[conf_col].idxmax()
-        manifest = manifest.loc[idx].reset_index(drop=True)
-
-    # Stage 1: split off test
-    trainval_df, test_df = train_test_split(manifest,
-                                            test_size=test_size,
-                                            stratify=manifest[label_col],
-                                            random_state=seed)
-
-    # Stage 2: split train/val from trainval (val_size is relative to the original dataset)
-    # Compute val fraction relative to trainval size
-    rel_val_size = val_size / (1.0 - test_size)
-    train_df, val_df = train_test_split(trainval_df,
-                                        test_size=rel_val_size,
-                                        stratify=trainval_df[label_col],
-                                        random_state=seed + 1)
+        manifest = manifest.loc[idx]
+
+    manifest=manifest.reset_index(drop=True)
+
+    if groupby_col is not None:
+        if groupby_col not in manifest.columns:
+            raise ValueError(f"groupby_col '{groupby_col}' not found in dataframe columns")
+
+        train_df, val_df, test_df = _stratified_grouped_datasplit(manifest=manifest,
+                                                                  groupby_col=groupby_col,
+                                                                  label_col=label_col,
+                                                                  val_size=val_size,
+                                                                  test_size=test_size,
+                                                                  seed=seed)
+    else:
+        # Stage 1: split off test
+        trainval_df, test_df = train_test_split(manifest,
+                                                test_size=test_size,
+                                                stratify=manifest[label_col],
+                                                random_state=seed)
+
+        # Stage 2: split train/val from trainval (val_size is relative to the original dataset)
+        # Compute val fraction relative to trainval size
+        rel_val_size = val_size / (1.0 - test_size)
+        train_df, val_df = train_test_split(trainval_df,
+                                            test_size=rel_val_size,
+                                            stratify=trainval_df[label_col],
+                                            random_state=seed + 1)
     # save to csv
     if out_dir is not None:
         save_data(train_df, out_dir + "/train_data.csv")

tests/test_export.py

@@ -20,6 +20,7 @@
     export_coco,
     export_folders,
     export_megadetector,
+    export_train_val_test,
     export_yolo,
     remove_link,
     update_labels_from_folders,
@@ -84,6 +85,16 @@ def _make_camptrapdp_manifest():
     })
 
 
+def _make_trainvaltest_manifest(n_files=20, n_groups=10):
+    """Return a minimal manifest for export_train_val_test"""
+    return pd.DataFrame({
+        'filepath': [f'/tmp/img{i % n_files}.jpg' for i in range(n_files)],
+        'class': ['deer' if i % 2 == 0 else 'elk' for i in range(n_files)],
+        'confidence': [0.5 + (i * 0.01) for i in range(n_files)],
+        'sequence': [f'seq_{i % n_groups}' for i in range(n_files)]
+    })
+
+
 class TestRemoveLink(unittest.TestCase):
 
     def test_files_are_deleted(self):
@@ -433,5 +444,75 @@ def test_label_files_written(self):
                 self.assertGreater(len(label_files), 0)
 
 
+class TestExportTrainValTest(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        cls.manifest=_make_trainvaltest_manifest()
+
+    def test_missing_column_raises_error(self):
+        with self.assertRaises(ValueError):
+            export_train_val_test(self.manifest, label_col="missing")
+        with self.assertRaises(ValueError):
+            export_train_val_test(self.manifest, file_col="missing")
+        with self.assertRaises(ValueError):
+            export_train_val_test(self.manifest, groupby_col="missing")
+
+    def test_invalid_split_sizes(self):
+        with self.assertRaises(AssertionError):
+            export_train_val_test(self.manifest, val_size=0.6, test_size=0.5)
+
+    def test_export_to_directory(self):
+        with tempfile.TemporaryDirectory() as out:
+            export_train_val_test(
+                self.manifest,
+                label_col="class",
+                file_col="filepath",
+                out_dir=out
+            )
+            self.assertTrue((Path(out) / "train_data.csv").exists())
+            self.assertTrue((Path(out) / "validate_data.csv").exists())
+            self.assertTrue((Path(out) / "test_data.csv").exists())
+
+    def test_confidence_deduplication(self):
+        temp_manifest=pd.concat([self.manifest,pd.DataFrame({
+            'filepath': ['/tmp/img0.jpg'],
+            'class': ['deer'],
+            'confidence': [0.2]
+        })])
+        train, val, test = export_train_val_test(
+            temp_manifest,
+            label_col="class",
+            file_col="filepath",
+            conf_col="confidence",
+            test_size=0.1,
+            val_size=0.1
+        )
+        combined = pd.concat([train, val, test])
+        # Verify total rows decreases by 1 due to deduplication
+        self.assertEqual(len(combined), len(temp_manifest)-1)
+
+        # Ensure the deduplciated file kept its higher
+        # original confidence (not 0.2)
+        a_row = combined[combined['filepath'] == '/tmp/img0.jpg']
+        self.assertNotAlmostEqual(a_row['confidence'].values[0], 0.2)
+
+    def test_insufficient_groups(self):
+        with self.assertRaises(AssertionError):
+            export_train_val_test(self.manifest, val_size=0.15,
+                                  test_size=0.1, groupby_col="sequence")
+
+    def test_grouped_no_leakage(self):
+        train, val, test = export_train_val_test(self.manifest, "class", "filepath",
+                                                 "confidence", val_size=0.1,
+                                                 test_size=0.1,
+                                                 groupby_col="sequence")
+        train_groups = set(train['sequence'])
+        val_groups = set(val['sequence'])
+        test_groups = set(test['sequence'])
+        self.assertTrue(train_groups.isdisjoint(val_groups))
+        self.assertTrue(train_groups.isdisjoint(test_groups))
+        self.assertTrue(val_groups.isdisjoint(test_groups))
+
 if __name__ == '__main__':
     unittest.main()