diff --git a/src/animl/export.py b/src/animl/export.py index eafc276..b7dbada 100644 --- a/src/animl/export.py +++ b/src/animl/export.py @@ -5,6 +5,7 @@ @ Kyra Swanson 2023 """ +import math import os import pandas as pd from typing import Optional, Union @@ -12,7 +13,7 @@ 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,9 +241,11 @@ 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: @@ -181,26 +253,40 @@ def export_train_val_test(manifest: pd.DataFrame, 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") diff --git a/tests/test_export.py b/tests/test_export.py index 46c90fb..d91aaa7 100644 --- a/tests/test_export.py +++ b/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()