[Datumaro] Reducing nesting of tests (#1875)

* Add `Dataset.from_iterable` constructor
* Simplify creation of `Dataset` objects in common simple cases
* Refactor tests

datumaro/datumaro/components/extractor.py

@@ -72,6 +72,33 @@ class Categories:
 class LabelCategories(Categories):
     Category = namedtuple('Category', ['name', 'parent', 'attributes'])
 
+    @classmethod
+    def from_iterable(cls, iterable):
+        """Generation of LabelCategories from iterable object
+
+        Args:
+            iterable ([type]): This iterable object can be:
+            1)simple str - will generate one Category with str as name
+            2)list of str - will interpreted as list of Category names
+            3)list of positional argumetns - will generate Categories
+            with this arguments
+
+
+        Returns:
+            LabelCategories: LabelCategories object
+        """
+        temp_categories = cls()
+
+        if isinstance(iterable, str):
+            iterable = [[iterable]]
+
+        for category in iterable:
+            if isinstance(category, str):
+                category = [category]
+            temp_categories.add(*category)
+
+        return temp_categories
+
     def __init__(self, items=None, attributes=None):
         super().__init__(attributes=attributes)
 
@@ -482,6 +509,31 @@ class Bbox(_Shape):
 class PointsCategories(Categories):
     Category = namedtuple('Category', ['labels', 'joints'])
 
+    @classmethod
+    def from_iterable(cls, iterable):
+        """Generation of PointsCategories from iterable object
+
+        Args:
+            iterable ([type]): This iterable object can be:
+            1)simple int - will generate one Category with int as label
+            2)list of int - will interpreted as list of Category labels
+            3)list of positional argumetns - will generate Categories
+            with this arguments
+
+        Returns:
+            PointsCategories: PointsCategories object
+        """
+        temp_categories = cls()
+
+        if isinstance(iterable, int):
+            iterable = [[iterable]]
+
+        for category in iterable:
+            if isinstance(category, int):
+                category = [category]
+            temp_categories.add(*category)
+        return temp_categories
+
     def __init__(self, items=None, attributes=None):
         super().__init__(attributes=attributes)
 

datumaro/datumaro/components/project.py

@@ -18,7 +18,8 @@ import sys
 from datumaro.components.config import Config, DEFAULT_FORMAT
 from datumaro.components.config_model import (Model, Source,
     PROJECT_DEFAULT_CONFIG, PROJECT_SCHEMA)
-from datumaro.components.extractor import Extractor
+from datumaro.components.extractor import Extractor, LabelCategories,\
+    AnnotationType
 from datumaro.components.launcher import ModelTransform
 from datumaro.components.dataset_filter import \
     XPathDatasetFilter, XPathAnnotationsFilter
@@ -319,6 +320,35 @@ class Subset(Extractor):
         return self._parent.categories()
 
 class Dataset(Extractor):
+    @classmethod
+    def from_iterable(cls, iterable, categories=None):
+        """Generation of Dataset from iterable object
+
+        Args:
+            iterable: Iterable object contains DatasetItems
+            categories (dict, optional): You can pass dict of categories or
+            you can pass list of names. It'll interpreted as list of names of
+            LabelCategories. Defaults to {}.
+
+        Returns:
+            Dataset: Dataset object
+        """
+
+        if isinstance(categories, list):
+            categories = {AnnotationType.label : LabelCategories.from_iterable(categories)}
+
+        if not categories:
+            categories = {}
+
+        class tmpExtractor(Extractor):
+            def __iter__(self):
+                return iter(iterable)
+
+            def categories(self):
+                return categories
+
+        return cls.from_extractors(tmpExtractor())
+
     @classmethod
     def from_extractors(cls, *sources):
         # merge categories

datumaro/tests/test_coco_format.py

@@ -4,7 +4,7 @@ import os.path as osp
 
 from unittest import TestCase
 
-from datumaro.components.project import Project
+from datumaro.components.project import (Project, Dataset)
 from datumaro.components.extractor import (Extractor, DatasetItem,
     AnnotationType, Label, Mask, Points, Polygon, Bbox, Caption,
     LabelCategories, PointsCategories
@@ -26,9 +26,8 @@ DUMMY_DATASET_DIR = osp.join(osp.dirname(__file__), 'assets', 'coco_dataset')
 
 class CocoImporterTest(TestCase):
     def test_can_import(self):
-        class DstExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+
+        expected_dataset = Dataset.from_iterable([
             DatasetItem(id='000000000001', image=np.ones((10, 5, 3)),
                 subset='val', attributes={'id': 1},
                 annotations=[
@@ -41,17 +40,12 @@ class CocoImporterTest(TestCase):
                         id=2, group=2, attributes={'is_crowd': True}),
                 ]
             ),
-                ])
-
-            def categories(self):
-                label_cat = LabelCategories()
-                label_cat.add('TEST')
-                return { AnnotationType.label: label_cat }
+        ], categories=['TEST',])
 
         dataset = Project.import_from(DUMMY_DATASET_DIR, 'coco') \
             .make_dataset()
 
-        compare_datasets(self, DstExtractor(), dataset)
+        compare_datasets(self, expected_dataset, dataset)
 
     def test_can_detect(self):
         self.assertTrue(CocoImporter.detect(DUMMY_DATASET_DIR))
@@ -71,9 +65,7 @@ class CocoConverterTest(TestCase):
         compare_datasets(self, expected=target_dataset, actual=parsed_dataset)
 
     def test_can_save_and_load_captions(self):
-        class TestExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        expected_dataset = Dataset.from_iterable([
             DatasetItem(id=1, subset='train',
                 annotations=[
                     Caption('hello', id=1, group=1),
@@ -87,23 +79,15 @@ class CocoConverterTest(TestCase):
             DatasetItem(id=3, subset='val',
                 annotations=[
                     Caption('word', id=1, group=1),
-                        ], attributes={'id': 1}
-                    ),
+                ], attributes={'id': 1}),
             ])
 
         with TestDir() as test_dir:
-            self._test_save_and_load(TestExtractor(),
+            self._test_save_and_load(expected_dataset,
                 CocoCaptionsConverter.convert, test_dir)
 
     def test_can_save_and_load_instances(self):
-        label_categories = LabelCategories()
-        for i in range(10):
-            label_categories.add(str(i))
-        categories = { AnnotationType.label: label_categories }
-
-        class TestExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        source_dataset = Dataset.from_iterable([
             DatasetItem(id=1, subset='train', image=np.ones((4, 4, 3)),
                 annotations=[
                     # Bbox + single polygon
@@ -143,14 +127,9 @@ class CocoConverterTest(TestCase):
                         attributes={ 'is_crowd': True },
                         label=4, group=3, id=3),
                 ], attributes={'id': 1}),
-                ])
-
-            def categories(self):
-                return categories
+            ], categories=[str(i) for i in range(10)])
 
-        class DstExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        target_dataset = Dataset.from_iterable([
             DatasetItem(id=1, subset='train', image=np.ones((4, 4, 3)),
                 annotations=[
                     Polygon([0, 1, 2, 1, 2, 3, 0, 3],
@@ -179,26 +158,16 @@ class CocoConverterTest(TestCase):
                         ),
                         attributes={ 'is_crowd': True },
                         label=4, group=3, id=3),
-                        ], attributes={'id': 1}),
-                ])
-
-            def categories(self):
-                return categories
+                ], attributes={'id': 1})
+            ], categories=[str(i) for i in range(10)])
 
         with TestDir() as test_dir:
