core/database/classes/dataset.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+import random
+
+from pydantic import BaseModel
+
+
+class Dataset(BaseModel):
+    train_names: list[str]
+    test_names: list[str]
+    validation_names: list[str]
+
+    @classmethod
+    def fraction_map(cls) -> dict[str, float]:
+        """Dict with train, test and validation fractions."""
+        # define map
+        split_map = {
+            'train': 0.7,
+            'test': 0.2,
+            'validation': 0.1,
+        }
+        # sanity check
+        assert sum(split_map.values()) == 1.0
+        return split_map
+
+    @classmethod
+    def new_from_name_list(
+        cls,
+        name_list: list[str],
+    ) -> Dataset:
+        """Generate new dataset from list of filenames."""
+        # calculate split fractions
+        fraction_map = cls.fraction_map()
+        num_total = len(name_list)
+        num_validation = round(num_total * fraction_map['validation'])
+        num_test = round(num_total * fraction_map['test'])
+        # split data
+        validation_name_list = random.choices(name_list, k=num_validation)
+        name_list = [
+            name for name in name_list if name not in validation_name_list
+        ]
+        test_name_list = random.choices(name_list, k=num_test)
+        train_name_list = [
+            name for name in name_list if name not in test_name_list
+        ]
+        # instantiate object
+        dataset = Dataset(
+            train_names=train_name_list,
+            test_names=test_name_list,
+            validation_names=validation_name_list,
+        )
+        return dataset