adding unit test for selections utils

voxel51 · brimoor · Oct 20, 2020 · Oct 20, 2020 · Oct 20, 2020 · Oct 20, 2020
commit 443ea65a377c9a33bd2b7ab1e16e6c85e6f6d8d0
diff --git a/tests/selection/selection_tests.py b/tests/selection/selection_tests.py
@@ -0,0 +1,73 @@
+"""
+Unit tests for the :mod:`fiftyone.utils.selections` module.
+
+| Copyright 2017-2020, Voxel51, Inc.
+| `voxel51.com <https://voxel51.com/>`_
+|
+"""
+import random
+import unittest
+
+import fiftyone as fo
+import fiftyone.core.dataset as fod
+import fiftyone.utils.selections as fous
+import fiftyone.zoo as foz
+
+
+class SelectionTests(unittest.TestCase):
+    def test_list_datasets(self):
+        num_samples_to_select = 3
+        max_objects_per_sample_to_select = 5
+
+        dataset = foz.load_zoo_dataset(
+            "coco-2017",
+            split="validation",
+            dataset_name=fod.get_default_dataset_name(),
+            shuffle=True,
+            max_samples=100,
+        )
+
+        # Generate some random selections
+        selected_objects = []
+        for sample in dataset.take(num_samples_to_select):
+            detections = sample.ground_truth.detections
+            num_objects = random.randint(
+                1, min(len(detections), max_objects_per_sample_to_select)
+            )
+            for detection in random.sample(detections, num_objects):
+                selected_objects.append(
+                    {
+                        "sample_id": sample.id,
+                        "field": "ground_truth",
+                        "object_id": detection.id,
+                    }
+                )
+
+        selected_view = fous.select_objects(dataset, selected_objects)
+        excluded_view = fous.exclude_objects(dataset, selected_objects)
+
+        total_objects = _count_detections(dataset, "ground_truth")
+        num_selected_objects = len(selected_objects)
+        num_objects_in_selected_view = _count_detections(
+            selected_view, "ground_truth"
+        )
+        num_objects_in_excluded_view = _count_detections(
+            excluded_view, "ground_truth"
+        )
+        num_objects_excluded = total_objects - num_objects_in_excluded_view
+
+        self.assertEqual(num_selected_objects, num_objects_in_selected_view)
+        self.assertEqual(num_selected_objects, num_objects_excluded)
+
+
+def _count_detections(sample_collection, label_field):
+    num_objects = 0
+    for sample in sample_collection:
+        num_objects += len(sample[label_field].detections)
+
+    return num_objects
+
+
+if __name__ == "__main__":
+    fo.config.show_progress_bars = False
+    unittest.main(verbosity=2)