-            self._test_save_and_load(TestExtractor(),
+            self._test_save_and_load(source_dataset,
                 CocoInstancesConverter.convert, test_dir,
-                target_dataset=DstExtractor())
+                target_dataset=target_dataset)
 
     def test_can_merge_polygons_on_loading(self):
-        label_categories = LabelCategories()
-        for i in range(10):
-            label_categories.add(str(i))
-        categories = { AnnotationType.label: label_categories }
-
-        class SrcExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        source_dataset = Dataset.from_iterable([
             DatasetItem(id=1, image=np.zeros((6, 10, 3)),
                 annotations=[
                     Polygon([0, 0, 4, 0, 4, 4],
@@ -207,14 +176,9 @@ class CocoConverterTest(TestCase):
                         label=3, id=4, group=4),
                 ]
             ),
-                ])
-
-            def categories(self):
-                return categories
+        ], categories=[str(i) for i in range(10)])
 
-        class DstExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        target_dataset = Dataset.from_iterable([
             DatasetItem(id=1, image=np.zeros((6, 10, 3)),
                 annotations=[
                     Mask(np.array([
@@ -231,25 +195,16 @@ class CocoConverterTest(TestCase):
                     attributes={ 'is_crowd': False }),
                 ], attributes={'id': 1}
             ),
-                ])
-
-            def categories(self):
-                return categories
+        ], categories=[str(i) for i in range(10)])
 
         with TestDir() as test_dir:
-            self._test_save_and_load(SrcExtractor(),
+            self._test_save_and_load(source_dataset,
                 CocoInstancesConverter.convert, test_dir,
                 importer_args={'merge_instance_polygons': True},
-                target_dataset=DstExtractor())
+                target_dataset=target_dataset)
 
     def test_can_crop_covered_segments(self):
-        label_categories = LabelCategories()
-        for i in range(10):
-            label_categories.add(str(i))
-
-        class SrcTestExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        source_dataset = Dataset.from_iterable([
             DatasetItem(id=1, image=np.zeros((5, 5, 3)),
                 annotations=[
                     Mask(np.array([
@@ -264,14 +219,9 @@ class CocoConverterTest(TestCase):
                         label=1, id=2, z_order=1),
                 ]
             ),
-                ])
+        ], categories=[str(i) for i in range(10)])
 
-            def categories(self):
-                return { AnnotationType.label: label_categories }
-
-        class DstTestExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        target_dataset = Dataset.from_iterable([
             DatasetItem(id=1, image=np.zeros((5, 5, 3)),
                 annotations=[
                     Mask(np.array([
@@ -289,24 +239,15 @@ class CocoConverterTest(TestCase):
                         attributes={ 'is_crowd': False }),
                 ], attributes={'id': 1}
             ),
-                ])
-
-            def categories(self):
-                return { AnnotationType.label: label_categories }
+        ], categories=[str(i) for i in range(10)])
 
         with TestDir() as test_dir:
-            self._test_save_and_load(SrcTestExtractor(),
+            self._test_save_and_load(source_dataset,
                  partial(CocoInstancesConverter.convert, crop_covered=True),
-                test_dir, target_dataset=DstTestExtractor())
+                 test_dir, target_dataset=target_dataset)
 
     def test_can_convert_polygons_to_mask(self):
-        label_categories = LabelCategories()
-        for i in range(10):
-            label_categories.add(str(i))
-
-        class SrcTestExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        source_dataset = Dataset.from_iterable([
             DatasetItem(id=1, image=np.zeros((6, 10, 3)),
                 annotations=[
                     Polygon([0, 0, 4, 0, 4, 4],
@@ -315,14 +256,9 @@ class CocoConverterTest(TestCase):
                         label=3, id=4, group=4),
                 ]
             ),
-                ])
+        ], categories=[str(i) for i in range(10)])
 
-            def categories(self):
-                return { AnnotationType.label: label_categories }
-
-        class DstTestExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        target_dataset = Dataset.from_iterable([
             DatasetItem(id=1, image=np.zeros((6, 10, 3)),
                 annotations=[
                     Mask(np.array([
@@ -339,24 +275,15 @@ class CocoConverterTest(TestCase):
                         label=3, id=4, group=4),
                 ], attributes={'id': 1}
             ),
-                ])
-
-            def categories(self):
-                return { AnnotationType.label: label_categories }
+        ], categories=[str(i) for i in range(10)])
 
         with TestDir() as test_dir:
-            self._test_save_and_load(SrcTestExtractor(),
+            self._test_save_and_load(source_dataset,
                 partial(CocoInstancesConverter.convert, segmentation_mode='mask'),
-                test_dir, target_dataset=DstTestExtractor())
+                test_dir, target_dataset=target_dataset)
 
     def test_can_convert_masks_to_polygons(self):
-        label_categories = LabelCategories()
-        for i in range(10):
-            label_categories.add(str(i))
-
-        class SrcExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        source_dataset = Dataset.from_iterable([
             DatasetItem(id=1, image=np.zeros((5, 10, 3)),
                 annotations=[
                     Mask(np.array([
@@ -369,14 +296,9 @@ class CocoConverterTest(TestCase):
                         label=3, id=4, group=4),
                 ]
             ),
-                ])
-
-            def categories(self):
-                return { AnnotationType.label: label_categories }
+        ], categories=[str(i) for i in range(10)])
 
-        class DstExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        target_dataset = Dataset.from_iterable([
             DatasetItem(id=1, image=np.zeros((5, 10, 3)),
                 annotations=[
                     Polygon(
@@ -389,20 +311,16 @@ class CocoConverterTest(TestCase):
                         attributes={ 'is_crowd': False }),
                 ], attributes={'id': 1}
             ),
-                ])
-
-            def categories(self):
-                return { AnnotationType.label: label_categories }
+        ], categories=[str(i) for i in range(10)])
 
         with TestDir() as test_dir:
-            self._test_save_and_load(SrcExtractor(),
+            self._test_save_and_load(source_dataset,
                 partial(CocoInstancesConverter.convert, segmentation_mode='polygons'),
-                test_dir, target_dataset=DstExtractor())
+                test_dir,
+                target_dataset=target_dataset)
 
     def test_can_save_and_load_images(self):
-        class TestExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        expected_dataset = Dataset.from_iterable([
             DatasetItem(id=1, subset='train', attributes={'id': 1}),
             DatasetItem(id=2, subset='train', attributes={'id': 2}),
 
@@ -414,47 +332,25 @@ class CocoConverterTest(TestCase):
         ])
 
         with TestDir() as test_dir:
-            self._test_save_and_load(TestExtractor(),
+            self._test_save_and_load(expected_dataset,
                 CocoImageInfoConverter.convert, test_dir)
 
     def test_can_save_and_load_labels(self):
-        class TestExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        expected_dataset = Dataset.from_iterable([
             DatasetItem(id=1, subset='train',
                 annotations=[
                     Label(4, id=1, group=1),
                     Label(9, id=2, group=2),
-                        ], attributes={'id': 1}
-                    ),
-                ])
-
-            def categories(self):
-                label_categories = LabelCategories()
-                for i in range(10):
-                    label_categories.add(str(i))
-                return {
-                    AnnotationType.label: label_categories,
-                }
+                ], attributes={'id': 1}),
+        ], categories=[str(i) for i in range(10)])
 
         with TestDir() as test_dir:
-            self._test_save_and_load(TestExtractor(),
+            self._test_save_and_load(expected_dataset,
                 CocoLabelsConverter.convert, test_dir)
 
     def test_can_save_and_load_keypoints(self):
-        label_categories = LabelCategories()
-        points_categories = PointsCategories()
-        for i in range(10):
-            label_categories.add(str(i))
-            points_categories.add(i, joints=[[0, 1], [1, 2]])
-        categories = {
-            AnnotationType.label: label_categories,
-            AnnotationType.points: points_categories,
-        }
-
-        class TestExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+
+        source_dataset = Dataset.from_iterable([
             DatasetItem(id=1, subset='train', image=np.zeros((5, 5, 3)),
                 annotations=[
                     # Full instance annotations: polygon + keypoints
@@ -475,15 +371,16 @@ class CocoConverterTest(TestCase):
 
                     # Solitary keypoints with no label
                     Points([0, 0, 1, 2, 3, 4], [0, 1, 2], id=5),
-                        ])
-                ])
-
-            def categories(self):
-                return categories
+                ]),
+            ], categories={
+                    AnnotationType.label: LabelCategories.from_iterable(
+                        str(i) for i in range(10)),
+                    AnnotationType.points: PointsCategories.from_iterable(
+                        (i, None, [[0, 1], [1, 2]]) for i in range(10)
+                    ),
+            })
 
-        class DstTestExtractor(TestExtractor):
-            def __iter__(self):
-                return iter([
+        target_dataset = Dataset.from_iterable([
             DatasetItem(id=1, subset='train', image=np.zeros((5, 5, 3)),
                 annotations=[
                     Points([0, 0, 0, 2, 4, 1], [0, 1, 2],
@@ -514,41 +411,41 @@ class CocoConverterTest(TestCase):
                         group=5, id=5,
                         attributes={'is_crowd': False}),
                 ], attributes={'id': 1}),
-                ])
+            ], categories={
+                    AnnotationType.label: LabelCategories.from_iterable(
+                        str(i) for i in range(10)),
+                    AnnotationType.points: PointsCategories.from_iterable(
+                        (i, None, [[0, 1], [1, 2]]) for i in range(10)
+                    ),
+            })
 
         with TestDir() as test_dir:
-            self._test_save_and_load(TestExtractor(),
+            self._test_save_and_load(source_dataset,
                 CocoPersonKeypointsConverter.convert, test_dir,
-                target_dataset=DstTestExtractor())
+                target_dataset=target_dataset)
 
     def test_can_save_dataset_with_no_subsets(self):
-        class TestExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        test_dataset = Dataset.from_iterable([
             DatasetItem(id=1, attributes={'id': 1}),
             DatasetItem(id=2, attributes={'id': 2}),
         ])
 
         with TestDir() as test_dir:
-            self._test_save_and_load(TestExtractor(),
+            self._test_save_and_load(test_dataset,
                 CocoConverter.convert, test_dir)
 
     def test_can_save_dataset_with_image_info(self):
-        class TestExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        expected_dataset = Dataset.from_iterable([
             DatasetItem(id=1, image=Image(path='1.jpg', size=(10, 15)),
                 attributes={'id': 1}),
         ])
 
         with TestDir() as test_dir:
-            self._test_save_and_load(TestExtractor(),
+            self._test_save_and_load(expected_dataset,
                 CocoImageInfoConverter.convert, test_dir)
 
     def test_relative_paths(self):
-        class TestExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        expected_dataset = Dataset.from_iterable([
             DatasetItem(id='1', image=np.ones((4, 2, 3)),
                 attributes={'id': 1}),
             DatasetItem(id='subdir1/1', image=np.ones((2, 6, 3)),
@@ -558,39 +455,27 @@ class CocoConverterTest(TestCase):
         ])
 
         with TestDir() as test_dir:
-            self._test_save_and_load(TestExtractor(),
-                partial(CocoImageInfoConverter.convert, save_images=True),
-                test_dir)
+            self._test_save_and_load(expected_dataset,
+                partial(CocoImageInfoConverter.convert, save_images=True), test_dir)
 
     def test_preserve_coco_ids(self):
-        class TestExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        expected_dataset = Dataset.from_iterable([
             DatasetItem(id='some/name1', image=np.ones((4, 2, 3)),
                 attributes={'id': 40}),
         ])
 
         with TestDir() as test_dir:
-            self._test_save_and_load(TestExtractor(),
-                partial(CocoImageInfoConverter.convert, save_images=True),
-                test_dir)
+            self._test_save_and_load(expected_dataset,
+                partial(CocoImageInfoConverter.convert, save_images=True), test_dir)
 
     def test_annotation_attributes(self):
-        class TestExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        expected_dataset = Dataset.from_iterable([
             DatasetItem(id=1, image=np.ones((4, 2, 3)), annotations=[
                 Polygon([0, 0, 4, 0, 4, 4], label=5, group=1, id=1,
                     attributes={'is_crowd': False, 'x': 5, 'y': 'abc'}),
             ], attributes={'id': 1})
-                ])
-
-            def categories(self):
-                label_categories = LabelCategories()
-                for i in range(10):
-                    label_categories.add(str(i))
-                return { AnnotationType.label: label_categories, }
+        ], categories=[str(i) for i in range(10)])
 
         with TestDir() as test_dir:
-            self._test_save_and_load(TestExtractor(),
+            self._test_save_and_load(expected_dataset,
                 CocoConverter.convert, test_dir)

datumaro/tests/test_cvat_format.py

@@ -3,7 +3,7 @@ import numpy as np
 import os.path as osp
 
 from unittest import TestCase
-
+from datumaro.components.project import Dataset
 from datumaro.components.extractor import (Extractor, DatasetItem,
     AnnotationType, Points, Polygon, PolyLine, Bbox, Label,
     LabelCategories,
@@ -28,9 +28,7 @@ class CvatImporterTest(TestCase):
         self.assertTrue(CvatImporter.detect(DUMMY_VIDEO_DATASET_DIR))
 
     def test_can_load_image(self):
-        class DstExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        expected_dataset = Dataset.from_iterable([
             DatasetItem(id='img0', subset='train',
                 image=np.ones((8, 8, 3)),
                 annotations=[
@@ -50,22 +48,19 @@ class CvatImporterTest(TestCase):
                     Points([1, 2, 3, 4, 5, 6], label=1, z_order=2,
                         attributes={'occluded': False}),
                 ], attributes={'frame': 1}),
+        ], categories={
+            AnnotationType.label: LabelCategories.from_iterable([
+                ['label1', '', {'a1', 'a2'}],
+                ['label2'],
             ])
-
-            def categories(self):
-                label_categories = LabelCategories()
-                label_categories.add('label1', attributes={'a1', 'a2'})
-                label_categories.add('label2')
-                return { AnnotationType.label: label_categories }
+        })
 
         parsed_dataset = CvatImporter()(DUMMY_IMAGE_DATASET_DIR).make_dataset()
 
-        compare_datasets(self, DstExtractor(), parsed_dataset)
+        compare_datasets(self, expected_dataset, parsed_dataset)
 
     def test_can_load_video(self):
-        class DstExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        expected_dataset = Dataset.from_iterable([
             DatasetItem(id='frame_000010', subset='annotations',
                 image=np.ones((20, 25, 3)),
                 annotations=[
@@ -131,18 +126,17 @@ class CvatImporterTest(TestCase):
                             'track_id': 2,
                         }),
                 ], attributes={'frame': 16}),
-                ])
-
-            def categories(self):
-                label_categories = LabelCategories()
-                label_categories.add('klhg', attributes={'hgl'})
-                label_categories.add('z U k')
-                label_categories.add('II')
-                return { AnnotationType.label: label_categories }
+        ], categories={
+            AnnotationType.label: LabelCategories.from_iterable([
+                ['klhg', '', {'hgl'}],
+                ['z U k'],
+                ['II']
+            ]),
+        })
 
         parsed_dataset = CvatImporter()(DUMMY_VIDEO_DATASET_DIR).make_dataset()
 
-        compare_datasets(self, DstExtractor(), parsed_dataset)
+        compare_datasets(self, expected_dataset, parsed_dataset)
 
 class CvatConverterTest(TestCase):
     def _test_save_and_load(self, source_dataset, converter, test_dir,
@@ -165,9 +159,7 @@ class CvatConverterTest(TestCase):
         label_categories.items[2].attributes.update(['a1', 'a2'])
         label_categories.attributes.update(['occluded'])
 
-        class SrcExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        source_dataset = Dataset.from_iterable([
             DatasetItem(id=0, subset='s1', image=np.zeros((5, 10, 3)),
                 annotations=[
                     Polygon([0, 0, 4, 0, 4, 4],
@@ -201,14 +193,11 @@ class CvatConverterTest(TestCase):
 
             DatasetItem(id=3, subset='s3', image=Image(
                 path='3.jpg', size=(2, 4))),
-                ])
-
-            def categories(self):
-                return { AnnotationType.label: label_categories }
+        ], categories={
+            AnnotationType.label: label_categories,
+        })
 
-        class DstExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        target_dataset = Dataset.from_iterable([
             DatasetItem(id=0, subset='s1', image=np.zeros((5, 10, 3)),
                 annotations=[
                     Polygon([0, 0, 4, 0, 4, 4],
@@ -244,58 +233,46 @@ class CvatConverterTest(TestCase):
             DatasetItem(id=3, subset='s3', image=Image(
                     path='3.jpg', size=(2, 4)),
                 attributes={'frame': 0}),
-                ])
-
-            def categories(self):
-                return { AnnotationType.label: label_categories }
+        ], categories={
+            AnnotationType.label: label_categories,
+        })
 
         with TestDir() as test_dir:
-            self._test_save_and_load(SrcExtractor(),
+            self._test_save_and_load(source_dataset,
                 partial(CvatConverter.convert, save_images=True), test_dir,
-                target_dataset=DstExtractor())
+                target_dataset=target_dataset)
 
     def test_relative_paths(self):
-        class SrcExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        source_dataset = Dataset.from_iterable([
             DatasetItem(id='1', image=np.ones((4, 2, 3))),
             DatasetItem(id='subdir1/1', image=np.ones((2, 6, 3))),
             DatasetItem(id='subdir2/1', image=np.ones((5, 4, 3))),
-                ])
-
-            def categories(self):
-                return { AnnotationType.label: LabelCategories() }
+        ], categories={ AnnotationType.label: LabelCategories() })
 
-        class DstExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        target_dataset = Dataset.from_iterable([
             DatasetItem(id='1', image=np.ones((4, 2, 3)),
                 attributes={'frame': 0}),
             DatasetItem(id='subdir1/1', image=np.ones((2, 6, 3)),
                 attributes={'frame': 1}),
             DatasetItem(id='subdir2/1', image=np.ones((5, 4, 3)),
                 attributes={'frame': 2}),
-                ])
-
-            def categories(self):
-                return { AnnotationType.label: LabelCategories() }
+        ], categories={
+            AnnotationType.label: LabelCategories()
+        })
 
         with TestDir() as test_dir:
-            self._test_save_and_load(SrcExtractor(),
+            self._test_save_and_load(source_dataset,
                 partial(CvatConverter.convert, save_images=True), test_dir,
-                target_dataset=DstExtractor())
+                target_dataset=target_dataset)
 
     def test_preserve_frame_ids(self):
-        class TestExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        expected_dataset = Dataset.from_iterable([
             DatasetItem(id='some/name1', image=np.ones((4, 2, 3)),
                 attributes={'frame': 40}),
-                ])
-
-            def categories(self):
-                return { AnnotationType.label: LabelCategories() }
+        ], categories={
+            AnnotationType.label: LabelCategories()
+        })
 
         with TestDir() as test_dir:
-            self._test_save_and_load(TestExtractor(),
+            self._test_save_and_load(expected_dataset,
                 CvatConverter.convert, test_dir)

datumaro/tests/test_datumaro_format.py

@@ -2,7 +2,7 @@ from functools import partial
 import numpy as np
 
 from unittest import TestCase
-
+from datumaro.components.project import Dataset
 from datumaro.components.project import Project
 from datumaro.components.extractor import (Extractor, DatasetItem,
     AnnotationType, Label, Mask, Points, Polygon,
@@ -32,9 +32,18 @@ class DatumaroConverterTest(TestCase):
         compare_datasets_strict(self,
             expected=target_dataset, actual=parsed_dataset)
 
-    class TestExtractor(Extractor):
-        def __iter__(self):
-            return iter([
+    label_categories = LabelCategories()
+    for i in range(5):
+        label_categories.add('cat' + str(i))
+
+    mask_categories = MaskCategories(
+        generate_colormap(len(label_categories.items)))
+
+    points_categories = PointsCategories()
+    for index, _ in enumerate(label_categories.items):
+        points_categories.add(index, ['cat1', 'cat2'], joints=[[0, 1]])
+
+    test_dataset = Dataset.from_iterable([
         DatasetItem(id=100, subset='train', image=np.ones((10, 6, 3)),
             annotations=[
                 Caption('hello', id=1),
@@ -68,46 +77,30 @@ class DatumaroConverterTest(TestCase):
 
         DatasetItem(id=42),
         DatasetItem(id=43, image=Image(path='1/b/c.qq', size=(2, 4))),
-            ])
-
-        def categories(self):
-            label_categories = LabelCategories()
-            for i in range(5):
-                label_categories.add('cat' + str(i))
-
-            mask_categories = MaskCategories(
-                generate_colormap(len(label_categories.items)))
-
-            points_categories = PointsCategories()
-            for index, _ in enumerate(label_categories.items):
-                points_categories.add(index, ['cat1', 'cat2'], joints=[[0, 1]])
-
-            return {
+    ], categories={
         AnnotationType.label: label_categories,
         AnnotationType.mask: mask_categories,
         AnnotationType.points: points_categories,
-            }
+    })
 
     def test_can_save_and_load(self):
         with TestDir() as test_dir:
-            self._test_save_and_load(self.TestExtractor(),
+            self._test_save_and_load(self.test_dataset,
                 partial(DatumaroConverter.convert, save_images=True), test_dir)
 
     def test_can_detect(self):
         with TestDir() as test_dir:
-            DatumaroConverter.convert(self.TestExtractor(), save_dir=test_dir)
+            DatumaroConverter.convert(self.test_dataset, save_dir=test_dir)
 
             self.assertTrue(DatumaroImporter.detect(test_dir))
 
     def test_relative_paths(self):
-        class TestExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        test_dataset = Dataset.from_iterable([
             DatasetItem(id='1', image=np.ones((4, 2, 3))),
             DatasetItem(id='subdir1/1', image=np.ones((2, 6, 3))),
             DatasetItem(id='subdir2/1', image=np.ones((5, 4, 3))),
         ])
 
         with TestDir() as test_dir:
-            self._test_save_and_load(TestExtractor(),
+            self._test_save_and_load(test_dataset,
                 partial(DatumaroConverter.convert, save_images=True), test_dir)

datumaro/tests/test_labelme_format.py

@@ -3,7 +3,7 @@ import numpy as np
 import os.path as osp
 
 from unittest import TestCase
-
+from datumaro.components.project import Dataset
 from datumaro.components.extractor import (Extractor, DatasetItem,
     AnnotationType, Bbox, Mask, Polygon, LabelCategories
 )
@@ -29,9 +29,7 @@ class LabelMeConverterTest(TestCase):
         compare_datasets(self, expected=target_dataset, actual=parsed_dataset)
 
     def test_can_save_and_load(self):
-        class SrcExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        source_dataset = Dataset.from_iterable([
             DatasetItem(id=1, subset='train',
                 image=np.ones((16, 16, 3)),
                 annotations=[
@@ -50,19 +48,12 @@ class LabelMeConverterTest(TestCase):
                     ),
                 ]
             ),
-                ])
-
-            def categories(self):
-                label_cat = LabelCategories()
-                for label in range(10):
-                    label_cat.add('label_' + str(label))
-                return {
-                    AnnotationType.label: label_cat,
-                }
+        ], categories={
+            AnnotationType.label: LabelCategories.from_iterable(
+                'label_' + str(label) for label in range(10)),
+        })
 
-        class DstExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        target_dataset = Dataset.from_iterable([
             DatasetItem(id=1, subset='train',
                 image=np.ones((16, 16, 3)),
                 annotations=[
@@ -94,36 +85,28 @@ class LabelMeConverterTest(TestCase):
                     ),
                 ]
             ),
-                ])
-
-            def categories(self):
-                label_cat = LabelCategories()
-                label_cat.add('label_2')
-                label_cat.add('label_3')
-                return {
-                    AnnotationType.label: label_cat,
-                }
+        ], categories={
+            AnnotationType.label: LabelCategories.from_iterable([
+                'label_2', 'label_3']),
+        })
 
         with TestDir() as test_dir:
-            self._test_save_and_load(SrcExtractor(),
+            self._test_save_and_load(
+                source_dataset,
                 partial(LabelMeConverter.convert, save_images=True),
-                test_dir, target_dataset=DstExtractor())
+                test_dir, target_dataset=target_dataset)
 
     def test_cant_save_dataset_with_relative_paths(self):
-        class SrcExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        expected_dataset = Dataset.from_iterable([
             DatasetItem(id='dir/1', image=np.ones((2, 6, 3))),
-                ])
-
-            def categories(self):
-                return { AnnotationType.label: LabelCategories() }
+        ], categories={
+            AnnotationType.label: LabelCategories(),
+        })
 
         with self.assertRaisesRegex(Exception, r'only supports flat'):
             with TestDir() as test_dir:
-                self._test_save_and_load(SrcExtractor(),
-                    partial(LabelMeConverter.convert, save_images=True),
-                    test_dir)
+                self._test_save_and_load(expected_dataset,
+                    LabelMeConverter.convert, test_dir)
 
 
 DUMMY_DATASET_DIR = osp.join(osp.dirname(__file__), 'assets', 'labelme_dataset')
@@ -133,8 +116,6 @@ class LabelMeImporterTest(TestCase):
         self.assertTrue(LabelMeImporter.detect(DUMMY_DATASET_DIR))
 
     def test_can_import(self):
-        class DstExtractor(Extractor):
-            def __iter__(self):
         img1 = np.ones((77, 102, 3)) * 255
         img1[6:32, 7:41] = 0
 
@@ -157,7 +138,7 @@ class LabelMeImporterTest(TestCase):
             [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
         ]
 
-                return iter([
+        target_dataset = Dataset.from_iterable([
             DatasetItem(id='img1', image=img1,
                 annotations=[
                     Polygon([43, 34, 45, 34, 45, 37, 43, 37],
@@ -213,21 +194,13 @@ class LabelMeImporterTest(TestCase):
                     ),
                 ]
             ),
-                ])
-
-            def categories(self):
-                label_cat = LabelCategories()
-                label_cat.add('window')
-                label_cat.add('license plate')
-                label_cat.add('o1')
-                label_cat.add('q1')
-                label_cat.add('b1')
-                label_cat.add('m1')
-                label_cat.add('hg')
-                return {
-                    AnnotationType.label: label_cat,
-                }
+        ], categories={
+            AnnotationType.label: LabelCategories.from_iterable([
+                'window', 'license plate', 'o1',
+                'q1', 'b1', 'm1', 'hg',
+            ]),
+        })
 
         parsed = Project.import_from(DUMMY_DATASET_DIR, 'label_me') \
             .make_dataset()
-        compare_datasets(self, expected=DstExtractor(), actual=parsed)
+        compare_datasets(self, expected=target_dataset, actual=parsed)

datumaro/tests/test_mot_format.py

@@ -3,7 +3,7 @@ import numpy as np
 import os.path as osp
 
 from unittest import TestCase
-
+from datumaro.components.project import Dataset
 from datumaro.components.extractor import (Extractor, DatasetItem,
     AnnotationType, Bbox, LabelCategories
 )
@@ -28,9 +28,7 @@ class MotConverterTest(TestCase):
         compare_datasets(self, expected=target_dataset, actual=parsed_dataset)
 
     def test_can_save_bboxes(self):
-        class SrcExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        source_dataset = Dataset.from_iterable([
             DatasetItem(id=1, subset='train',
                 image=np.ones((16, 16, 3)),
                 annotations=[
@@ -56,19 +54,12 @@ class MotConverterTest(TestCase):
             DatasetItem(id=3, subset='test',
                 image=np.ones((5, 4, 3)) * 3,
             ),
-                ])
-
-            def categories(self):
-                label_cat = LabelCategories()
-                for label in range(10):
-                    label_cat.add('label_' + str(label))
-                return {
-                    AnnotationType.label: label_cat,
-                }
-
-        class DstExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        ], categories={
+            AnnotationType.label: LabelCategories.from_iterable(
+                'label_' + str(label) for label in range(10)),
+        })
+
+        target_dataset = Dataset.from_iterable([
             DatasetItem(id=1,
                 image=np.ones((16, 16, 3)),
                 annotations=[
@@ -104,20 +95,16 @@ class MotConverterTest(TestCase):
             DatasetItem(id=3,
                 image=np.ones((5, 4, 3)) * 3,
             ),
-                ])
-
-            def categories(self):
-                label_cat = LabelCategories()
-                for label in range(10):
-                    label_cat.add('label_' + str(label))
-                return {
-                    AnnotationType.label: label_cat,
-                }
+        ], categories={
+            AnnotationType.label: LabelCategories.from_iterable(
+                'label_' + str(label) for label in range(10)),
+        })
 
         with TestDir() as test_dir:
-            self._test_save_and_load(SrcExtractor(),
+            self._test_save_and_load(
+                source_dataset,
                 partial(MotSeqGtConverter.convert, save_images=True),
-                test_dir, target_dataset=DstExtractor())
+                test_dir, target_dataset=target_dataset)
 
 
 DUMMY_DATASET_DIR = osp.join(osp.dirname(__file__), 'assets', 'mot_dataset')
@@ -127,9 +114,7 @@ class MotImporterTest(TestCase):
         self.assertTrue(MotSeqImporter.detect(DUMMY_DATASET_DIR))
 
     def test_can_import(self):
-        class DstExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        expected_dataset = Dataset.from_iterable([
             DatasetItem(id=1,
                 image=np.ones((16, 16, 3)),
                 annotations=[
@@ -140,17 +125,12 @@ class MotImporterTest(TestCase):
                     }),
                 ]
             ),
-                ])
-
-            def categories(self):
-                label_cat = LabelCategories()
-                for label in range(10):
-                    label_cat.add('label_' + str(label))
-                return {
-                    AnnotationType.label: label_cat,
-                }
+        ], categories={
+            AnnotationType.label: LabelCategories.from_iterable(
+                'label_' + str(label) for label in range(10)),
+        })
 
         dataset = Project.import_from(DUMMY_DATASET_DIR, 'mot_seq') \
             .make_dataset()
 
-        compare_datasets(self, DstExtractor(), dataset)
+        compare_datasets(self, expected_dataset, dataset)

datumaro/tests/test_tfrecord_format.py

@@ -3,7 +3,7 @@ import numpy as np
 import os.path as osp
 
 from unittest import TestCase, skipIf
-
+from datumaro.components.project import Dataset
 from datumaro.components.extractor import (Extractor, DatasetItem,
     AnnotationType, Bbox, Mask, LabelCategories
 )
@@ -48,9 +48,7 @@ class TfrecordConverterTest(TestCase):
         compare_datasets(self, expected=target_dataset, actual=parsed_dataset)
 
     def test_can_save_bboxes(self):
-        class TestExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        test_dataset = Dataset.from_iterable([
             DatasetItem(id=1, subset='train',
                 image=np.ones((16, 16, 3)),
                 annotations=[
@@ -59,25 +57,19 @@ class TfrecordConverterTest(TestCase):
                     Bbox(2, 4, 4, 4),
                 ], attributes={'source_id': ''}
             ),
-                ])
-
-            def categories(self):
-                label_cat = LabelCategories()
-                for label in range(10):
-                    label_cat.add('label_' + str(label))
-                return {
-                    AnnotationType.label: label_cat,
-                }
+        ], categories={
+            AnnotationType.label: LabelCategories.from_iterable(
+                'label_' + str(label) for label in range(10)),
+        })
 
         with TestDir() as test_dir:
-            self._test_save_and_load(TestExtractor(),
+            self._test_save_and_load(
+                test_dataset,
                 partial(TfDetectionApiConverter.convert, save_images=True),
                 test_dir)
 
     def test_can_save_masks(self):
-        class TestExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        test_dataset = Dataset.from_iterable([
             DatasetItem(id=1, subset='train', image=np.ones((4, 5, 3)),
                 annotations=[
                     Mask(image=np.array([
@@ -89,25 +81,19 @@ class TfrecordConverterTest(TestCase):
                 ],
                 attributes={'source_id': ''}
             ),
-                ])
-
-            def categories(self):
-                label_cat = LabelCategories()
-                for label in range(10):
-                    label_cat.add('label_' + str(label))
-                return {
-                    AnnotationType.label: label_cat,
-                }
+        ], categories={
+            AnnotationType.label: LabelCategories.from_iterable(
+                'label_' + str(label) for label in range(10)),
+        })
 
         with TestDir() as test_dir:
-            self._test_save_and_load(TestExtractor(),
+            self._test_save_and_load(
+                test_dataset,
                 partial(TfDetectionApiConverter.convert, save_masks=True),
                 test_dir)
 
     def test_can_save_dataset_with_no_subsets(self):
-        class TestExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        test_dataset = Dataset.from_iterable([
             DatasetItem(id=1,
                 image=np.ones((16, 16, 3)),
                 annotations=[
@@ -129,36 +115,29 @@ class TfrecordConverterTest(TestCase):
                 image=np.ones((8, 4, 3)) * 3,
                 attributes={'source_id': ''}
             ),
-                ])
-
-            def categories(self):
-                label_cat = LabelCategories()
-                for label in range(10):
-                    label_cat.add('label_' + str(label))
-                return {
-                    AnnotationType.label: label_cat,
-                }
+        ], categories={
+            AnnotationType.label: LabelCategories.from_iterable(
+                'label_' + str(label) for label in range(10)),
+        })
 
         with TestDir() as test_dir:
-            self._test_save_and_load(TestExtractor(),
+            self._test_save_and_load(
+                test_dataset,
                 partial(TfDetectionApiConverter.convert, save_images=True),
                 test_dir)
 
     def test_can_save_dataset_with_image_info(self):
-        class TestExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        test_dataset = Dataset.from_iterable([
             DatasetItem(id='1/q.e',
                 image=Image(path='1/q.e', size=(10, 15)),
                 attributes={'source_id': ''}
             )
-                ])
-
-            def categories(self):
-                return { AnnotationType.label: LabelCategories() }
+        ], categories={
+            AnnotationType.label: LabelCategories(),
+        })
 
         with TestDir() as test_dir:
-            self._test_save_and_load(TestExtractor(),
+            self._test_save_and_load(test_dataset,
                 TfDetectionApiConverter.convert, test_dir)
 
     def test_labelmap_parsing(self):
@@ -197,9 +176,7 @@ class TfrecordImporterTest(TestCase):
         self.assertTrue(TfDetectionApiImporter.detect(DUMMY_DATASET_DIR))
 
     def test_can_import(self):
-        class DstExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        target_dataset = Dataset.from_iterable([
             DatasetItem(id=1, subset='train',
                 image=np.ones((16, 16, 3)),
                 annotations=[
@@ -222,17 +199,12 @@ class TfrecordImporterTest(TestCase):
                 image=np.ones((5, 4, 3)) * 3,
                 attributes={'source_id': '3'}
             ),
-                ])
-
-            def categories(self):
-                label_cat = LabelCategories()
-                for label in range(10):
-                    label_cat.add('label_' + str(label))
-                return {
-                    AnnotationType.label: label_cat,
-                }
+        ], categories={
+            AnnotationType.label: LabelCategories.from_iterable(
+                'label_' + str(label) for label in range(10)),
+        })
 
         dataset = Project.import_from(DUMMY_DATASET_DIR, 'tf_detection_api') \
             .make_dataset()
 
-        compare_datasets(self, DstExtractor(), dataset)
+        compare_datasets(self, target_dataset, dataset)

datumaro/tests/test_transforms.py

@@ -2,7 +2,7 @@ import logging as log
 import numpy as np
 
 from unittest import TestCase
-
+from datumaro.components.project import Dataset
 from datumaro.components.extractor import (Extractor, DatasetItem,
     Mask, Polygon, PolyLine, Points, Bbox, Label,
     LabelCategories, MaskCategories, AnnotationType
@@ -67,9 +67,7 @@ class TransformsTest(TestCase):
         compare_datasets(self, DstExtractor(), actual)
 
     def test_mask_to_polygons_small_polygons_message(self):
-        class SrcExtractor(Extractor):
-            def __iter__(self):
-                items = [
+        source_dataset = Dataset.from_iterable([
             DatasetItem(id=1, image=np.zeros((5, 10, 3)),
                 annotations=[
                     Mask(np.array([
@@ -80,23 +78,19 @@ class TransformsTest(TestCase):
                     ),
                 ]
             ),
-                ]
-                return iter(items)
+        ])
 
-        class DstExtractor(Extractor):
-            def __iter__(self):
-                return iter([ DatasetItem(id=1, image=np.zeros((5, 10, 3))), ])
+        target_dataset = Dataset.from_iterable([
+            DatasetItem(id=1, image=np.zeros((5, 10, 3))), ])
 
         with self.assertLogs(level=log.DEBUG) as logs:
-            actual = transforms.MasksToPolygons(SrcExtractor())
+            actual = transforms.MasksToPolygons(source_dataset)
 
-            compare_datasets(self, DstExtractor(), actual)
+            compare_datasets(self, target_dataset, actual)
             self.assertRegex('\n'.join(logs.output), 'too small polygons')
 
     def test_polygons_to_masks(self):
-        class SrcExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        source_dataset = Dataset.from_iterable([
             DatasetItem(id=1, image=np.zeros((5, 10, 3)),
                 annotations=[
                     Polygon([0, 0, 4, 0, 4, 4]),
@@ -105,9 +99,7 @@ class TransformsTest(TestCase):
             ),
         ])
 
-        class DstExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        target_dataset = Dataset.from_iterable([
             DatasetItem(id=1, image=np.zeros((5, 10, 3)),
                 annotations=[
                     Mask(np.array([
@@ -130,13 +122,11 @@ class TransformsTest(TestCase):
             ),
         ])
 
-        actual = transforms.PolygonsToMasks(SrcExtractor())
-        compare_datasets(self, DstExtractor(), actual)
+        actual = transforms.PolygonsToMasks(source_dataset)
+        compare_datasets(self, target_dataset, actual)
 
     def test_crop_covered_segments(self):
-        class SrcExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        source_dataset = Dataset.from_iterable([
             DatasetItem(id=1, image=np.zeros((5, 5, 3)),
                 annotations=[
                     # The mask is partially covered by the polygon
@@ -154,9 +144,7 @@ class TransformsTest(TestCase):
             ),
         ])
 
-        class DstExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        target_dataset = Dataset.from_iterable([
             DatasetItem(id=1, image=np.zeros((5, 5, 3)),
                 annotations=[
                     Mask(np.array([
@@ -173,13 +161,11 @@ class TransformsTest(TestCase):
             ),
         ])
 
-        actual = transforms.CropCoveredSegments(SrcExtractor())
-        compare_datasets(self, DstExtractor(), actual)
+        actual = transforms.CropCoveredSegments(source_dataset)
+        compare_datasets(self, target_dataset, actual)
 
     def test_merge_instance_segments(self):
-        class SrcExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        source_dataset = Dataset.from_iterable([
             DatasetItem(id=1, image=np.zeros((5, 5, 3)),
                 annotations=[
                     Mask(np.array([
@@ -198,9 +184,7 @@ class TransformsTest(TestCase):
             ),
         ])
 
-        class DstExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        target_dataset = Dataset.from_iterable([
             DatasetItem(id=1, image=np.zeros((5, 5, 3)),
                 annotations=[
                     Mask(np.array([
@@ -223,35 +207,29 @@ class TransformsTest(TestCase):
             ),
         ])
 
-        actual = transforms.MergeInstanceSegments(SrcExtractor(),
+        actual = transforms.MergeInstanceSegments(source_dataset,
             include_polygons=True)
-        compare_datasets(self, DstExtractor(), actual)
+        compare_datasets(self, target_dataset, actual)
 
     def test_map_subsets(self):
-        class SrcExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        source_dataset = Dataset.from_iterable([
             DatasetItem(id=1, subset='a'),
             DatasetItem(id=2, subset='b'),
             DatasetItem(id=3, subset='c'),
         ])
 
-        class DstExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        target_dataset = Dataset.from_iterable([
             DatasetItem(id=1, subset=''),
             DatasetItem(id=2, subset='a'),
             DatasetItem(id=3, subset='c'),
         ])
 
-        actual = transforms.MapSubsets(SrcExtractor(),
+        actual = transforms.MapSubsets(source_dataset,
             { 'a': '', 'b': 'a' })
-        compare_datasets(self, DstExtractor(), actual)
+        compare_datasets(self, target_dataset, actual)
 
     def test_shapes_to_boxes(self):
-        class SrcExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        source_dataset = Dataset.from_iterable([
             DatasetItem(id=1, image=np.zeros((5, 5, 3)),
                 annotations=[
                     Mask(np.array([
@@ -268,9 +246,7 @@ class TransformsTest(TestCase):
             ),
         ])
 
-        class DstExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        target_dataset = Dataset.from_iterable([
             DatasetItem(id=1, image=np.zeros((5, 5, 3)),
                 annotations=[
                     Bbox(0, 0, 4, 4, id=1),
@@ -281,31 +257,24 @@ class TransformsTest(TestCase):
             ),
         ])
 
-        actual = transforms.ShapesToBoxes(SrcExtractor())
-        compare_datasets(self, DstExtractor(), actual)
+        actual = transforms.ShapesToBoxes(source_dataset)
+        compare_datasets(self, target_dataset, actual)
 
     def test_id_from_image(self):
-        class SrcExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        source_dataset = Dataset.from_iterable([
             DatasetItem(id=1, image='path.jpg'),
             DatasetItem(id=2),
         ])
-
-        class DstExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        target_dataset = Dataset.from_iterable([
             DatasetItem(id='path', image='path.jpg'),
             DatasetItem(id=2),
         ])
 
-        actual = transforms.IdFromImageName(SrcExtractor())
-        compare_datasets(self, DstExtractor(), actual)
+        actual = transforms.IdFromImageName(source_dataset)
+        compare_datasets(self, target_dataset, actual)
 
     def test_boxes_to_masks(self):
-        class SrcExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        source_dataset = Dataset.from_iterable([
             DatasetItem(id=1, image=np.zeros((5, 5, 3)),
                 annotations=[
                     Bbox(0, 0, 3, 3, z_order=1),
@@ -315,9 +284,7 @@ class TransformsTest(TestCase):
             ),
         ])
 
-        class DstExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        target_dataset = Dataset.from_iterable([
             DatasetItem(id=1, image=np.zeros((5, 5, 3)),
                 annotations=[
                     Mask(np.array([
@@ -348,13 +315,11 @@ class TransformsTest(TestCase):
             ),
         ])
 
-        actual = transforms.BoxesToMasks(SrcExtractor())
-        compare_datasets(self, DstExtractor(), actual)
+        actual = transforms.BoxesToMasks(source_dataset)
+        compare_datasets(self, target_dataset, actual)
 
     def test_random_split(self):
-        class SrcExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        source_dataset = Dataset.from_iterable([
             DatasetItem(id=1, subset="a"),
             DatasetItem(id=2, subset="a"),
             DatasetItem(id=3, subset="b"),
@@ -364,7 +329,7 @@ class TransformsTest(TestCase):
             DatasetItem(id=7, subset=""),
         ])
 
-        actual = transforms.RandomSplit(SrcExtractor(), splits=[
+        actual = transforms.RandomSplit(source_dataset, splits=[
             ('train', 4.0 / 7.0),
             ('test', 3.0 / 7.0),
         ])
@@ -373,21 +338,19 @@ class TransformsTest(TestCase):
         self.assertEqual(3, len(actual.get_subset('test')))
 
     def test_random_split_gives_error_on_wrong_ratios(self):
-        class SrcExtractor(Extractor):
-            def __iter__(self):
-                return iter([DatasetItem(id=1)])
+        source_dataset = Dataset.from_iterable([DatasetItem(id=1)])
 
         with self.assertRaises(Exception):
-            transforms.RandomSplit(SrcExtractor(), splits=[
+            transforms.RandomSplit(source_dataset, splits=[
                 ('train', 0.5),
                 ('test', 0.7),
             ])
 
         with self.assertRaises(Exception):
-            transforms.RandomSplit(SrcExtractor(), splits=[])
+            transforms.RandomSplit(source_dataset, splits=[])
 
         with self.assertRaises(Exception):
-            transforms.RandomSplit(SrcExtractor(), splits=[
+            transforms.RandomSplit(source_dataset, splits=[
                 ('train', -0.5),
                 ('test', 1.5),
             ])
@@ -462,24 +425,19 @@ class TransformsTest(TestCase):
         compare_datasets(self, DstExtractor(), actual)
 
     def test_remap_labels_delete_unspecified(self):
-        class SrcExtractor(Extractor):
-            def __iter__(self):
-                return iter([ DatasetItem(id=1, annotations=[ Label(0) ]) ])
-
-            def categories(self):
-                label_cat = LabelCategories()
-                label_cat.add('label0')
-
-                return { AnnotationType.label: label_cat }
-
-        class DstExtractor(Extractor):
-            def __iter__(self):
-                return iter([ DatasetItem(id=1, annotations=[]) ])
+        source_dataset = Dataset.from_iterable([
+            DatasetItem(id=1, annotations=[ Label(0) ])
+        ], categories={
+            AnnotationType.label: LabelCategories.from_iterable('label0'),
+        })
 
-            def categories(self):
-                return { AnnotationType.label: LabelCategories() }
+        target_dataset = Dataset.from_iterable([
+                DatasetItem(id=1, annotations=[]),
+            ], categories={
+                AnnotationType.label: LabelCategories(),
+            })
 
-        actual = transforms.RemapLabels(SrcExtractor(),
+        actual = transforms.RemapLabels(source_dataset,
             mapping={}, default='delete')
 
-        compare_datasets(self, DstExtractor(), actual)
+        compare_datasets(self, target_dataset, actual)

datumaro/tests/test_yolo_format.py

@@ -6,7 +6,7 @@ from unittest import TestCase
 from datumaro.components.extractor import (Extractor, DatasetItem,
     AnnotationType, Bbox, LabelCategories,
 )
-from datumaro.components.project import Project
+from datumaro.components.project import Project, Dataset
 from datumaro.plugins.yolo_format.importer import YoloImporter
 from datumaro.plugins.yolo_format.converter import YoloConverter
 from datumaro.util.image import Image, save_image
@@ -15,9 +15,7 @@ from datumaro.util.test_utils import TestDir, compare_datasets
 
 class YoloFormatTest(TestCase):
     def test_can_save_and_load(self):
-        class TestExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        source_dataset = Dataset.from_iterable([
             DatasetItem(id=1, subset='train', image=np.ones((8, 8, 3)),
                 annotations=[
                     Bbox(0, 2, 4, 2, label=2),
@@ -37,18 +35,12 @@ class YoloFormatTest(TestCase):
                     Bbox(0, 2, 4, 2, label=6),
                     Bbox(0, 7, 3, 2, label=7),
                 ]),
-                ])
-
-            def categories(self):
-                label_categories = LabelCategories()
-                for i in range(10):
-                    label_categories.add('label_' + str(i))
-                return {
-                    AnnotationType.label: label_categories,
-                }
+        ], categories={
+            AnnotationType.label: LabelCategories.from_iterable(
+                'label_' + str(i) for i in range(10)),
+        })
 
         with TestDir() as test_dir:
-            source_dataset = TestExtractor()
 
             YoloConverter.convert(source_dataset, test_dir, save_images=True)
             parsed_dataset = YoloImporter()(test_dir).make_dataset()
@@ -56,27 +48,19 @@ class YoloFormatTest(TestCase):
             compare_datasets(self, source_dataset, parsed_dataset)
 
     def test_can_save_dataset_with_image_info(self):
-        class TestExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        source_dataset = Dataset.from_iterable([
             DatasetItem(id=1, subset='train',
                 image=Image(path='1.jpg', size=(10, 15)),
                 annotations=[
                     Bbox(0, 2, 4, 2, label=2),
                     Bbox(3, 3, 2, 3, label=4),
                 ]),
-                ])
-
-            def categories(self):
-                label_categories = LabelCategories()
-                for i in range(10):
-                    label_categories.add('label_' + str(i))
-                return {
-                    AnnotationType.label: label_categories,
-                }
+        ], categories={
+            AnnotationType.label: LabelCategories.from_iterable(
+                'label_' + str(i) for i in range(10)),
+        })
 
         with TestDir() as test_dir:
-            source_dataset = TestExtractor()
 
             YoloConverter.convert(source_dataset, test_dir)
 
@@ -87,27 +71,19 @@ class YoloFormatTest(TestCase):
             compare_datasets(self, source_dataset, parsed_dataset)
 
     def test_can_load_dataset_with_exact_image_info(self):
-        class TestExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        source_dataset = Dataset.from_iterable([
             DatasetItem(id=1, subset='train',
                 image=Image(path='1.jpg', size=(10, 15)),
                 annotations=[
                     Bbox(0, 2, 4, 2, label=2),
                     Bbox(3, 3, 2, 3, label=4),
                 ]),
-                ])
-
-            def categories(self):
-                label_categories = LabelCategories()
-                for i in range(10):
-                    label_categories.add('label_' + str(i))
-                return {
-                    AnnotationType.label: label_categories,
-                }
+        ], categories={
+            AnnotationType.label: LabelCategories.from_iterable(
+                'label_' + str(i) for i in range(10)),
+        })
 
         with TestDir() as test_dir:
-            source_dataset = TestExtractor()
 
             YoloConverter.convert(source_dataset, test_dir)
 
@@ -117,24 +93,20 @@ class YoloFormatTest(TestCase):
             compare_datasets(self, source_dataset, parsed_dataset)
 
     def test_relative_paths(self):
-        class TestExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        source_dataset = Dataset.from_iterable([
             DatasetItem(id='1', subset='train',
                 image=np.ones((4, 2, 3))),
             DatasetItem(id='subdir1/1', subset='train',
                 image=np.ones((2, 6, 3))),
             DatasetItem(id='subdir2/1', subset='train',
                 image=np.ones((5, 4, 3))),
-                ])
-
-            def categories(self):
-                return { AnnotationType.label: LabelCategories() }
+        ], categories={
+            AnnotationType.label: LabelCategories(),
+        })
 
         for save_images in {True, False}:
             with self.subTest(save_images=save_images):
                 with TestDir() as test_dir:
-                    source_dataset = TestExtractor()
 
                     YoloConverter.convert(source_dataset, test_dir,
                         save_images=save_images)
@@ -150,26 +122,19 @@ class YoloImporterTest(TestCase):
         self.assertTrue(YoloImporter.detect(DUMMY_DATASET_DIR))
 
     def test_can_import(self):
-        class DstExtractor(Extractor):
-            def __iter__(self):
-                return iter([
+        expected_dataset = Dataset.from_iterable([
             DatasetItem(id=1, subset='train',
                 image=np.ones((10, 15, 3)),
                 annotations=[
                     Bbox(0, 2, 4, 2, label=2),
                     Bbox(3, 3, 2, 3, label=4),
                 ]),
-                ])
-
-            def categories(self):
-                label_categories = LabelCategories()
-                for i in range(10):
-                    label_categories.add('label_' + str(i))
-                return {
-                    AnnotationType.label: label_categories,
-                }
+        ], categories={
+            AnnotationType.label: LabelCategories.from_iterable(
+                'label_' + str(i) for i in range(10)),
+        })
 
         dataset = Project.import_from(DUMMY_DATASET_DIR, 'yolo') \
             .make_dataset()
 
-        compare_datasets(self, DstExtractor(), dataset)
+        compare_datasets(self, expected_dataset, dataset)