From bf6580bbbdf778efcb0d96b38fd99a9d41c4659e Mon Sep 17 00:00:00 2001
From: Zhiltsov Max <zhiltsov.max35@gmail.com>
Date: Fri, 7 Feb 2020 17:40:44 +0300
Subject: [PATCH 01/25] Employ transforms and item wrapper

---
 .../datumaro/components/dataset_filter.py     | 52 +++--------------
 datumaro/datumaro/components/extractor.py     | 23 ++++----
 datumaro/datumaro/components/project.py       | 57 +++----------------
 3 files changed, 27 insertions(+), 105 deletions(-)

diff --git a/datumaro/datumaro/components/dataset_filter.py b/datumaro/datumaro/components/dataset_filter.py
index 5037331f07a..e0f337e222a 100644
--- a/datumaro/datumaro/components/dataset_filter.py
+++ b/datumaro/datumaro/components/dataset_filter.py
@@ -4,7 +4,7 @@
 # SPDX-License-Identifier: MIT
 
 from lxml import etree as ET # NOTE: lxml has proper XPath implementation
-from datumaro.components.extractor import (DatasetItem, Extractor,
+from datumaro.components.extractor import (Transform,
     Annotation, AnnotationType,
     Label, Mask, Points, Polygon, PolyLine, Bbox, Caption,
 )
@@ -218,39 +218,9 @@ def XPathDatasetFilter(extractor, xpath=None):
         DatasetItemEncoder.encode(item, extractor.categories())))
     return extractor.select(f)
 
-class XPathAnnotationsFilter(Extractor): # NOTE: essentially, a transform
-    class ItemWrapper(DatasetItem):
-        def __init__(self, item, annotations):
-            self._item = item
-            self._annotations = annotations
-
-        @DatasetItem.id.getter
-        def id(self):
-            return self._item.id
-
-        @DatasetItem.subset.getter
-        def subset(self):
-            return self._item.subset
-
-        @DatasetItem.path.getter
-        def path(self):
-            return self._item.path
-
-        @DatasetItem.annotations.getter
-        def annotations(self):
-            return self._annotations
-
-        @DatasetItem.has_image.getter
-        def has_image(self):
-            return self._item.has_image
-
-        @DatasetItem.image.getter
-        def image(self):
-            return self._item.image
-
+class XPathAnnotationsFilter(Transform):
     def __init__(self, extractor, xpath=None, remove_empty=False):
-        super().__init__()
-        self._extractor = extractor
+        super().__init__(extractor)
 
         if xpath is not None:
             xpath = ET.XPath(xpath)
@@ -258,24 +228,16 @@ def __init__(self, extractor, xpath=None, remove_empty=False):
 
         self._remove_empty = remove_empty
 
-    def __len__(self):
-        return len(self._extractor)
-
     def __iter__(self):
         for item in self._extractor:
-            item = self._filter_item(item)
+            item = self.transform_item(item)
             if item is not None:
                 yield item
 
-    def subsets(self):
-        return self._extractor.subsets()
-
-    def categories(self):
-        return self._extractor.categories()
-
-    def _filter_item(self, item):
+    def transform_item(self, item):
         if self._filter is None:
             return item
+
         encoded = DatasetItemEncoder.encode(item, self._extractor.categories())
         filtered = self._filter(encoded)
         filtered = [elem for elem in filtered if elem.tag == 'annotation']
@@ -285,4 +247,4 @@ def _filter_item(self, item):
 
         if self._remove_empty and len(annotations) == 0:
             return None
-        return self.ItemWrapper(item, annotations)
\ No newline at end of file
+        return self.wrap_item(item, annotations=annotations)
\ No newline at end of file
diff --git a/datumaro/datumaro/components/extractor.py b/datumaro/datumaro/components/extractor.py
index b6d7be0cb2c..ffd33fc9ade 100644
--- a/datumaro/datumaro/components/extractor.py
+++ b/datumaro/datumaro/components/extractor.py
@@ -637,6 +637,16 @@ def __eq__(self, other):
             (self.has_image and np.array_equal(self.image, other.image) or \
                 not self.has_image)
 
+    def wrap(item, **kwargs):
+        expected_args = {'id', 'annotations', 'subset', 'path', 'image'}
+        for k in expected_args:
+            if k not in kwargs:
+                if k == 'image' and item.has_image:
+                    kwargs[k] = lambda: item.image
+                else:
+                    kwargs[k] = getattr(item, k)
+        return DatasetItem(**kwargs)
+
 class IExtractor:
     def __iter__(self):
         raise NotImplementedError()
@@ -741,16 +751,9 @@ def __call__(self, path, **extra_params):
         raise NotImplementedError()
 
 class Transform(Extractor):
-    @classmethod
-    def wrap_item(cls, item, **kwargs):
-        expected_args = {'id', 'annotations', 'subset', 'path', 'image'}
-        for k in expected_args:
-            if k not in kwargs:
-                if k == 'image' and item.has_image:
-                    kwargs[k] = lambda: item.image
-                else:
-                    kwargs[k] = getattr(item, k)
-        return DatasetItem(**kwargs)
+    @staticmethod
+    def wrap_item(item, **kwargs):
+        return item.wrap(**kwargs)
 
     def __init__(self, extractor):
         super().__init__()
diff --git a/datumaro/datumaro/components/project.py b/datumaro/datumaro/components/project.py
index f9ee8c313ee..a4dafe0f9d6 100644
--- a/datumaro/datumaro/components/project.py
+++ b/datumaro/datumaro/components/project.py
@@ -302,45 +302,6 @@ def __len__(self):
     def categories(self):
         return self._parent.categories()
 
-class DatasetItemWrapper(DatasetItem):
-    def __init__(self, item, path, annotations, image=None):
-        self._item = item
-        if path is None:
-            path = []
-        self._path = path
-        self._annotations = annotations
-        self._image = image
-
-    @DatasetItem.id.getter
-    def id(self):
-        return self._item.id
-
-    @DatasetItem.subset.getter
-    def subset(self):
-        return self._item.subset
-
-    @DatasetItem.path.getter
-    def path(self):
-        return self._path
-
-    @DatasetItem.annotations.getter
-    def annotations(self):
-        return self._annotations
-
-    @DatasetItem.has_image.getter
-    def has_image(self):
-        if self._image is not None:
-            return True
-        return self._item.has_image
-
-    @DatasetItem.image.getter
-    def image(self):
-        if self._image is not None:
-            if callable(self._image):
-                return self._image()
-            return self._image
-        return self._item.image
-
 class Dataset(Extractor):
     @classmethod
     def from_extractors(cls, *sources):
@@ -369,12 +330,11 @@ def from_extractors(cls, *sources):
                         # TODO: think of image comparison
                         image = cls._lazy_image(existing_item)
 
-                    item = DatasetItemWrapper(item=item, path=path,
+                    item = item.wrap(path=path,
                         image=image, annotations=self._merge_anno(
                             existing_item.annotations, item.annotations))
                 else:
-                    item = DatasetItemWrapper(item=item, path=path,
-                        annotations=item.annotations)
+                    item = item.wrap(path=path, annotations=item.annotations)
 
                 subsets[item.subset].items[item.id] = item
 
@@ -423,8 +383,7 @@ def put(self, item, item_id=None, subset=None, path=None):
         if subset is None:
             subset = item.subset
 
-        item = DatasetItemWrapper(item=item, path=None,
-            annotations=item.annotations)
+        item = item.wrap(path=None, annotations=item.annotations)
         if item.subset not in self._subsets:
             self._subsets[item.subset] = Subset(self)
         self._subsets[subset].items[item_id] = item
@@ -528,7 +487,7 @@ def __init__(self, project):
                     path = existing_item.path
                     if item.path != path:
                         path = None # NOTE: move to our own dataset
-                    item = DatasetItemWrapper(item=item, path=path,
+                    item = item.wrap(path=path,
                         image=image, annotations=self._merge_anno(
                             existing_item.annotations, item.annotations))
                 else:
@@ -542,8 +501,7 @@ def __init__(self, project):
                         if path is None:
                             path = []
                         path = [source_name] + path
-                    item = DatasetItemWrapper(item=item, path=path,
-                        annotations=item.annotations)
+                    item = item.wrap(path=path, annotations=item.annotations)
 
                 subsets[item.subset].items[item.id] = item
 
@@ -558,7 +516,7 @@ def __init__(self, project):
                         if existing_item.has_image:
                             # TODO: think of image comparison
                             image = self._lazy_image(existing_item)
-                        item = DatasetItemWrapper(item=item, path=None,
+                        item = item.wrap(path=None,
                             annotations=item.annotations, image=image)
 
                 subsets[item.subset].items[item.id] = item
@@ -597,8 +555,7 @@ def put(self, item, item_id=None, subset=None, path=None):
         if subset is None:
             subset = item.subset
 
-        item = DatasetItemWrapper(item=item, path=path,
-            annotations=item.annotations)
+        item = item.wrap(path=path, annotations=item.annotations)
         if item.subset not in self._subsets:
             self._subsets[item.subset] = Subset(self)
         self._subsets[subset].items[item_id] = item

From baeaf869a474021f8161c2a2ecf260ec79f4346b Mon Sep 17 00:00:00 2001
From: Zhiltsov Max <zhiltsov.max35@gmail.com>
Date: Wed, 12 Feb 2020 14:32:33 +0300
Subject: [PATCH 02/25] Add image class and tests

---
 datumaro/datumaro/util/image.py         | 73 ++++++++++++++++++++++---
 datumaro/tests/test_coco_format.py      | 59 ++++++++++++--------
 datumaro/tests/test_cvat_format.py      |  8 ++-
 datumaro/tests/test_datumaro_format.py  |  5 +-
 datumaro/tests/test_image.py            |  4 +-
 datumaro/tests/test_image_dir_format.py |  3 +-
 datumaro/tests/test_images.py           | 42 ++++++++++++--
 datumaro/tests/test_masks.py            | 19 +++++++
 datumaro/tests/test_project.py          | 49 ++++++++++++-----
 datumaro/tests/test_tfrecord_format.py  | 15 +++++
 datumaro/tests/test_voc_format.py       | 17 ++++--
 datumaro/tests/test_yolo_format.py      | 36 +++++++++++-
 12 files changed, 268 insertions(+), 62 deletions(-)

diff --git a/datumaro/datumaro/util/image.py b/datumaro/datumaro/util/image.py
index 395c6f5080d..2d9a4d43326 100644
--- a/datumaro/datumaro/util/image.py
+++ b/datumaro/datumaro/util/image.py
@@ -123,14 +123,16 @@ def decode_image(image_bytes):
 
 
 class lazy_image:
-    def __init__(self, path, loader=load_image, cache=None):
+    def __init__(self, path, loader=None, cache=None):
+        if loader is None:
+            loader = load_image
         self.path = path
         self.loader = loader
 
         # Cache:
         # - False: do not cache
-        # - None: use default (don't store in a class variable)
-        # - object: use this object as a cache
+        # - None: use the global cache
+        # - object: an object to be used as cache
         assert cache in {None, False} or isinstance(cache, object)
         self.cache = cache
 
@@ -138,9 +140,9 @@ def __call__(self):
         image = None
         image_id = hash(self) # path is not necessary hashable or a file path
 
-        cache = self._get_cache()
+        cache = self._get_cache(self.cache)
         if cache is not None:
-            image = self._get_cache().get(image_id)
+            image = cache.get(image_id)
 
         if image is None:
             image = self.loader(self.path)
@@ -148,8 +150,8 @@ def __call__(self):
                 cache.push(image_id, image)
         return image
 
-    def _get_cache(self):
-        cache = self.cache
+    @staticmethod
+    def _get_cache(cache):
         if cache is None:
             cache = _ImageCache.get_instance()
         elif cache == False:
@@ -157,4 +159,59 @@ def _get_cache(self):
         return cache
 
     def __hash__(self):
-        return hash((id(self), self.path, self.loader))
\ No newline at end of file
+        return hash((id(self), self.path, self.loader))
+
+class Image:
+    def __init__(self, data=None, path=None, loader=None, cache=None,
+            size=None):
+        assert size is None or len(size) == 2
+        if size is not None:
+            assert 0 < size[0] and 0 < size[1], size
+            size = tuple(size)
+        else:
+            size = None
+        self._size = size # (H, W)
+
+        assert path is None or isinstance(path, str)
+        if not path:
+            path = ''
+        self._path = path
+
+        assert any(e is not None for e in (data, path, loader, size)), "Image can not be empty"
+        if data is None and (path or loader is not None):
+            data = lazy_image(path, loader=loader, cache=cache)
+        self._data = data
+
+    @property
+    def path(self):
+        return self._path
+
+    @property
+    def data(self):
+        if callable(self._data):
+            return self._data()
+        return self._data
+
+    @property
+    def has_data(self):
+        return self._data is not None
+
+    @property
+    def size(self):
+        if self._size is None:
+            data = self.data
+            if data is not None:
+                self._size = data.shape[:2]
+        return self._size
+
+    def __eq__(self, other):
+        if isinstance(other, np.ndarray):
+            return self.has_data and np.array_equal(self.data, other)
+
+        if not isinstance(other, __class__):
+            return False
+        return \
+            (np.array_equal(self.size, other.size)) and \
+            (self.has_data == other.has_data) and \
+            (self.has_data and np.array_equal(self.data, other.data) or \
+                not self.has_data)
\ No newline at end of file
diff --git a/datumaro/tests/test_coco_format.py b/datumaro/tests/test_coco_format.py
index 9dd64f878fd..9283f475910 100644
--- a/datumaro/tests/test_coco_format.py
+++ b/datumaro/tests/test_coco_format.py
@@ -19,8 +19,7 @@
     CocoLabelsConverter,
 )
 from datumaro.plugins.coco_format.importer import CocoImporter
-from datumaro.util.image import save_image
-from datumaro.util import find
+from datumaro.util.image import save_image, Image
 from datumaro.util.test_utils import TestDir, compare_datasets
 
 
@@ -100,7 +99,7 @@ def COCO_dataset_generate(self, path):
         os.makedirs(img_dir)
         os.makedirs(ann_dir)
 
-        image = np.ones((10, 5, 3), dtype=np.uint8)
+        image = np.ones((10, 5, 3))
         save_image(osp.join(img_dir, '000000000001.jpg'), image)
 
         annotation = self.generate_annotation()
@@ -109,30 +108,33 @@ def COCO_dataset_generate(self, path):
             json.dump(annotation, outfile)
 
     def test_can_import(self):
-        with TestDir() as test_dir:
-            self.COCO_dataset_generate(test_dir)
-            project = Project.import_from(test_dir, 'coco')
-            dataset = project.make_dataset()
+        class DstExtractor(Extractor):
+            def __iter__(self):
+                return iter([
+                    DatasetItem(id=1, image=np.ones((10, 5, 3)), subset='val',
+                        annotations=[
+                            Polygon([0, 0, 1, 0, 1, 2, 0, 2], label=0,
+                                id=1, group=1, attributes={'is_crowd': False}),
+                            Mask(np.array(
+                                [[1, 0, 0, 1, 0]] * 5 +
+                                [[1, 1, 1, 1, 0]] * 5
+                                ), label=0,
+                                id=2, group=2, attributes={'is_crowd': True}),
+                        ]
+                    ),
+                ])
 
-            self.assertListEqual(['val'], sorted(dataset.subsets()))
-            self.assertEqual(1, len(dataset))
+            def categories(self):
+                label_cat = LabelCategories()
+                label_cat.add('TEST')
+                return { AnnotationType.label: label_cat }
 
-            item = next(iter(dataset))
-            self.assertTrue(item.has_image)
-            self.assertEqual(np.sum(item.image), np.prod(item.image.shape))
-            self.assertEqual(2, len(item.annotations))
+        with TestDir() as test_dir:
+            self.COCO_dataset_generate(test_dir)
 
-            ann_1 = find(item.annotations, lambda x: x.id == 1)
-            ann_1_poly = find(item.annotations, lambda x: \
-                x.group == ann_1.id and x.type == AnnotationType.polygon)
-            self.assertFalse(ann_1 is None)
-            self.assertFalse(ann_1_poly is None)
+            dataset = Project.import_from(test_dir, 'coco').make_dataset()
 
-            ann_2 = find(item.annotations, lambda x: x.id == 2)
-            ann_2_mask = find(item.annotations, lambda x: \
-                x.group == ann_2.id and x.type == AnnotationType.mask)
-            self.assertFalse(ann_2 is None)
-            self.assertFalse(ann_2_mask is None)
+            compare_datasets(self, DstExtractor(), dataset)
 
 class CocoConverterTest(TestCase):
     def _test_save_and_load(self, source_dataset, converter, test_dir,
@@ -630,6 +632,17 @@ def categories(self):
                     AnnotationType.label: label_cat,
                 }
 
+        with TestDir() as test_dir:
+            self._test_save_and_load(TestExtractor(),
+                CocoConverter(), test_dir)
+
+    def test_can_save_dataset_with_image_info(self):
+        class TestExtractor(Extractor):
+            def __iter__(self):
+                return iter([
+                    DatasetItem(id=1, image=Image(size=(10, 15))),
+                ])
+
         with TestDir() as test_dir:
             self._test_save_and_load(TestExtractor(),
                 CocoConverter(), test_dir)
\ No newline at end of file
diff --git a/datumaro/tests/test_cvat_format.py b/datumaro/tests/test_cvat_format.py
index dc4a2dc4397..43d9520508f 100644
--- a/datumaro/tests/test_cvat_format.py
+++ b/datumaro/tests/test_cvat_format.py
@@ -12,7 +12,7 @@
 from datumaro.plugins.cvat_format.importer import CvatImporter
 from datumaro.plugins.cvat_format.converter import CvatConverter
 from datumaro.plugins.cvat_format.format import CvatPath
-from datumaro.util.image import save_image
+from datumaro.util.image import save_image, Image
 from datumaro.util.test_utils import TestDir, compare_datasets
 
 
@@ -82,7 +82,7 @@ def generate_dummy_cvat(path):
         save_image(osp.join(images_dir, 'img1.jpg'), np.ones((10, 10, 3)))
         item2_elem = ET.SubElement(root_elem, 'image')
         item2_elem.attrib.update({
-            'id': '1', 'name': 'img1', 'width': '8', 'height': '8'
+            'id': '1', 'name': 'img1', 'width': '10', 'height': '10'
         })
 
         item2_ann1_elem = ET.SubElement(item2_elem, 'polygon')
@@ -192,6 +192,8 @@ def __iter__(self):
                             PolyLine([5, 0, 9, 0, 5, 5]), # will be skipped as no label
                         ]
                     ),
+
+                    DatasetItem(id=3, subset='s3', image=Image(size=(2, 4))),
                 ])
 
             def categories(self):
@@ -232,6 +234,8 @@ def __iter__(self):
                                 attributes={ 'z_order': 1, 'occluded': False }),
                         ]
                     ),
+
+                    DatasetItem(id=3, subset='s3', image=Image(size=(2, 4))),
                 ])
 
             def categories(self):
diff --git a/datumaro/tests/test_datumaro_format.py b/datumaro/tests/test_datumaro_format.py
index 4a168a6deaa..0bc3911af23 100644
--- a/datumaro/tests/test_datumaro_format.py
+++ b/datumaro/tests/test_datumaro_format.py
@@ -9,8 +9,9 @@
     LabelCategories, MaskCategories, PointsCategories
 )
 from datumaro.plugins.datumaro_format.converter import DatumaroConverter
-from datumaro.util.test_utils import TestDir, item_to_str
 from datumaro.util.mask_tools import generate_colormap
+from datumaro.util.image import Image
+from datumaro.util.test_utils import TestDir, item_to_str
 
 
 class DatumaroConverterTest(TestCase):
@@ -48,7 +49,7 @@ def __iter__(self):
                 DatasetItem(id=42, subset='test'),
 
                 DatasetItem(id=42),
-                DatasetItem(id=43),
+                DatasetItem(id=43, image=Image(size=(2, 4))),
             ])
 
         def categories(self):
diff --git a/datumaro/tests/test_image.py b/datumaro/tests/test_image.py
index 9614fcc90d7..efb7aea2969 100644
--- a/datumaro/tests/test_image.py
+++ b/datumaro/tests/test_image.py
@@ -8,7 +8,7 @@
 from datumaro.util.test_utils import TestDir
 
 
-class ImageTest(TestCase):
+class ImageOperationsTest(TestCase):
     def setUp(self):
         self.default_backend = image_module._IMAGE_BACKEND
 
@@ -49,4 +49,4 @@ def test_encode_and_decode_backends(self):
             dst_image = image_module.decode_image(buffer)
 
             self.assertTrue(np.array_equal(src_image, dst_image),
-                'save: %s, load: %s' % (save_backend, load_backend))
\ No newline at end of file
+                'save: %s, load: %s' % (save_backend, load_backend))
diff --git a/datumaro/tests/test_image_dir_format.py b/datumaro/tests/test_image_dir_format.py
index 6b382c212cb..f7875df0630 100644
--- a/datumaro/tests/test_image_dir_format.py
+++ b/datumaro/tests/test_image_dir_format.py
@@ -22,7 +22,8 @@ def test_can_load(self):
             source_dataset = self.TestExtractor()
 
             for item in source_dataset:
-                save_image(osp.join(test_dir, '%s.jpg' % item.id), item.image)
+                save_image(osp.join(test_dir, '%s.jpg' % item.id),
+                    item.image.data)
 
             project = Project.import_from(test_dir, 'image_dir')
             parsed_dataset = project.make_dataset()
diff --git a/datumaro/tests/test_images.py b/datumaro/tests/test_images.py
index e3f12a3df89..5f6dcb48a38 100644
--- a/datumaro/tests/test_images.py
+++ b/datumaro/tests/test_images.py
@@ -1,11 +1,10 @@
 import numpy as np
 import os.path as osp
-from PIL import Image
 
 from unittest import TestCase
 
 from datumaro.util.test_utils import TestDir
-from datumaro.util.image import lazy_image
+from datumaro.util.image import lazy_image, load_image, save_image, Image
 from datumaro.util.image_cache import ImageCache
 
 
@@ -13,10 +12,8 @@ class LazyImageTest(TestCase):
     def test_cache_works(self):
         with TestDir() as test_dir:
             image = np.ones((100, 100, 3), dtype=np.uint8)
-            image = Image.fromarray(image).convert('RGB')
-
             image_path = osp.join(test_dir, 'image.jpg')
-            image.save(image_path)
+            save_image(image_path, image)
 
             caching_loader = lazy_image(image_path, cache=None)
             self.assertTrue(caching_loader() is caching_loader())
@@ -46,4 +43,37 @@ def test_global_cache_is_accessible(self):
 
         ImageCache.get_instance().clear()
         self.assertTrue(loader() is loader())
-        self.assertEqual(ImageCache.get_instance().size(), 1)
\ No newline at end of file
+        self.assertEqual(ImageCache.get_instance().size(), 1)
+
+class ImageTest(TestCase):
+    def test_lazy_image_shape(self):
+        loader = lambda _: np.ones((5, 6, 7))
+
+        image_lazy = Image(loader=loader, size=(2, 4))
+        image_eager = Image(loader=loader)
+
+        self.assertEqual((2, 4), image_lazy.size)
+        self.assertEqual((5, 6), image_eager.size)
+
+    @staticmethod
+    def test_ctors():
+        with TestDir() as test_dir:
+            path = osp.join(test_dir, 'path.png')
+            image = np.ones([2, 4, 3])
+            save_image(path, image)
+
+            for args in [
+                { 'data': image },
+                { 'data': image, 'path': path },
+                { 'data': image, 'path': path, 'size': (2, 4) },
+                { 'path': path },
+                { 'path': path, 'loader': load_image },
+                { 'path': path, 'size': (2, 4) },
+                { 'size': (2, 4) },
+            ]:
+                img = Image(**args)
+                # pylint: disable=pointless-statement
+                if img.has_data:
+                    img.data
+                img.size
+                # pylint: enable=pointless-statement
diff --git a/datumaro/tests/test_masks.py b/datumaro/tests/test_masks.py
index 1619f1db1ad..274181eaa23 100644
--- a/datumaro/tests/test_masks.py
+++ b/datumaro/tests/test_masks.py
@@ -68,6 +68,25 @@ def test_can_crop_covered_segments(self):
             self.assertTrue(np.array_equal(e_mask, c_mask),
                 '#%s: %s\n%s\n' % (i, e_mask, c_mask))
 
+    def test_mask_to_rle(self):
+        source_mask = np.array([
+            [0, 1, 1, 1, 0, 1, 1, 1, 1, 0],
+            [0, 0, 1, 1, 0, 1, 0, 1, 0, 0],
+            [0, 0, 0, 1, 0, 1, 1, 0, 0, 0],
+            [0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
+            [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+        ])
+
+        rle_uncompressed = mask_tools.mask_to_rle(source_mask)
+
+        from pycocotools import mask as mask_utils
+        resulting_mask = mask_utils.frPyObjects(
+            rle_uncompressed, *rle_uncompressed['size'])
+        resulting_mask = mask_utils.decode(resulting_mask)
+
+        self.assertTrue(np.array_equal(source_mask, resulting_mask),
+            '%s\n%s\n' % (source_mask, resulting_mask))
+
 class ColormapOperationsTest(TestCase):
     def test_can_paint_mask(self):
         mask = np.zeros((1, 3), dtype=np.uint8)
diff --git a/datumaro/tests/test_project.py b/datumaro/tests/test_project.py
index 8c4106687d7..1a0038d02bc 100644
--- a/datumaro/tests/test_project.py
+++ b/datumaro/tests/test_project.py
@@ -11,10 +11,11 @@
 from datumaro.components.extractor import (Extractor, DatasetItem,
     Label, Mask, Points, Polygon, PolyLine, Bbox, Caption,
 )
+from datumaro.util.image import Image
 from datumaro.components.config import Config, DefaultConfig, SchemaBuilder
 from datumaro.components.dataset_filter import \
     XPathDatasetFilter, XPathAnnotationsFilter, DatasetItemEncoder
-from datumaro.util.test_utils import TestDir
+from datumaro.util.test_utils import TestDir, compare_datasets
 
 
 class ProjectTest(TestCase):
@@ -135,12 +136,12 @@ def test_can_batch_launch_custom_model(self):
         class TestExtractor(Extractor):
             def __iter__(self):
                 for i in range(5):
-                    yield DatasetItem(id=i, subset='train', image=i)
+                    yield DatasetItem(id=i, subset='train', image=np.array([i]))
 
         class TestLauncher(Launcher):
             def launch(self, inputs):
                 for i, inp in enumerate(inputs):
-                    yield [ Label(attributes={'idx': i, 'data': inp}) ]
+                    yield [ Label(attributes={'idx': i, 'data': inp.item()}) ]
 
         model_name = 'model'
         launcher_name = 'custom_launcher'
@@ -165,19 +166,18 @@ def test_can_do_transform_with_custom_model(self):
         class TestExtractorSrc(Extractor):
             def __iter__(self):
                 for i in range(2):
-                    yield DatasetItem(id=i, subset='train', image=i,
-                        annotations=[ Label(i) ])
+                    yield DatasetItem(id=i, image=np.ones([2, 2, 3]) * i,
+                        annotations=[Label(i)])
 
         class TestLauncher(Launcher):
             def launch(self, inputs):
                 for inp in inputs:
-                    yield [ Label(inp) ]
+                    yield [ Label(inp[0, 0, 0]) ]
 
         class TestConverter(Converter):
             def __call__(self, extractor, save_dir):
                 for item in extractor:
                     with open(osp.join(save_dir, '%s.txt' % item.id), 'w') as f:
-                        f.write(str(item.subset) + '\n')
                         f.write(str(item.annotations[0].label) + '\n')
 
         class TestExtractorDst(Extractor):
@@ -189,11 +189,8 @@ def __iter__(self):
                 for path in self.items:
                     with open(path, 'r') as f:
                         index = osp.splitext(osp.basename(path))[0]
-                        subset = f.readline().strip()
                         label = int(f.readline().strip())
-                        assert subset == 'train'
-                        yield DatasetItem(id=index, subset=subset,
-                            annotations=[ Label(label) ])
+                        yield DatasetItem(id=index, annotations=[Label(label)])
 
         model_name = 'model'
         launcher_name = 'custom_launcher'
@@ -476,6 +473,11 @@ def __iter__(self):
                     DatasetItem(id=2, subset='train'),
 
                     DatasetItem(id=3, subset='test'),
+                    DatasetItem(id=4, subset='test'),
+
+                    DatasetItem(id=1),
+                    DatasetItem(id=2),
+                    DatasetItem(id=3),
                 ])
 
         extractor_name = 'ext1'
@@ -485,8 +487,29 @@ def __iter__(self):
             'url': 'path',
             'format': extractor_name,
         })
-        project.set_subsets(['train'])
 
         dataset = project.make_dataset()
 
-        self.assertEqual(3, len(dataset))
+        compare_datasets(self, CustomExtractor(), dataset)
+
+class DatasetItemTest(TestCase):
+    def test_ctor_requires_id(self):
+        has_error = False
+        try:
+            # pylint: disable=no-value-for-parameter
+            DatasetItem()
+            # pylint: enable=no-value-for-parameter
+        except TypeError:
+            has_error = True
+
+        self.assertTrue(has_error)
+
+    @staticmethod
+    def test_ctors_with_image():
+        for args in [
+            { 'id': 0, 'image': None },
+            { 'id': 0, 'image': np.array([1, 2, 3]) },
+            { 'id': 0, 'image': lambda f: np.array([1, 2, 3]) },
+            { 'id': 0, 'image': Image(data=np.array([1, 2, 3])) },
+        ]:
+            DatasetItem(**args)
\ No newline at end of file
diff --git a/datumaro/tests/test_tfrecord_format.py b/datumaro/tests/test_tfrecord_format.py
index 664359d0bd0..31d93f34f6c 100644
--- a/datumaro/tests/test_tfrecord_format.py
+++ b/datumaro/tests/test_tfrecord_format.py
@@ -8,6 +8,7 @@
 from datumaro.plugins.tf_detection_api_format.importer import TfDetectionApiImporter
 from datumaro.plugins.tf_detection_api_format.extractor import TfDetectionApiExtractor
 from datumaro.plugins.tf_detection_api_format.converter import TfDetectionApiConverter
+from datumaro.util.image import Image
 from datumaro.util.test_utils import TestDir, compare_datasets
 
 
@@ -101,6 +102,20 @@ def categories(self):
                 TestExtractor(), TfDetectionApiConverter(save_images=True),
                 test_dir)
 
+    def test_can_save_dataset_with_image_info(self):
+        class TestExtractor(Extractor):
+            def __iter__(self):
+                return iter([
+                    DatasetItem(id=1, image=Image(size=(10, 15))),
+                ])
+
+            def categories(self):
+                return { AnnotationType.label: LabelCategories() }
+
+        with TestDir() as test_dir:
+            self._test_save_and_load(TestExtractor(),
+                TfDetectionApiConverter(), test_dir)
+
     def test_labelmap_parsing(self):
         text = """
             {
diff --git a/datumaro/tests/test_voc_format.py b/datumaro/tests/test_voc_format.py
index ca2733fc426..30473e21f14 100644
--- a/datumaro/tests/test_voc_format.py
+++ b/datumaro/tests/test_voc_format.py
@@ -27,7 +27,7 @@
 )
 from datumaro.plugins.voc_format.importer import VocImporter
 from datumaro.components.project import Project
-from datumaro.util.image import save_image
+from datumaro.util.image import save_image, Image
 from datumaro.util.test_utils import TestDir, compare_datasets
 
 
@@ -171,13 +171,11 @@ def generate_dummy_voc(path):
     return subsets
 
 class TestExtractorBase(Extractor):
-    _categories = VOC.make_voc_categories()
-
     def _label(self, voc_label):
         return self.categories()[AnnotationType.label].find(voc_label)[0]
 
     def categories(self):
-        return self._categories
+        return VOC.make_voc_categories()
 
 class VocExtractorTest(TestCase):
     def test_can_load_voc_cls(self):
@@ -694,6 +692,17 @@ def categories(self):
                 SrcExtractor(), VocConverter(label_map=label_map),
                 test_dir, target_dataset=DstExtractor())
 
+    def test_can_save_dataset_with_image_info(self):
+        class TestExtractor(TestExtractorBase):
+            def __iter__(self):
+                return iter([
+                    DatasetItem(id=1, image=Image(size=(10, 15))),
+                ])
+
+        with TestDir() as test_dir:
+            self._test_save_and_load(TestExtractor(),
+                VocConverter(label_map='voc'), test_dir)
+
 class VocImportTest(TestCase):
     def test_can_import(self):
         with TestDir() as test_dir:
diff --git a/datumaro/tests/test_yolo_format.py b/datumaro/tests/test_yolo_format.py
index 9ce972a206f..bad4a5e4c18 100644
--- a/datumaro/tests/test_yolo_format.py
+++ b/datumaro/tests/test_yolo_format.py
@@ -1,4 +1,5 @@
 import numpy as np
+import os.path as osp
 
 from unittest import TestCase
 
@@ -6,7 +7,9 @@
     AnnotationType, Bbox, LabelCategories,
 )
 from datumaro.plugins.yolo_format.importer import YoloImporter
+from datumaro.plugins.yolo_format.format import YoloPath
 from datumaro.plugins.yolo_format.converter import YoloConverter
+from datumaro.util.image import Image, save_image
 from datumaro.util.test_utils import TestDir, compare_datasets
 
 
@@ -50,4 +53,35 @@ def categories(self):
             YoloConverter(save_images=True)(source_dataset, test_dir)
             parsed_dataset = YoloImporter()(test_dir).make_dataset()
 
-            compare_datasets(self, source_dataset, parsed_dataset)
\ No newline at end of file
+            compare_datasets(self, source_dataset, parsed_dataset)
+
+    def test_can_save_dataset_with_image_info(self):
+        class TestExtractor(Extractor):
+            def __iter__(self):
+                return iter([
+                    DatasetItem(id=1, subset='train',
+                        image=Image(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,
+                }
+
+        with TestDir() as test_dir:
+            source_dataset = TestExtractor()
+
+            YoloConverter()(source_dataset, test_dir)
+
+            save_image(osp.join(test_dir, 'obj_train_data', '1.jpg'),
+                np.ones((10, 15, 3))) # put the image for dataset
+            parsed_dataset = YoloImporter()(test_dir).make_dataset()
+
+            compare_datasets(self, source_dataset, parsed_dataset)

From f79768c77d1ce720bdeae9b05118d1d9fe66937e Mon Sep 17 00:00:00 2001
From: Zhiltsov Max <zhiltsov.max35@gmail.com>
Date: Wed, 12 Feb 2020 18:14:22 +0300
Subject: [PATCH 03/25] Add image info support to formats

---
 .../datumaro/components/dataset_filter.py     | 11 ++--
 datumaro/datumaro/components/extractor.py     | 45 +++++++++------
 datumaro/datumaro/components/launcher.py      | 54 ++++--------------
 datumaro/datumaro/components/project.py       | 35 ++++++------
 .../datumaro/plugins/coco_format/converter.py | 30 ++++++----
 .../datumaro/plugins/cvat_format/converter.py | 31 ++++++----
 .../datumaro/plugins/cvat_format/extractor.py | 57 ++++++++++++++-----
 .../plugins/datumaro_format/converter.py      | 12 +++-
 .../plugins/datumaro_format/extractor.py      | 19 +++++--
 datumaro/datumaro/plugins/image_dir.py        | 16 +++++-
 .../tf_detection_api_format/converter.py      | 23 +++++---
 .../tf_detection_api_format/extractor.py      | 19 +++++--
 datumaro/datumaro/plugins/transforms.py       |  6 +-
 .../datumaro/plugins/voc_format/converter.py  |  7 ++-
 .../datumaro/plugins/yolo_format/converter.py | 12 ++--
 datumaro/tests/test_image_dir_format.py       |  7 +--
 16 files changed, 227 insertions(+), 157 deletions(-)

diff --git a/datumaro/datumaro/components/dataset_filter.py b/datumaro/datumaro/components/dataset_filter.py
index e0f337e222a..f5d923bef79 100644
--- a/datumaro/datumaro/components/dataset_filter.py
+++ b/datumaro/datumaro/components/dataset_filter.py
@@ -31,11 +31,12 @@ def encode(cls, item, categories=None):
     def encode_image(cls, image):
         image_elem = ET.Element('image')
 
-        h, w = image.shape[:2]
-        c = 1 if len(image.shape) == 2 else image.shape[2]
+        h, w = image.size
         ET.SubElement(image_elem, 'width').text = str(w)
         ET.SubElement(image_elem, 'height').text = str(h)
-        ET.SubElement(image_elem, 'depth').text = str(c)
+
+        ET.SubElement(image_elem, 'has_data').text = '%d' % int(image.has_data)
+        ET.SubElement(image_elem, 'path').text = image.path
 
         return image_elem
 
@@ -82,10 +83,6 @@ def encode_mask_object(cls, obj, categories):
             str(cls._get_label(obj.label, categories))
         ET.SubElement(ann_elem, 'label_id').text = str(obj.label)
 
-        mask = obj.image
-        if mask is not None:
-            ann_elem.append(cls.encode_image(mask))
-
         return ann_elem
 
     @classmethod
diff --git a/datumaro/datumaro/components/extractor.py b/datumaro/datumaro/components/extractor.py
index ffd33fc9ade..c1b23ef6903 100644
--- a/datumaro/datumaro/components/extractor.py
+++ b/datumaro/datumaro/components/extractor.py
@@ -7,6 +7,7 @@
 from enum import Enum
 import numpy as np
 
+from datumaro.util.image import Image
 
 AnnotationType = Enum('AnnotationType',
     [
@@ -418,8 +419,8 @@ def get_area(self):
 
 class Polygon(_Shape):
     # pylint: disable=redefined-builtin
-    def __init__(self, points=None, z_order=None,
-            label=None, id=None, attributes=None, group=None):
+    def __init__(self, points=None, label=None,
+            z_order=None, id=None, attributes=None, group=None):
         if points is not None:
             # keep the message on the single line to produce
             # informative output
@@ -576,26 +577,39 @@ def __eq__(self, other):
 class DatasetItem:
     # pylint: disable=redefined-builtin
     def __init__(self, id, annotations=None,
-            subset=None, path=None, image=None):
+            subset=None, path=None, image=None, name=None):
         assert id is not None
-        if not isinstance(id, str):
-            id = str(id)
-        assert len(id) != 0
+        id = str(id)
+        # TODO: if not isinstance(id, int):
+        #     id = int(id)
         self._id = id
 
+        assert name is None or isinstance(name, str)
+        if not name:
+            name = str(id)
+        self._name = name
+
         if subset is None:
             subset = ''
-        assert isinstance(subset, str)
+        else:
+            subset = str(subset)
         self._subset = subset
 
         if path is None:
             path = []
+        else:
+            path = list(path)
         self._path = path
 
         if annotations is None:
             annotations = []
+        else:
+            annotations = list(annotations)
         self._annotations = annotations
 
+        if callable(image) or isinstance(image, np.ndarray):
+            image = Image(data=image)
+        assert image is None or isinstance(image, Image)
         self._image = image
     # pylint: enable=redefined-builtin
 
@@ -603,6 +617,10 @@ def __init__(self, id, annotations=None,
     def id(self):
         return self._id
 
+    @property
+    def name(self):
+        return self._name
+
     @property
     def subset(self):
         return self._subset
@@ -617,8 +635,6 @@ def annotations(self):
 
     @property
     def image(self):
-        if callable(self._image):
-            return self._image()
         return self._image
 
     @property
@@ -633,18 +649,13 @@ def __eq__(self, other):
             (self.subset == other.subset) and \
             (self.annotations == other.annotations) and \
             (self.path == other.path) and \
-            (self.has_image == other.has_image) and \
-            (self.has_image and np.array_equal(self.image, other.image) or \
-                not self.has_image)
+            (self.image == other.image)
 
     def wrap(item, **kwargs):
-        expected_args = {'id', 'annotations', 'subset', 'path', 'image'}
+        expected_args = {'id', 'annotations', 'subset', 'path', 'image', 'name'}
         for k in expected_args:
             if k not in kwargs:
-                if k == 'image' and item.has_image:
-                    kwargs[k] = lambda: item.image
-                else:
-                    kwargs[k] = getattr(item, k)
+                kwargs[k] = getattr(item, k)
         return DatasetItem(**kwargs)
 
 class IExtractor:
diff --git a/datumaro/datumaro/components/launcher.py b/datumaro/datumaro/components/launcher.py
index 3ac1e1fb672..eb36295584d 100644
--- a/datumaro/datumaro/components/launcher.py
+++ b/datumaro/datumaro/components/launcher.py
@@ -5,7 +5,7 @@
 
 import numpy as np
 
-from datumaro.components.extractor import DatasetItem, Extractor
+from datumaro.components.extractor import Transform
 
 
 # pylint: disable=no-self-use
@@ -23,37 +23,9 @@ def get_categories(self):
         return None
 # pylint: enable=no-self-use
 
-class InferenceWrapper(Extractor):
-    class ItemWrapper(DatasetItem):
-        def __init__(self, item, annotations, path=None):
-            super().__init__(id=item.id)
-            self._annotations = annotations
-            self._item = item
-            self._path = path
-
-        @DatasetItem.id.getter
-        def id(self):
-            return self._item.id
-
-        @DatasetItem.subset.getter
-        def subset(self):
-            return self._item.subset
-
-        @DatasetItem.path.getter
-        def path(self):
-            return self._path
-
-        @DatasetItem.annotations.getter
-        def annotations(self):
-            return self._annotations
-
-        @DatasetItem.image.getter
-        def image(self):
-            return self._item.image
-
+class InferenceWrapper(Transform):
     def __init__(self, extractor, launcher, batch_size=1):
-        super().__init__()
-        self._extractor = extractor
+        super().__init__(extractor)
         self._launcher = launcher
         self._batch_size = batch_size
 
@@ -71,25 +43,23 @@ def __iter__(self):
                 if len(batch_items) == 0:
                     break
 
-            inputs = np.array([item.image for item in batch_items])
+            inputs = np.array([item.image.data for item in batch_items])
             inference = self._launcher.launch(inputs)
 
             for item, annotations in zip(batch_items, inference):
-                yield self.ItemWrapper(item, annotations)
-
-    def __len__(self):
-        return len(self._extractor)
-
-    def subsets(self):
-        return self._extractor.subsets()
+                yield self.wrap_item(item, annotations=annotations)
 
     def get_subset(self, name):
         subset = self._extractor.get_subset(name)
-        return InferenceWrapper(subset,
-            self._launcher, self._batch_size)
+        return InferenceWrapper(subset, self._launcher, self._batch_size)
 
     def categories(self):
         launcher_override = self._launcher.get_categories()
         if launcher_override is not None:
             return launcher_override
-        return self._extractor.categories()
\ No newline at end of file
+        return self._extractor.categories()
+
+    def transform_item(self, item):
+        inputs = np.expand_dims(item.image, axis=0)
+        annotations = self._launcher.launch(inputs)[0]
+        return self.wrap_item(item, annotations=annotations)
\ No newline at end of file
diff --git a/datumaro/datumaro/components/project.py b/datumaro/datumaro/components/project.py
index a4dafe0f9d6..7688fe996f9 100644
--- a/datumaro/datumaro/components/project.py
+++ b/datumaro/datumaro/components/project.py
@@ -325,14 +325,7 @@ def from_extractors(cls, *sources):
 
                 existing_item = subsets[item.subset].items.get(item.id)
                 if existing_item is not None:
-                    image = None
-                    if existing_item.has_image:
-                        # TODO: think of image comparison
-                        image = cls._lazy_image(existing_item)
-
-                    item = item.wrap(path=path,
-                        image=image, annotations=self._merge_anno(
-                            existing_item.annotations, item.annotations))
+                    item = self._merge_items(existing_item, item, path=path)
                 else:
                     item = item.wrap(path=path, annotations=item.annotations)
 
@@ -412,6 +405,23 @@ def _lazy_image(item):
         # NOTE: avoid https://docs.python.org/3/faq/programming.html#why-do-lambdas-defined-in-a-loop-with-different-values-all-return-the-same-result
         return lambda: item.image
 
+    @classmethod
+    def _merge_items(cls, existing_item, current_item, path=None):
+        image = None
+        if existing_item.has_image and current_item.has_image:
+            if existing_item.image.has_data:
+                image = existing_item.image
+            elif current_item.image.has_data:
+                image = current_item.image
+        elif existing_item.has_image:
+            image = existing_item
+        elif current_item.has_image:
+            image = current_item.image
+
+        return existing_item.wrap(path=path,
+            image=image, annotations=cls._merge_anno(
+                existing_item.annotations, current_item.annotations))
+
     @staticmethod
     def _merge_anno(a, b):
         from itertools import chain
@@ -479,17 +489,10 @@ def __init__(self, project):
             for item in source:
                 existing_item = subsets[item.subset].items.get(item.id)
                 if existing_item is not None:
-                    image = None
-                    if existing_item.has_image:
-                        # TODO: think of image comparison
-                        image = self._lazy_image(existing_item)
-
                     path = existing_item.path
                     if item.path != path:
                         path = None # NOTE: move to our own dataset
-                    item = item.wrap(path=path,
-                        image=image, annotations=self._merge_anno(
-                            existing_item.annotations, item.annotations))
+                    item = self._merge_items(existing_item, item, path=path)
                 else:
                     s_config = config.sources[source_name]
                     if s_config and \
diff --git a/datumaro/datumaro/plugins/coco_format/converter.py b/datumaro/datumaro/plugins/coco_format/converter.py
index e6895e96de0..02d5c6e3f91 100644
--- a/datumaro/datumaro/plugins/coco_format/converter.py
+++ b/datumaro/datumaro/plugins/coco_format/converter.py
@@ -69,7 +69,7 @@ def is_empty(self):
 
     def save_image_info(self, item, filename):
         if item.has_image:
-            h, w = item.image.shape[:2]
+            h, w = item.image.size
         else:
             h = 0
             w = 0
@@ -130,7 +130,7 @@ def save_categories(self, dataset):
         pass
 
     def save_annotations(self, item):
-        for ann in item.annotations:
+        for ann_idx, ann in enumerate(item.annotations):
             if ann.type != AnnotationType.caption:
                 continue
 
@@ -144,7 +144,8 @@ def save_annotations(self, item):
                 try:
                     elem['score'] = float(ann.attributes['score'])
                 except Exception as e:
-                    log.warning("Failed to convert attribute 'score': %e" % e)
+                    log.warning("Item '%s', ann #%s: failed to convert "
+                        "attribute 'score': %e" % (item.id, ann_idx, e))
 
             self.annotations.append(elem)
 
@@ -293,10 +294,10 @@ def save_annotations(self, item):
             return
 
         if not item.has_image:
-            log.warn("Skipping writing instances for "
-                "item '%s' as it has no image info" % item.id)
+            log.warn("Item '%s': skipping writing instances "
+                "since no image info available" % item.id)
             return
-        h, w, _ = item.image.shape
+        h, w = item.image.size
         instances = [self.find_instance_parts(i, w, h) for i in instances]
 
         if self._context._crop_covered:
@@ -319,7 +320,7 @@ def convert_instance(self, instance, item):
         area = 0
         if segmentation:
             if item.has_image:
-                h, w, _ = item.image.shape
+                h, w = item.image.size
             else:
                 # NOTE: here we can guess the image size as
                 # it is only needed for the area computation
@@ -350,7 +351,8 @@ def convert_instance(self, instance, item):
             try:
                 elem['score'] = float(ann.attributes['score'])
             except Exception as e:
-                log.warning("Failed to convert attribute 'score': %e" % e)
+                log.warning("Item '%s': failed to convert attribute "
+                    "'score': %e" % (item.id, e))
 
         return elem
 
@@ -459,7 +461,8 @@ def save_annotations(self, item):
                 try:
                     elem['score'] = float(ann.attributes['score'])
                 except Exception as e:
-                    log.warning("Failed to convert attribute 'score': %e" % e)
+                    log.warning("Item '%s': failed to convert attribute "
+                        "'score': %e" % (item.id, e))
 
             self.annotations.append(elem)
 
@@ -519,8 +522,13 @@ def make_task_converters(self):
         }
 
     def save_image(self, item, filename):
+        image = item.image.data
+        if image is None:
+            log.warning("Item '%s' has no image" % item.id)
+            return
+
         path = osp.join(self._images_dir, filename)
-        save_image(path, item.image)
+        save_image(path, image)
 
         return path
 
@@ -549,7 +557,7 @@ def convert(self):
                     if item.has_image:
                         self.save_image(item, filename)
                     else:
-                        log.debug("Item '%s' has no image" % item.id)
+                        log.debug("Item '%s' has no image info" % item.id)
                 for task_conv in task_converters.values():
                     task_conv.save_image_info(item, filename)
                     task_conv.save_annotations(item)
diff --git a/datumaro/datumaro/plugins/cvat_format/converter.py b/datumaro/datumaro/plugins/cvat_format/converter.py
index 948916fee65..edd72a1cab8 100644
--- a/datumaro/datumaro/plugins/cvat_format/converter.py
+++ b/datumaro/datumaro/plugins/cvat_format/converter.py
@@ -162,30 +162,37 @@ def write(self):
                 if item.has_image:
                     self._save_image(item)
                 else:
-                    log.debug("Item '%s' has no image" % item.id)
+                    log.debug("Item '%s' has no image info" % item.id)
             self._write_item(item)
 
         self._writer.close_root()
 
     def _save_image(self, item):
-        image = item.image
+        image = item.image.data
         if image is None:
+            log.warning("Item '%s' has no image" % item.id)
             return
 
-        image_path = osp.join(self._context._images_dir,
-            str(item.id) + CvatPath.IMAGE_EXT)
+        file_name = item.name
+        if not file_name:
+            file_name = str(item.id)
+        file_name += CvatPath.IMAGE_EXT
+        image_path = osp.join(self._context._images_dir, file_name)
         save_image(image_path, image)
 
     def _write_item(self, item):
-        h, w = 0, 0
-        if item.has_image:
-            h, w = item.image.shape[:2]
-        self._writer.open_image(OrderedDict([
+        image_name = item.name
+        if not image_name:
+            image_name = str(item.id)
+        image_info = OrderedDict([
             ("id", str(item.id)),
-            ("name", str(item.id)),
-            ("width", str(w)),
-            ("height", str(h))
-        ]))
+            ("name", image_name),
+        ])
+        if item.has_image:
+            h, w = item.image.size
+            image_info["width"] = str(w)
+            image_info["height"] = str(h)
+        self._writer.open_image(image_info)
 
         for ann in item.annotations:
             if ann.type in {AnnotationType.points, AnnotationType.polyline,
diff --git a/datumaro/datumaro/plugins/cvat_format/extractor.py b/datumaro/datumaro/plugins/cvat_format/extractor.py
index 2ef8d8280d7..04e0d2f7c5c 100644
--- a/datumaro/datumaro/plugins/cvat_format/extractor.py
+++ b/datumaro/datumaro/plugins/cvat_format/extractor.py
@@ -12,7 +12,7 @@
     AnnotationType, Points, Polygon, PolyLine, Bbox,
     LabelCategories
 )
-from datumaro.util.image import lazy_image
+from datumaro.util.image import Image
 
 from .format import CvatPath
 
@@ -67,7 +67,7 @@ def _parse(cls, path):
         context = ET.ElementTree.iterparse(path, events=("start", "end"))
         context = iter(context)
 
-        categories = cls._parse_meta(context)
+        categories, frame_size = cls._parse_meta(context)
 
         items = OrderedDict()
 
@@ -81,11 +81,15 @@ def _parse(cls, path):
                         'id': el.attrib.get('id'),
                         'label': el.attrib.get('label'),
                         'group': int(el.attrib.get('group_id', 0)),
+                        'height': frame_size[0],
+                        'width': frame_size[1],
                     }
                 elif el.tag == 'image':
                     image = {
                         'name': el.attrib.get('name'),
                         'frame': el.attrib['id'],
+                        'width': el.attrib.get('width'),
+                        'height': el.attrib.get('height'),
                     }
                 elif el.tag in cls._SUPPORTED_SHAPES and (track or image):
                     shape = {
@@ -130,10 +134,7 @@ def _parse(cls, path):
                         for pair in el.attrib['points'].split(';'):
                             shape['points'].extend(map(float, pair.split(',')))
 
-                    frame_desc = items.get(shape['frame'], {
-                        'name': shape.get('name'),
-                        'annotations': [],
-                    })
+                    frame_desc = items.get(shape['frame'], {'annotations': []})
                     frame_desc['annotations'].append(
                         cls._parse_ann(shape, categories))
                     items[shape['frame']] = frame_desc
@@ -142,6 +143,13 @@ def _parse(cls, path):
                 elif el.tag == 'track':
                     track = None
                 elif el.tag == 'image':
+                    frame_desc = items.get(image['frame'], {'annotations': []})
+                    frame_desc.update({
+                        'name': image.get('name'),
+                        'height': image.get('height'),
+                        'width': image.get('width'),
+                    })
+                    items[image['frame']] = frame_desc
                     image = None
                 el.clear()
 
@@ -155,6 +163,7 @@ def _parse_meta(context):
 
         categories = {}
 
+        frame_size = None
         has_z_order = False
         mode = 'annotation'
         labels = OrderedDict()
@@ -183,6 +192,10 @@ def consumed(expected_state, tag):
                 if accepted('annotations', 'meta'): pass
                 elif accepted('meta', 'task'): pass
                 elif accepted('task', 'z_order'): pass
+                elif accepted('task', 'original_size'):
+                    frame_size = [None, None]
+                elif accepted('original_size', 'height', next_state='frame_height'): pass
+                elif accepted('original_size', 'width', next_state='frame_width'): pass
                 elif accepted('task', 'labels'): pass
                 elif accepted('labels', 'label'):
                     label = { 'name': None, 'attributes': set() }
@@ -202,6 +215,11 @@ def consumed(expected_state, tag):
                 elif consumed('task', 'task'): pass
                 elif consumed('z_order', 'z_order'):
                     has_z_order = (el.text == 'True')
+                elif consumed('original_size', 'original_size'): pass
+                elif consumed('frame_height', 'height'):
+                    frame_size[0] = int(el.text)
+                elif consumed('frame_width', 'width'):
+                    frame_size[1] = int(el.text)
                 elif consumed('label_name', 'name'):
                     label['name'] = el.text
                 elif consumed('attr_name', 'name'):
@@ -231,7 +249,7 @@ def consumed(expected_state, tag):
 
         categories[AnnotationType.label] = label_cat
 
-        return categories
+        return categories, frame_size
 
     @classmethod
     def _parse_ann(cls, ann, categories):
@@ -278,26 +296,37 @@ def _parse_ann(cls, ann, categories):
 
     def _load_items(self, parsed):
         for item_id, item_desc in parsed.items():
+            image = None
             file_name = item_desc.get('name')
             if not file_name:
                 file_name = item_id
-            image = self._find_image(file_name)
+            image_path = self._find_image(file_name)
+            image_size = (item_desc.get('height'), item_desc.get('width'))
+            if all(image_size):
+                image_size = (int(image_size[0]), int(image_size[1]))
+            else:
+                image_size = None
+            if image_path:
+                image = Image(path=image_path, size=image_size)
+            elif image_size:
+                image = Image(size=image_size)
 
             parsed[item_id] = DatasetItem(id=item_id, subset=self._subset,
-                image=image, annotations=item_desc.get('annotations', None))
+                image=image, annotations=item_desc.get('annotations'))
         return parsed
 
     def _find_image(self, file_name):
-        search_paths = [
-            osp.join(osp.dirname(self._path), file_name)
-        ]
+        search_paths = []
         if self._images_dir:
             search_paths += [
                 osp.join(self._images_dir, file_name),
                 osp.join(self._images_dir, self._subset or DEFAULT_SUBSET_NAME,
                     file_name),
             ]
+        search_paths += [
+            osp.join(osp.dirname(self._path), file_name)
+        ]
         for image_path in search_paths:
             if osp.isfile(image_path):
-                return lazy_image(image_path)
-        return None
\ No newline at end of file
+                return image_path
+        return None
diff --git a/datumaro/datumaro/plugins/datumaro_format/converter.py b/datumaro/datumaro/plugins/datumaro_format/converter.py
index 32f31e4acd3..256e235ae3b 100644
--- a/datumaro/datumaro/plugins/datumaro_format/converter.py
+++ b/datumaro/datumaro/plugins/datumaro_format/converter.py
@@ -58,6 +58,11 @@ def write_item(self, item):
         }
         if item.path:
             item_desc['path'] = item.path
+        if item.has_image:
+            item_desc['image'] = {
+                'size': item.image.size,
+                'path': item.image.path,
+            }
         self.items.append(item_desc)
 
         for ann in item.annotations:
@@ -226,7 +231,7 @@ def _convert_points_categories(self, obj):
         return converted
 
 class _Converter:
-    def __init__(self, extractor, save_dir, save_images=False,):
+    def __init__(self, extractor, save_dir, save_images=False):
         self._extractor = extractor
         self._save_dir = save_dir
         self._save_images = save_images
@@ -258,14 +263,15 @@ def convert(self):
             writer = subsets[subset]
 
             if self._save_images:
-                self._save_image(item)
+                if item.has_image:
+                    self._save_image(item)
             writer.write_item(item)
 
         for subset, writer in subsets.items():
             writer.write(annotations_dir)
 
     def _save_image(self, item):
-        image = item.image
+        image = item.image.data
         if image is None:
             return
 
diff --git a/datumaro/datumaro/plugins/datumaro_format/extractor.py b/datumaro/datumaro/plugins/datumaro_format/extractor.py
index 04b9af0b948..a666629ac2d 100644
--- a/datumaro/datumaro/plugins/datumaro_format/extractor.py
+++ b/datumaro/datumaro/plugins/datumaro_format/extractor.py
@@ -12,7 +12,7 @@
     AnnotationType, Label, Mask, Points, Polygon, PolyLine, Bbox, Caption,
     LabelCategories, MaskCategories, PointsCategories
 )
-from datumaro.util.image import lazy_image
+from datumaro.util.image import Image
 from datumaro.util.mask_tools import lazy_mask
 
 from .format import DatumaroPath
@@ -93,11 +93,20 @@ def _load_items(self, parsed):
         items = []
         for item_desc in parsed['items']:
             item_id = item_desc['id']
+
             image = None
-            image_path = osp.join(self._path, DatumaroPath.IMAGES_DIR,
-                item_id + DatumaroPath.IMAGE_EXT)
-            if osp.exists(image_path):
-                image = lazy_image(image_path)
+            image_info = item_desc.get('image', {})
+            if image_info.get('path'):
+                image_path = osp.join(self._path, DatumaroPath.IMAGES_DIR,
+                    image_info['path'])
+            else:
+                image_path = osp.join(self._path, DatumaroPath.IMAGES_DIR,
+                    item_id + DatumaroPath.IMAGE_EXT)
+            image_size = image_info.get('size')
+            if osp.isfile(image_path):
+                image = Image(path=image_path, size=image_size)
+            elif image_size:
+                image = Image(size=image_size)
 
             annotations = self._load_annotations(item_desc)
 
diff --git a/datumaro/datumaro/plugins/image_dir.py b/datumaro/datumaro/plugins/image_dir.py
index c2e0e687bce..6bc5e215656 100644
--- a/datumaro/datumaro/plugins/image_dir.py
+++ b/datumaro/datumaro/plugins/image_dir.py
@@ -8,7 +8,8 @@
 import os.path as osp
 
 from datumaro.components.extractor import DatasetItem, SourceExtractor, Importer
-from datumaro.util.image import lazy_image
+from datumaro.components.converter import Converter
+from datumaro.util.image import lazy_image, save_image
 
 
 class ImageDirImporter(Importer):
@@ -73,3 +74,16 @@ def _is_image(self, path):
             if osp.isfile(path) and path.endswith(ext):
                 return True
         return False
+
+
+class ImageDirConverter(Converter):
+    def __call__(self, extractor, save_dir):
+        os.makedirs(save_dir, exist_ok=True)
+
+        for item in extractor:
+            if item.has_image and item.image.has_data:
+                file_name = item.name
+                if not file_name:
+                    file_name = str(item.id)
+                file_name += '.jpg'
+                save_image(osp.join(save_dir, file_name), item.image.data)
\ No newline at end of file
diff --git a/datumaro/datumaro/plugins/tf_detection_api_format/converter.py b/datumaro/datumaro/plugins/tf_detection_api_format/converter.py
index 8761807906f..340492638fe 100644
--- a/datumaro/datumaro/plugins/tf_detection_api_format/converter.py
+++ b/datumaro/datumaro/plugins/tf_detection_api_format/converter.py
@@ -10,7 +10,9 @@
 import os.path as osp
 import string
 
-from datumaro.components.extractor import AnnotationType, DEFAULT_SUBSET_NAME
+from datumaro.components.extractor import (AnnotationType, DEFAULT_SUBSET_NAME,
+    LabelCategories
+)
 from datumaro.components.converter import Converter
 from datumaro.components.cli_plugin import CliPlugin
 from datumaro.util.image import encode_image
@@ -49,26 +51,30 @@ def float_list_feature(value):
     }
 
     if not item.has_image:
-        raise Exception(
-            "Failed to export dataset item '%s': item has no image" % item.id)
-    height, width = item.image.shape[:2]
+        raise Exception("Failed to export dataset item '%s': "
+            "item has no image info" % item.id)
+    height, width = item.image.size
 
     features.update({
         'image/height': int64_feature(height),
         'image/width': int64_feature(width),
     })
 
+    features.update({
+        'image/encoded': bytes_feature(b''),
+        'image/format': bytes_feature(b'')
+    })
     if save_images:
-        if item.has_image:
+        if item.has_image and item.image.has_data:
             fmt = DetectionApiPath.IMAGE_FORMAT
-            buffer = encode_image(item.image, DetectionApiPath.IMAGE_EXT)
+            buffer = encode_image(item.image.data, DetectionApiPath.IMAGE_EXT)
 
             features.update({
                 'image/encoded': bytes_feature(buffer),
                 'image/format': bytes_feature(fmt.encode('utf-8')),
             })
         else:
-            log.debug("Item '%s' has no image" % item.id)
+            log.warning("Item '%s' has no image" % item.id)
 
     xmins = [] # List of normalized left x coordinates in bounding box (1 per box)
     xmaxs = [] # List of normalized right x coordinates in bounding box (1 per box)
@@ -130,7 +136,8 @@ def __call__(self, extractor, save_dir):
                 subset_name = DEFAULT_SUBSET_NAME
                 subset = extractor
 
-            label_categories = subset.categories()[AnnotationType.label]
+            label_categories = subset.categories().get(AnnotationType.label,
+                LabelCategories())
             get_label = lambda label_id: label_categories.items[label_id].name \
                 if label_id is not None else ''
             label_ids = OrderedDict((label.name, 1 + idx)
diff --git a/datumaro/datumaro/plugins/tf_detection_api_format/extractor.py b/datumaro/datumaro/plugins/tf_detection_api_format/extractor.py
index 0d86d0f3324..7e571321d16 100644
--- a/datumaro/datumaro/plugins/tf_detection_api_format/extractor.py
+++ b/datumaro/datumaro/plugins/tf_detection_api_format/extractor.py
@@ -12,7 +12,7 @@
     DEFAULT_SUBSET_NAME, DatasetItem,
     AnnotationType, Bbox, LabelCategories
 )
-from datumaro.util.image import lazy_image, decode_image
+from datumaro.util.image import Image, decode_image, lazy_image
 from datumaro.util.tf_util import import_tf as _import_tf
 
 from .format import DetectionApiPath
@@ -174,13 +174,20 @@ def _parse_tfrecord_file(cls, filepath, subset_name, images_dir):
                     label=dataset_labels.get(label, None), id=index
                 ))
 
-            image = None
-            if image is None and frame_image and frame_format:
-                image = lazy_image(frame_image, loader=decode_image)
-            if image is None and frame_filename and images_dir:
+            image_params = {}
+            if frame_height and frame_width:
+                image_params['size'] = (frame_height, frame_width)
+
+            if frame_image and frame_format:
+                image_params['data'] = lazy_image(frame_image, decode_image)
+            elif frame_filename and images_dir:
                 image_path = osp.join(images_dir, frame_filename)
                 if osp.exists(image_path):
-                    image = lazy_image(image_path)
+                    image_params['path'] = image_path
+
+            image=None
+            if image_params:
+                image = Image(**image_params)
 
             dataset_items.append(DatasetItem(id=item_id, subset=subset_name,
                 image=image, annotations=annotations))
diff --git a/datumaro/datumaro/plugins/transforms.py b/datumaro/datumaro/plugins/transforms.py
index 6c418387b76..44ec703b478 100644
--- a/datumaro/datumaro/plugins/transforms.py
+++ b/datumaro/datumaro/plugins/transforms.py
@@ -28,7 +28,7 @@ def transform_item(self, item):
 
         if not item.has_image:
             raise Exception("Image info is required for this transform")
-        h, w = item.image.shape[:2]
+        h, w = item.image.size
         segments = self.crop_segments(segments, w, h)
 
         annotations += segments
@@ -107,7 +107,7 @@ def transform_item(self, item):
 
         if not item.has_image:
             raise Exception("Image info is required for this transform")
-        h, w = item.image.shape[:2]
+        h, w = item.image.size
         instances = self.find_instances(segments)
         segments = [self.merge_segments(i, w, h, self._include_polygons)
             for i in instances]
@@ -196,7 +196,7 @@ def transform_item(self, item):
             if ann.type == AnnotationType.polygon:
                 if not item.has_image:
                     raise Exception("Image info is required for this transform")
-                h, w = item.image.shape[:2]
+                h, w = item.image.size
                 annotations.append(self.convert_polygon(ann, h, w))
             else:
                 annotations.append(ann)
diff --git a/datumaro/datumaro/plugins/voc_format/converter.py b/datumaro/datumaro/plugins/voc_format/converter.py
index 5de9602ff73..79bacd6afbc 100644
--- a/datumaro/datumaro/plugins/voc_format/converter.py
+++ b/datumaro/datumaro/plugins/voc_format/converter.py
@@ -136,10 +136,10 @@ def save_subsets(self):
                 log.debug("Converting item '%s'", item.id)
 
                 if self._save_images:
-                    if item.has_image:
+                    if item.has_image and item.image.has_data:
                         save_image(osp.join(self._images_dir,
                                 item.id + VocPath.IMAGE_EXT),
-                            item.image)
+                            item.image.data)
                     else:
                         log.debug("Item '%s' has no image" % item.id)
 
@@ -467,7 +467,8 @@ def _get_actions(self, label):
     def _make_label_id_map(self):
         source_labels = {
             id: label.name for id, label in
-            enumerate(self._extractor.categories()[AnnotationType.label].items)
+            enumerate(self._extractor.categories().get(
+                AnnotationType.label, LabelCategories()).items)
         }
         target_labels = {
             label.name: id for id, label in
diff --git a/datumaro/datumaro/plugins/yolo_format/converter.py b/datumaro/datumaro/plugins/yolo_format/converter.py
index d50c2a0aab2..59a46a6260c 100644
--- a/datumaro/datumaro/plugins/yolo_format/converter.py
+++ b/datumaro/datumaro/plugins/yolo_format/converter.py
@@ -80,12 +80,16 @@ def __call__(self, extractor, save_dir):
                     osp.basename(subset_dir), image_name)
 
                 if self._save_images:
-                    if item.has_image:
-                        save_image(osp.join(subset_dir, image_name), item.image)
+                    if item.has_image and item.image.has_data:
+                        save_image(osp.join(subset_dir, image_name),
+                            item.image.data)
                     else:
-                        log.debug("Item '%s' has no images" % item.id)
+                        log.warning("Item '%s' has no image" % item.id)
 
-                height, width = item.image.shape[:2]
+                if not item.has_image:
+                    raise Exception("Failed to export item '%s': "
+                        "item has no image info" % item.id)
+                height, width = item.image.size
 
                 yolo_annotation = ''
                 for bbox in item.annotations:
diff --git a/datumaro/tests/test_image_dir_format.py b/datumaro/tests/test_image_dir_format.py
index f7875df0630..30dd05b1c43 100644
--- a/datumaro/tests/test_image_dir_format.py
+++ b/datumaro/tests/test_image_dir_format.py
@@ -1,12 +1,11 @@
 import numpy as np
-import os.path as osp
 
 from unittest import TestCase
 
 from datumaro.components.project import Project
 from datumaro.components.extractor import Extractor, DatasetItem
+from datumaro.plugins.image_dir import ImageDirConverter
 from datumaro.util.test_utils import TestDir, compare_datasets
-from datumaro.util.image import save_image
 
 
 class ImageDirFormatTest(TestCase):
@@ -21,9 +20,7 @@ def test_can_load(self):
         with TestDir() as test_dir:
             source_dataset = self.TestExtractor()
 
-            for item in source_dataset:
-                save_image(osp.join(test_dir, '%s.jpg' % item.id),
-                    item.image.data)
+            ImageDirConverter()(source_dataset, save_dir=test_dir)
 
             project = Project.import_from(test_dir, 'image_dir')
             parsed_dataset = project.make_dataset()

From fc94473ebf784c265cbca678f2b0f7ebf46b839b Mon Sep 17 00:00:00 2001
From: Zhiltsov Max <zhiltsov.max35@gmail.com>
Date: Wed, 12 Feb 2020 18:57:08 +0300
Subject: [PATCH 04/25] Fix cli

---
 .../datumaro/cli/contexts/project/__init__.py | 21 +++++++++++--------
 .../datumaro/plugins/voc_format/converter.py  |  6 +++++-
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/datumaro/datumaro/cli/contexts/project/__init__.py b/datumaro/datumaro/cli/contexts/project/__init__.py
index 15bc3a0c05e..3ca79e04782 100644
--- a/datumaro/datumaro/cli/contexts/project/__init__.py
+++ b/datumaro/datumaro/cli/contexts/project/__init__.py
@@ -156,16 +156,17 @@ def import_command(args):
     if project_name is None:
         project_name = osp.basename(project_dir)
 
-    extra_args = {}
     try:
         env = Environment()
         importer = env.make_importer(args.format)
-        if hasattr(importer, 'from_cmdline'):
-            extra_args = importer.from_cmdline(args.extra_args)
     except KeyError:
         raise CliException("Importer for format '%s' is not found" % \
             args.format)
 
+    extra_args = {}
+    if hasattr(importer, 'from_cmdline'):
+        extra_args = importer.from_cmdline(args.extra_args)
+
     log.info("Importing project from '%s' as '%s'" % \
         (args.source, args.format))
 
@@ -293,13 +294,14 @@ def export_command(args):
 
     try:
         converter = project.env.converters.get(args.format)
-        if hasattr(converter, 'from_cmdline'):
-            extra_args = converter.from_cmdline(args.extra_args)
-            converter = converter(**extra_args)
     except KeyError:
         raise CliException("Converter for format '%s' is not found" % \
             args.format)
 
+    if hasattr(converter, 'from_cmdline'):
+        extra_args = converter.from_cmdline(args.extra_args)
+        converter = converter(**extra_args)
+
     filter_args = FilterModes.make_filter_args(args.filter_mode)
 
     log.info("Loading the project...")
@@ -559,14 +561,15 @@ def transform_command(args):
             (project.config.project_name, make_file_name(args.transform)))
     dst_dir = osp.abspath(dst_dir)
 
-    extra_args = {}
     try:
         transform = project.env.transforms.get(args.transform)
-        if hasattr(transform, 'from_cmdline'):
-            extra_args = transform.from_cmdline(args.extra_args)
     except KeyError:
         raise CliException("Transform '%s' is not found" % args.transform)
 
+    extra_args = {}
+    if hasattr(transform, 'from_cmdline'):
+        extra_args = transform.from_cmdline(args.extra_args)
+
     log.info("Loading the project...")
     dataset = project.make_dataset()
 
diff --git a/datumaro/datumaro/plugins/voc_format/converter.py b/datumaro/datumaro/plugins/voc_format/converter.py
index 79bacd6afbc..03fa806824d 100644
--- a/datumaro/datumaro/plugins/voc_format/converter.py
+++ b/datumaro/datumaro/plugins/voc_format/converter.py
@@ -511,7 +511,11 @@ def _split_tasks_string(s):
     def _get_labelmap(s):
         if osp.isfile(s):
             return s
-        return LabelmapType[s].name
+        try:
+            return LabelmapType[s].name
+        except KeyError:
+            import argparse
+            raise argparse.ArgumentTypeError()
 
     @classmethod
     def build_cmdline_parser(cls, **kwargs):

From c3c16024d6280578be2a6edc2197826577e858a6 Mon Sep 17 00:00:00 2001
From: Zhiltsov Max <zhiltsov.max35@gmail.com>
Date: Wed, 12 Feb 2020 19:10:44 +0300
Subject: [PATCH 05/25] Fix merge and voc converte

---
 datumaro/datumaro/components/project.py           | 2 +-
 datumaro/datumaro/plugins/voc_format/converter.py | 9 ++++++---
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/datumaro/datumaro/components/project.py b/datumaro/datumaro/components/project.py
index 7688fe996f9..89668869e12 100644
--- a/datumaro/datumaro/components/project.py
+++ b/datumaro/datumaro/components/project.py
@@ -414,7 +414,7 @@ def _merge_items(cls, existing_item, current_item, path=None):
             elif current_item.image.has_data:
                 image = current_item.image
         elif existing_item.has_image:
-            image = existing_item
+            image = existing_item.image
         elif current_item.has_image:
             image = current_item.image
 
diff --git a/datumaro/datumaro/plugins/voc_format/converter.py b/datumaro/datumaro/plugins/voc_format/converter.py
index 03fa806824d..c7ecbdee039 100644
--- a/datumaro/datumaro/plugins/voc_format/converter.py
+++ b/datumaro/datumaro/plugins/voc_format/converter.py
@@ -170,9 +170,12 @@ def save_subsets(self):
                     ET.SubElement(source_elem, 'image').text = 'Unknown'
 
                     if item.has_image:
-                        image_shape = item.image.shape
-                        h, w = image_shape[:2]
-                        c = 1 if len(image_shape) == 2 else image_shape[2]
+                        h, w = item.image.size
+                        if item.image.has_data:
+                            image_shape = item.image.data.shape
+                            c = 1 if len(image_shape) == 2 else image_shape[2]
+                        else:
+                            c = 3
                         size_elem = ET.SubElement(root_elem, 'size')
                         ET.SubElement(size_elem, 'width').text = str(w)
                         ET.SubElement(size_elem, 'height').text = str(h)

From ff57202fe827e865bf1d8e42ec30ebf174bff347 Mon Sep 17 00:00:00 2001
From: Zhiltsov Max <zhiltsov.max35@gmail.com>
Date: Wed, 12 Feb 2020 19:16:22 +0300
Subject: [PATCH 06/25] Update remote images extractor

---
 .../plugins/cvat_rest_api_task_images.py                 | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/cvat/apps/dataset_manager/export_templates/plugins/cvat_rest_api_task_images.py b/cvat/apps/dataset_manager/export_templates/plugins/cvat_rest_api_task_images.py
index c7cf8fbbd22..e073f3edf0e 100644
--- a/cvat/apps/dataset_manager/export_templates/plugins/cvat_rest_api_task_images.py
+++ b/cvat/apps/dataset_manager/export_templates/plugins/cvat_rest_api_task_images.py
@@ -13,7 +13,7 @@
     SchemaBuilder as _SchemaBuilder,
 )
 import datumaro.components.extractor as datumaro
-from datumaro.util.image import lazy_image, load_image
+from datumaro.util.image import lazy_image, load_image, Image
 
 from cvat.utils.cli.core import CLI as CVAT_CLI, CVAT_API_V1
 
@@ -103,8 +103,11 @@ def __init__(self, url):
         items = []
         for entry in image_list:
             item_id = entry['id']
-            item = datumaro.DatasetItem(
-                id=item_id, image=self._make_image_loader(item_id))
+            size = None
+            if entry.get('height') and entry.get('width'):
+                size = (entry['height'], entry['width'])
+            image = Image(data=self._make_image_loader(item_id), size=size)
+            item = datumaro.DatasetItem(id=item_id, image=image)
             items.append((item.id, item))
 
         items = sorted(items, key=lambda e: int(e[0]))

From 82c3a56b6b0909651d0f18f5f8a802e40ec7a5f3 Mon Sep 17 00:00:00 2001
From: Zhiltsov Max <zhiltsov.max35@gmail.com>
Date: Thu, 13 Feb 2020 13:13:14 +0300
Subject: [PATCH 07/25] Codacy

---
 datumaro/tests/test_yolo_format.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/datumaro/tests/test_yolo_format.py b/datumaro/tests/test_yolo_format.py
index bad4a5e4c18..3af08781c36 100644
--- a/datumaro/tests/test_yolo_format.py
+++ b/datumaro/tests/test_yolo_format.py
@@ -7,7 +7,6 @@
     AnnotationType, Bbox, LabelCategories,
 )
 from datumaro.plugins.yolo_format.importer import YoloImporter
-from datumaro.plugins.yolo_format.format import YoloPath
 from datumaro.plugins.yolo_format.converter import YoloConverter
 from datumaro.util.image import Image, save_image
 from datumaro.util.test_utils import TestDir, compare_datasets

From 986ce386941e44e9407903d65e34a1c95bcc0271 Mon Sep 17 00:00:00 2001
From: Zhiltsov Max <zhiltsov.max35@gmail.com>
Date: Mon, 17 Feb 2020 18:24:03 +0300
Subject: [PATCH 08/25] Remove item name, require path in Image

---
 datumaro/datumaro/components/extractor.py | 27 ++++++-----------------
 datumaro/datumaro/util/image.py           | 16 +++++++++-----
 2 files changed, 18 insertions(+), 25 deletions(-)

diff --git a/datumaro/datumaro/components/extractor.py b/datumaro/datumaro/components/extractor.py
index c1b23ef6903..9813f2900d8 100644
--- a/datumaro/datumaro/components/extractor.py
+++ b/datumaro/datumaro/components/extractor.py
@@ -576,18 +576,10 @@ def __eq__(self, other):
 
 class DatasetItem:
     # pylint: disable=redefined-builtin
-    def __init__(self, id, annotations=None,
-            subset=None, path=None, image=None, name=None):
+    def __init__(self, id=None, annotations=None,
+            subset=None, path=None, image=None):
         assert id is not None
-        id = str(id)
-        # TODO: if not isinstance(id, int):
-        #     id = int(id)
-        self._id = id
-
-        assert name is None or isinstance(name, str)
-        if not name:
-            name = str(id)
-        self._name = name
+        self._id = str(id)
 
         if subset is None:
             subset = ''
@@ -609,6 +601,8 @@ def __init__(self, id, annotations=None,
 
         if callable(image) or isinstance(image, np.ndarray):
             image = Image(data=image)
+        elif isinstance(image, str):
+            image = Image(path=image)
         assert image is None or isinstance(image, Image)
         self._image = image
     # pylint: enable=redefined-builtin
@@ -617,10 +611,6 @@ def __init__(self, id, annotations=None,
     def id(self):
         return self._id
 
-    @property
-    def name(self):
-        return self._name
-
     @property
     def subset(self):
         return self._subset
@@ -647,12 +637,12 @@ def __eq__(self, other):
         return \
             (self.id == other.id) and \
             (self.subset == other.subset) and \
-            (self.annotations == other.annotations) and \
             (self.path == other.path) and \
+            (self.annotations == other.annotations) and \
             (self.image == other.image)
 
     def wrap(item, **kwargs):
-        expected_args = {'id', 'annotations', 'subset', 'path', 'image', 'name'}
+        expected_args = {'id', 'annotations', 'subset', 'path', 'image'}
         for k in expected_args:
             if k not in kwargs:
                 kwargs[k] = getattr(item, k)
@@ -677,9 +667,6 @@ def categories(self):
     def select(self, pred):
         raise NotImplementedError()
 
-    def get(self, item_id, subset=None, path=None):
-        raise NotImplementedError()
-
 class _DatasetFilter:
     def __init__(self, iterable, predicate):
         self.iterable = iterable
diff --git a/datumaro/datumaro/util/image.py b/datumaro/datumaro/util/image.py
index 2d9a4d43326..712a4f789ea 100644
--- a/datumaro/datumaro/util/image.py
+++ b/datumaro/datumaro/util/image.py
@@ -7,6 +7,7 @@
 
 from io import BytesIO
 import numpy as np
+import os.path as osp
 
 from enum import Enum
 _IMAGE_BACKENDS = Enum('_IMAGE_BACKENDS', ['cv2', 'PIL'])
@@ -166,26 +167,31 @@ def __init__(self, data=None, path=None, loader=None, cache=None,
             size=None):
         assert size is None or len(size) == 2
         if size is not None:
-            assert 0 < size[0] and 0 < size[1], size
+            assert len(size) == 2 and 0 < size[0] and 0 < size[1], size
             size = tuple(size)
         else:
             size = None
         self._size = size # (H, W)
 
         assert path is None or isinstance(path, str)
-        if not path:
+        if path is None:
             path = ''
         self._path = path
 
-        assert any(e is not None for e in (data, path, loader, size)), "Image can not be empty"
-        if data is None and (path or loader is not None):
-            data = lazy_image(path, loader=loader, cache=cache)
+        assert data is not None or path, "Image can not be empty"
+        if data is None and path:
+            if osp.isfile(path):
+                data = lazy_image(path, loader=loader, cache=cache)
         self._data = data
 
     @property
     def path(self):
         return self._path
 
+    @property
+    def filename(self):
+        return osp.basename(self._path)
+
     @property
     def data(self):
         if callable(self._data):

From a7824de17e8e36385f1ab34f22d95d1114d9a179 Mon Sep 17 00:00:00 2001
From: Zhiltsov Max <zhiltsov.max35@gmail.com>
Date: Mon, 17 Feb 2020 18:24:35 +0300
Subject: [PATCH 09/25] Merge images of dataset items

---
 datumaro/datumaro/components/project.py | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/datumaro/datumaro/components/project.py b/datumaro/datumaro/components/project.py
index 89668869e12..881bb3dc962 100644
--- a/datumaro/datumaro/components/project.py
+++ b/datumaro/datumaro/components/project.py
@@ -411,11 +411,22 @@ def _merge_items(cls, existing_item, current_item, path=None):
         if existing_item.has_image and current_item.has_image:
             if existing_item.image.has_data:
                 image = existing_item.image
-            elif current_item.image.has_data:
+            else:
                 image = current_item.image
+
+            if existing_item.image.path != current_item.image.path:
+                if not existing_item.image.path:
+                    image._path = current_item.image.path
+
+            if all([existing_item.image._size, current_item.image._size]):
+                assert existing_item.image._size == current_item.image._size, "Image info differs for item '%s'" % item.id
+            elif existing_item.image._size:
+                image._size = existing_item.image._size
+            else:
+                image._size = current_item.image._size
         elif existing_item.has_image:
             image = existing_item.image
-        elif current_item.has_image:
+        else:
             image = current_item.image
 
         return existing_item.wrap(path=path,

From 5a8677f4468de4ee98bb58ac15215c7371275cbe Mon Sep 17 00:00:00 2001
From: Zhiltsov Max <zhiltsov.max35@gmail.com>
Date: Mon, 17 Feb 2020 18:25:28 +0300
Subject: [PATCH 10/25] Update tests

---
 datumaro/tests/test_images.py  | 8 ++++----
 datumaro/tests/test_project.py | 3 ++-
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/datumaro/tests/test_images.py b/datumaro/tests/test_images.py
index 5f6dcb48a38..e77d73d4fac 100644
--- a/datumaro/tests/test_images.py
+++ b/datumaro/tests/test_images.py
@@ -47,10 +47,10 @@ def test_global_cache_is_accessible(self):
 
 class ImageTest(TestCase):
     def test_lazy_image_shape(self):
-        loader = lambda _: np.ones((5, 6, 7))
+        data = np.ones((5, 6, 7))
 
-        image_lazy = Image(loader=loader, size=(2, 4))
-        image_eager = Image(loader=loader)
+        image_lazy = Image(data=data, size=(2, 4))
+        image_eager = Image(data=data)
 
         self.assertEqual((2, 4), image_lazy.size)
         self.assertEqual((5, 6), image_eager.size)
@@ -66,10 +66,10 @@ def test_ctors():
                 { 'data': image },
                 { 'data': image, 'path': path },
                 { 'data': image, 'path': path, 'size': (2, 4) },
+                { 'data': image, 'path': path, 'loader': load_image, 'size': (2, 4) },
                 { 'path': path },
                 { 'path': path, 'loader': load_image },
                 { 'path': path, 'size': (2, 4) },
-                { 'size': (2, 4) },
             ]:
                 img = Image(**args)
                 # pylint: disable=pointless-statement
diff --git a/datumaro/tests/test_project.py b/datumaro/tests/test_project.py
index 1a0038d02bc..5b3c47f1d69 100644
--- a/datumaro/tests/test_project.py
+++ b/datumaro/tests/test_project.py
@@ -499,7 +499,7 @@ def test_ctor_requires_id(self):
             # pylint: disable=no-value-for-parameter
             DatasetItem()
             # pylint: enable=no-value-for-parameter
-        except TypeError:
+        except AssertionError:
             has_error = True
 
         self.assertTrue(has_error)
@@ -508,6 +508,7 @@ def test_ctor_requires_id(self):
     def test_ctors_with_image():
         for args in [
             { 'id': 0, 'image': None },
+            { 'id': 0, 'image': 'path.jpg' },
             { 'id': 0, 'image': np.array([1, 2, 3]) },
             { 'id': 0, 'image': lambda f: np.array([1, 2, 3]) },
             { 'id': 0, 'image': Image(data=np.array([1, 2, 3])) },

From d198fe9e3a5b0adffd4f3ed39a4feabd09791fa1 Mon Sep 17 00:00:00 2001
From: Zhiltsov Max <zhiltsov.max35@gmail.com>
Date: Mon, 17 Feb 2020 18:25:58 +0300
Subject: [PATCH 11/25] Add image dir converter

---
 datumaro/datumaro/plugins/image_dir.py | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/datumaro/datumaro/plugins/image_dir.py b/datumaro/datumaro/plugins/image_dir.py
index 6bc5e215656..c719c546f52 100644
--- a/datumaro/datumaro/plugins/image_dir.py
+++ b/datumaro/datumaro/plugins/image_dir.py
@@ -9,7 +9,7 @@
 
 from datumaro.components.extractor import DatasetItem, SourceExtractor, Importer
 from datumaro.components.converter import Converter
-from datumaro.util.image import lazy_image, save_image
+from datumaro.util.image import save_image
 
 
 class ImageDirImporter(Importer):
@@ -45,7 +45,7 @@ def __init__(self, url):
             path = osp.join(url, name)
             if self._is_image(path):
                 item_id = osp.splitext(name)[0]
-                item = DatasetItem(id=item_id, image=lazy_image(path))
+                item = DatasetItem(id=item_id, image=path)
                 items.append((item.id, item))
 
         items = sorted(items, key=lambda e: e[0])
@@ -82,8 +82,10 @@ def __call__(self, extractor, save_dir):
 
         for item in extractor:
             if item.has_image and item.image.has_data:
-                file_name = item.name
-                if not file_name:
-                    file_name = str(item.id)
-                file_name += '.jpg'
-                save_image(osp.join(save_dir, file_name), item.image.data)
\ No newline at end of file
+                filename = item.image.filename
+                if filename:
+                    filename = osp.splitext(filename)[0]
+                else:
+                    filename = item.id
+                filename += '.jpg'
+                save_image(osp.join(save_dir, filename), item.image.data)
\ No newline at end of file

From a7f61982ddc77ec94413b57751081724c34c31b6 Mon Sep 17 00:00:00 2001
From: Zhiltsov Max <zhiltsov.max35@gmail.com>
Date: Mon, 17 Feb 2020 18:28:37 +0300
Subject: [PATCH 12/25] Update Datumaro format

---
 .../plugins/datumaro_format/converter.py      | 29 ++++++++++++-------
 .../plugins/datumaro_format/extractor.py      | 13 ++-------
 datumaro/tests/test_datumaro_format.py        |  2 +-
 3 files changed, 22 insertions(+), 22 deletions(-)

diff --git a/datumaro/datumaro/plugins/datumaro_format/converter.py b/datumaro/datumaro/plugins/datumaro_format/converter.py
index 256e235ae3b..fe239fe68bf 100644
--- a/datumaro/datumaro/plugins/datumaro_format/converter.py
+++ b/datumaro/datumaro/plugins/datumaro_format/converter.py
@@ -30,9 +30,9 @@ def _cast(value, type_conv, default=None):
         return default
 
 class _SubsetWriter:
-    def __init__(self, name, converter):
+    def __init__(self, name, context):
         self._name = name
-        self._converter = converter
+        self._context = context
 
         self._data = {
             'info': {},
@@ -59,9 +59,13 @@ def write_item(self, item):
         if item.path:
             item_desc['path'] = item.path
         if item.has_image:
+            path = item.image.path
+            if self._context._save_images:
+                path = self._context._save_image(item)
+
             item_desc['image'] = {
                 'size': item.image.size,
-                'path': item.image.path,
+                'path': path,
             }
         self.items.append(item_desc)
 
@@ -128,7 +132,7 @@ def _save_mask(self, mask):
         self._next_mask_id += 1
 
         filename = '%d%s' % (mask_id, DatumaroPath.MASK_EXT)
-        masks_dir = osp.join(self._converter._annotations_dir,
+        masks_dir = osp.join(self._context._annotations_dir,
             DatumaroPath.MASKS_DIR)
         os.makedirs(masks_dir, exist_ok=True)
         path = osp.join(masks_dir, filename)
@@ -262,9 +266,6 @@ def convert(self):
                 subset = DEFAULT_SUBSET_NAME
             writer = subsets[subset]
 
-            if self._save_images:
-                if item.has_image:
-                    self._save_image(item)
             writer.write_item(item)
 
         for subset, writer in subsets.items():
@@ -273,11 +274,17 @@ def convert(self):
     def _save_image(self, item):
         image = item.image.data
         if image is None:
-            return
-
-        image_path = osp.join(self._images_dir,
-            str(item.id) + DatumaroPath.IMAGE_EXT)
+            return ''
+
+        filename = item.image.filename
+        if filename:
+            filename = osp.splitext(filename)[0]
+        else:
+            filename = item.id
+        filename += DatumaroPath.IMAGE_EXT
+        image_path = osp.join(self._images_dir, filename)
         save_image(image_path, image)
+        return filename
 
 class DatumaroConverter(Converter, CliPlugin):
     @classmethod
diff --git a/datumaro/datumaro/plugins/datumaro_format/extractor.py b/datumaro/datumaro/plugins/datumaro_format/extractor.py
index a666629ac2d..361df4d8466 100644
--- a/datumaro/datumaro/plugins/datumaro_format/extractor.py
+++ b/datumaro/datumaro/plugins/datumaro_format/extractor.py
@@ -96,17 +96,10 @@ def _load_items(self, parsed):
 
             image = None
             image_info = item_desc.get('image', {})
-            if image_info.get('path'):
+            if image_info:
                 image_path = osp.join(self._path, DatumaroPath.IMAGES_DIR,
-                    image_info['path'])
-            else:
-                image_path = osp.join(self._path, DatumaroPath.IMAGES_DIR,
-                    item_id + DatumaroPath.IMAGE_EXT)
-            image_size = image_info.get('size')
-            if osp.isfile(image_path):
-                image = Image(path=image_path, size=image_size)
-            elif image_size:
-                image = Image(size=image_size)
+                    image_info.get('path', '')) # relative or absolute fits
+                image = Image(path=image_path, size=image_info.get('size'))
 
             annotations = self._load_annotations(item_desc)
 
diff --git a/datumaro/tests/test_datumaro_format.py b/datumaro/tests/test_datumaro_format.py
index 0bc3911af23..e17da2a7b95 100644
--- a/datumaro/tests/test_datumaro_format.py
+++ b/datumaro/tests/test_datumaro_format.py
@@ -49,7 +49,7 @@ def __iter__(self):
                 DatasetItem(id=42, subset='test'),
 
                 DatasetItem(id=42),
-                DatasetItem(id=43, image=Image(size=(2, 4))),
+                DatasetItem(id=43, image=Image(path='1/b/c.qq', size=(2, 4))),
             ])
 
         def categories(self):

From 8ada47e7e503084f08de4e9e241bd935ef8c6f53 Mon Sep 17 00:00:00 2001
From: Zhiltsov Max <zhiltsov.max35@gmail.com>
Date: Mon, 17 Feb 2020 18:29:38 +0300
Subject: [PATCH 13/25] Update COCO format with image info

---
 .../datumaro/plugins/coco_format/converter.py | 47 +++++++++++++------
 .../datumaro/plugins/coco_format/extractor.py | 15 ++++--
 datumaro/tests/test_coco_format.py            |  2 +-
 3 files changed, 45 insertions(+), 19 deletions(-)

diff --git a/datumaro/datumaro/plugins/coco_format/converter.py b/datumaro/datumaro/plugins/coco_format/converter.py
index 02d5c6e3f91..4af56b75ff6 100644
--- a/datumaro/datumaro/plugins/coco_format/converter.py
+++ b/datumaro/datumaro/plugins/coco_format/converter.py
@@ -67,6 +67,9 @@ def __init__(self, context):
     def is_empty(self):
         return len(self._data['annotations']) == 0
 
+    def _get_image_id(self, item):
+        return self._context._get_image_id(item)
+
     def save_image_info(self, item, filename):
         if item.has_image:
             h, w = item.image.size
@@ -75,7 +78,7 @@ def save_image_info(self, item, filename):
             w = 0
 
         self._data['images'].append({
-            'id': _cast(item.id, int, 0),
+            'id': self._get_image_id(item),
             'width': int(w),
             'height': int(h),
             'file_name': _cast(filename, str, ''),
@@ -136,7 +139,7 @@ def save_annotations(self, item):
 
             elem = {
                 'id': self._get_ann_id(ann),
-                'image_id': _cast(item.id, int, 0),
+                'image_id': self._get_image_id(item),
                 'category_id': 0, # NOTE: workaround for a bug in cocoapi
                 'caption': ann.caption,
             }
@@ -340,7 +343,7 @@ def convert_instance(self, instance, item):
 
         elem = {
             'id': self._get_ann_id(ann),
-            'image_id': _cast(item.id, int, 0),
+            'image_id': self._get_image_id(item),
             'category_id': _cast(ann.label, int, -1) + 1,
             'segmentation': segmentation,
             'area': float(area),
@@ -454,7 +457,7 @@ def save_annotations(self, item):
 
             elem = {
                 'id': self._get_ann_id(ann),
-                'image_id': _cast(item.id, int, 0),
+                'image_id': self._get_image_id(item),
                 'category_id': int(ann.label) + 1,
             }
             if 'score' in ann.attributes:
@@ -504,36 +507,50 @@ def __init__(self, extractor, save_dir,
 
         self._crop_covered = crop_covered
 
-    def make_dirs(self):
+        self._image_ids = {}
+
+    def _make_dirs(self):
         self._images_dir = osp.join(self._save_dir, CocoPath.IMAGES_DIR)
         os.makedirs(self._images_dir, exist_ok=True)
 
         self._ann_dir = osp.join(self._save_dir, CocoPath.ANNOTATIONS_DIR)
         os.makedirs(self._ann_dir, exist_ok=True)
 
-    def make_task_converter(self, task):
+    def _make_task_converter(self, task):
         if task not in self._TASK_CONVERTER:
             raise NotImplementedError()
         return self._TASK_CONVERTER[task](self)
 
-    def make_task_converters(self):
+    def _make_task_converters(self):
         return {
-            task: self.make_task_converter(task) for task in self._tasks
+            task: self._make_task_converter(task) for task in self._tasks
         }
 
-    def save_image(self, item, filename):
+    def _get_image_id(self, item):
+        image_id = self._image_ids.get(item.id)
+        if image_id is None:
+            image_id = _cast(item.id, int, len(self._image_ids) + 1)
+            self._image_ids[item.id] = image_id
+        return image_id
+
+    def _save_image(self, item):
         image = item.image.data
         if image is None:
             log.warning("Item '%s' has no image" % item.id)
-            return
+            return ''
 
+        filename = item.image.filename
+        if filename:
+            filename = osp.splitext(filename)[0]
+        else:
+            filename = item.id
+        filename += CocoPath.IMAGE_EXT
         path = osp.join(self._images_dir, filename)
         save_image(path, image)
-
         return path
 
     def convert(self):
-        self.make_dirs()
+        self._make_dirs()
 
         subsets = self._extractor.subsets()
         if len(subsets) == 0:
@@ -546,16 +563,16 @@ def convert(self):
                 subset_name = DEFAULT_SUBSET_NAME
                 subset = self._extractor
 
-            task_converters = self.make_task_converters()
+            task_converters = self._make_task_converters()
             for task_conv in task_converters.values():
                 task_conv.save_categories(subset)
             for item in subset:
                 filename = ''
                 if item.has_image:
-                    filename = str(item.id) + CocoPath.IMAGE_EXT
+                    filename = item.image.filename
                 if self._save_images:
                     if item.has_image:
-                        self.save_image(item, filename)
+                        filename = self._save_image(item)
                     else:
                         log.debug("Item '%s' has no image info" % item.id)
                 for task_conv in task_converters.values():
diff --git a/datumaro/datumaro/plugins/coco_format/extractor.py b/datumaro/datumaro/plugins/coco_format/extractor.py
index dca7d8532a5..1e0fde72d26 100644
--- a/datumaro/datumaro/plugins/coco_format/extractor.py
+++ b/datumaro/datumaro/plugins/coco_format/extractor.py
@@ -15,7 +15,7 @@
     AnnotationType, Label, RleMask, Points, Polygon, Bbox, Caption,
     LabelCategories, PointsCategories
 )
-from datumaro.util.image import lazy_image
+from datumaro.util.image import Image
 
 from .format import CocoTask, CocoPath
 
@@ -117,7 +117,15 @@ def _load_items(self, loader):
 
         for img_id in loader.getImgIds():
             image_info = loader.loadImgs(img_id)[0]
-            image = self._find_image(image_info['file_name'])
+            image_path = self._find_image(image_info['file_name'])
+            if not image_path:
+                image_path = image_info['file_name']
+            image_size = (image_info.get('height'), image_info.get('width'))
+            if all(image_size):
+                image_size = (int(image_size[0]), int(image_size[1]))
+            else:
+                image_size = None
+            image = Image(path=image_path, size=image_size)
 
             anns = loader.getAnnIds(imgIds=img_id)
             anns = loader.loadAnns(anns)
@@ -232,7 +240,8 @@ def _find_image(self, file_name):
         ]
         for image_path in search_paths:
             if osp.exists(image_path):
-                return lazy_image(image_path)
+                return image_path
+        return None
 
 class CocoImageInfoExtractor(_CocoExtractor):
     def __init__(self, path, **kwargs):
diff --git a/datumaro/tests/test_coco_format.py b/datumaro/tests/test_coco_format.py
index 9283f475910..2caa03a7c09 100644
--- a/datumaro/tests/test_coco_format.py
+++ b/datumaro/tests/test_coco_format.py
@@ -640,7 +640,7 @@ def test_can_save_dataset_with_image_info(self):
         class TestExtractor(Extractor):
             def __iter__(self):
                 return iter([
-                    DatasetItem(id=1, image=Image(size=(10, 15))),
+                    DatasetItem(id=1, image=Image(path='1.jpg', size=(10, 15))),
                 ])
 
         with TestDir() as test_dir:

From 7d47ac998e2d7ca2176742456ca21a629a6bd7e2 Mon Sep 17 00:00:00 2001
From: Zhiltsov Max <zhiltsov.max35@gmail.com>
Date: Mon, 17 Feb 2020 18:31:10 +0300
Subject: [PATCH 14/25] Update CVAT format with image info

---
 .../datumaro/plugins/cvat_format/converter.py | 47 ++++++++++---------
 .../datumaro/plugins/cvat_format/extractor.py | 21 ++++-----
 datumaro/tests/test_cvat_format.py            | 14 +++---
 3 files changed, 43 insertions(+), 39 deletions(-)

diff --git a/datumaro/datumaro/plugins/cvat_format/converter.py b/datumaro/datumaro/plugins/cvat_format/converter.py
index edd72a1cab8..d18699d1d09 100644
--- a/datumaro/datumaro/plugins/cvat_format/converter.py
+++ b/datumaro/datumaro/plugins/cvat_format/converter.py
@@ -157,13 +157,8 @@ def write(self):
         self._writer.open_root()
         self._write_meta()
 
-        for item in self._extractor:
-            if self._context._save_images:
-                if item.has_image:
-                    self._save_image(item)
-                else:
-                    log.debug("Item '%s' has no image info" % item.id)
-            self._write_item(item)
+        for index, item in enumerate(self._extractor):
+            self._write_item(item, index)
 
         self._writer.close_root()
 
@@ -171,27 +166,35 @@ def _save_image(self, item):
         image = item.image.data
         if image is None:
             log.warning("Item '%s' has no image" % item.id)
-            return
+            return ''
 
-        file_name = item.name
-        if not file_name:
-            file_name = str(item.id)
-        file_name += CvatPath.IMAGE_EXT
-        image_path = osp.join(self._context._images_dir, file_name)
+        filename = item.image.filename
+        if filename:
+            filename = osp.splitext(filename)[0]
+        else:
+            filename = item.id
+        filename += CvatPath.IMAGE_EXT
+        image_path = osp.join(self._context._images_dir, filename)
         save_image(image_path, image)
+        return filename
 
-    def _write_item(self, item):
-        image_name = item.name
-        if not image_name:
-            image_name = str(item.id)
+    def _write_item(self, item, index):
         image_info = OrderedDict([
-            ("id", str(item.id)),
-            ("name", image_name),
+            ("id", str(_cast(item.id, int, index))),
         ])
         if item.has_image:
-            h, w = item.image.size
-            image_info["width"] = str(w)
-            image_info["height"] = str(h)
+            size = item.image.size
+            if size:
+                h, w = size
+                image_info["width"] = str(w)
+                image_info["height"] = str(h)
+
+            filename = item.image.filename
+            if self._context._save_images:
+                filename = self._save_image(item)
+            image_info["name"] = filename
+        else:
+            log.debug("Item '%s' has no image info" % item.id)
         self._writer.open_image(image_info)
 
         for ann in item.annotations:
diff --git a/datumaro/datumaro/plugins/cvat_format/extractor.py b/datumaro/datumaro/plugins/cvat_format/extractor.py
index 04e0d2f7c5c..1f36f70cc77 100644
--- a/datumaro/datumaro/plugins/cvat_format/extractor.py
+++ b/datumaro/datumaro/plugins/cvat_format/extractor.py
@@ -295,23 +295,22 @@ def _parse_ann(cls, ann, categories):
             raise NotImplementedError("Unknown annotation type '%s'" % ann_type)
 
     def _load_items(self, parsed):
-        for item_id, item_desc in parsed.items():
-            image = None
-            file_name = item_desc.get('name')
-            if not file_name:
-                file_name = item_id
-            image_path = self._find_image(file_name)
+        for frame_id, item_desc in parsed.items():
+            filename = item_desc.get('name')
+            if filename:
+                filename = self._find_image(filename)
+            if not filename:
+                filename = item_desc.get('name')
             image_size = (item_desc.get('height'), item_desc.get('width'))
             if all(image_size):
                 image_size = (int(image_size[0]), int(image_size[1]))
             else:
                 image_size = None
-            if image_path:
-                image = Image(path=image_path, size=image_size)
-            elif image_size:
-                image = Image(size=image_size)
+            image = None
+            if filename:
+                image = Image(path=filename, size=image_size)
 
-            parsed[item_id] = DatasetItem(id=item_id, subset=self._subset,
+            parsed[frame_id] = DatasetItem(id=frame_id, subset=self._subset,
                 image=image, annotations=item_desc.get('annotations'))
         return parsed
 
diff --git a/datumaro/tests/test_cvat_format.py b/datumaro/tests/test_cvat_format.py
index 43d9520508f..17061217845 100644
--- a/datumaro/tests/test_cvat_format.py
+++ b/datumaro/tests/test_cvat_format.py
@@ -159,7 +159,7 @@ def test_can_save_and_load(self):
         label_categories.items[2].attributes.update(['a1', 'a2'])
         label_categories.attributes.update(['z_order', 'occluded'])
 
-        class SrcTestExtractor(Extractor):
+        class SrcExtractor(Extractor):
             def __iter__(self):
                 return iter([
                     DatasetItem(id=0, subset='s1', image=np.zeros((5, 10, 3)),
@@ -193,13 +193,14 @@ def __iter__(self):
                         ]
                     ),
 
-                    DatasetItem(id=3, subset='s3', image=Image(size=(2, 4))),
+                    DatasetItem(id=3, subset='s3', image=Image(
+                        path='3.jpg', size=(2, 4))),
                 ])
 
             def categories(self):
                 return { AnnotationType.label: label_categories }
 
-        class DstTestExtractor(Extractor):
+        class DstExtractor(Extractor):
             def __iter__(self):
                 return iter([
                     DatasetItem(id=0, subset='s1', image=np.zeros((5, 10, 3)),
@@ -235,13 +236,14 @@ def __iter__(self):
                         ]
                     ),
 
-                    DatasetItem(id=3, subset='s3', image=Image(size=(2, 4))),
+                    DatasetItem(id=3, subset='s3', image=Image(
+                        path='3.jpg', size=(2, 4))),
                 ])
 
             def categories(self):
                 return { AnnotationType.label: label_categories }
 
         with TestDir() as test_dir:
-            self._test_save_and_load(SrcTestExtractor(),
+            self._test_save_and_load(SrcExtractor(),
                 CvatConverter(save_images=True), test_dir,
-                target_dataset=DstTestExtractor())
+                target_dataset=DstExtractor())

From c11ee19f80f1fdfa0596399837eac7c03db9f59e Mon Sep 17 00:00:00 2001
From: Zhiltsov Max <zhiltsov.max35@gmail.com>
Date: Mon, 17 Feb 2020 18:33:08 +0300
Subject: [PATCH 15/25] Update TFrecord format with image info

---
 .../tf_detection_api_format/extractor.py      | 20 +++++++++----------
 .../datumaro/plugins/voc_format/converter.py  | 14 +++++++++----
 datumaro/tests/test_tfrecord_format.py        | 16 +++++++--------
 3 files changed, 27 insertions(+), 23 deletions(-)

diff --git a/datumaro/datumaro/plugins/tf_detection_api_format/extractor.py b/datumaro/datumaro/plugins/tf_detection_api_format/extractor.py
index 7e571321d16..8974c65d805 100644
--- a/datumaro/datumaro/plugins/tf_detection_api_format/extractor.py
+++ b/datumaro/datumaro/plugins/tf_detection_api_format/extractor.py
@@ -163,31 +163,29 @@ def _parse_tfrecord_file(cls, filepath, subset_name, images_dir):
                 item_id = osp.splitext(frame_filename)[0]
 
             annotations = []
-            for index, shape in enumerate(
-                    np.dstack((labels, xmins, ymins, xmaxs, ymaxs))[0]):
+            for shape in np.dstack((labels, xmins, ymins, xmaxs, ymaxs))[0]:
                 label = shape[0].decode('utf-8')
                 x = clamp(shape[1] * frame_width, 0, frame_width)
                 y = clamp(shape[2] * frame_height, 0, frame_height)
                 w = clamp(shape[3] * frame_width, 0, frame_width) - x
                 h = clamp(shape[4] * frame_height, 0, frame_height) - y
                 annotations.append(Bbox(x, y, w, h,
-                    label=dataset_labels.get(label, None), id=index
+                    label=dataset_labels.get(label)
                 ))
 
-            image_params = {}
+            image_size = None
             if frame_height and frame_width:
-                image_params['size'] = (frame_height, frame_width)
+                image_size = (frame_height, frame_width)
 
+            image_params = {}
             if frame_image and frame_format:
                 image_params['data'] = lazy_image(frame_image, decode_image)
-            elif frame_filename and images_dir:
-                image_path = osp.join(images_dir, frame_filename)
-                if osp.exists(image_path):
-                    image_params['path'] = image_path
+            if frame_filename and images_dir:
+                image_params['path'] = osp.join(images_dir, frame_filename)
 
-            image=None
+            image = None
             if image_params:
-                image = Image(**image_params)
+                image = Image(**image_params, size=image_size)
 
             dataset_items.append(DatasetItem(id=item_id, subset=subset_name,
                 image=image, annotations=annotations))
diff --git a/datumaro/datumaro/plugins/voc_format/converter.py b/datumaro/datumaro/plugins/voc_format/converter.py
index c7ecbdee039..37b181f7c48 100644
--- a/datumaro/datumaro/plugins/voc_format/converter.py
+++ b/datumaro/datumaro/plugins/voc_format/converter.py
@@ -135,10 +135,17 @@ def save_subsets(self):
             for item in subset:
                 log.debug("Converting item '%s'", item.id)
 
+                image_filename = ''
+                if item.has_image:
+                    image_filename = item.image.filename
                 if self._save_images:
                     if item.has_image and item.image.has_data:
-                        save_image(osp.join(self._images_dir,
-                                item.id + VocPath.IMAGE_EXT),
+                        if image_filename:
+                            image_filename = osp.splitext(image_filename)[0]
+                        else:
+                            image_filename = item.id
+                        image_filename += VocPath.IMAGE_EXT
+                        save_image(osp.join(self._images_dir, image_filename),
                             item.image.data)
                     else:
                         log.debug("Item '%s' has no image" % item.id)
@@ -161,8 +168,7 @@ def save_subsets(self):
                     else:
                         folder = ''
                     ET.SubElement(root_elem, 'folder').text = folder
-                    ET.SubElement(root_elem, 'filename').text = \
-                        item.id + VocPath.IMAGE_EXT
+                    ET.SubElement(root_elem, 'filename').text = image_filename
 
                     source_elem = ET.SubElement(root_elem, 'source')
                     ET.SubElement(source_elem, 'database').text = 'Unknown'
diff --git a/datumaro/tests/test_tfrecord_format.py b/datumaro/tests/test_tfrecord_format.py
index 31d93f34f6c..efbef0fd2b8 100644
--- a/datumaro/tests/test_tfrecord_format.py
+++ b/datumaro/tests/test_tfrecord_format.py
@@ -34,16 +34,16 @@ def __iter__(self):
                     DatasetItem(id=1, subset='train',
                         image=np.ones((16, 16, 3)),
                         annotations=[
-                            Bbox(0, 4, 4, 8, label=2, id=0),
-                            Bbox(0, 4, 4, 4, label=3, id=1),
-                            Bbox(2, 4, 4, 4, id=2),
+                            Bbox(0, 4, 4, 8, label=2),
+                            Bbox(0, 4, 4, 4, label=3),
+                            Bbox(2, 4, 4, 4),
                         ]
                     ),
 
                     DatasetItem(id=2, subset='val',
                         image=np.ones((8, 8, 3)),
                         annotations=[
-                            Bbox(1, 2, 4, 2, label=3, id=0),
+                            Bbox(1, 2, 4, 2, label=3),
                         ]
                     ),
 
@@ -72,15 +72,15 @@ def __iter__(self):
                     DatasetItem(id=1,
                         image=np.ones((16, 16, 3)),
                         annotations=[
-                            Bbox(2, 1, 4, 4, label=2, id=0),
-                            Bbox(4, 2, 8, 4, label=3, id=1),
+                            Bbox(2, 1, 4, 4, label=2),
+                            Bbox(4, 2, 8, 4, label=3),
                         ]
                     ),
 
                     DatasetItem(id=2,
                         image=np.ones((8, 8, 3)) * 2,
                         annotations=[
-                            Bbox(4, 4, 4, 4, label=3, id=0),
+                            Bbox(4, 4, 4, 4, label=3),
                         ]
                     ),
 
@@ -106,7 +106,7 @@ def test_can_save_dataset_with_image_info(self):
         class TestExtractor(Extractor):
             def __iter__(self):
                 return iter([
-                    DatasetItem(id=1, image=Image(size=(10, 15))),
+                    DatasetItem(id=1, image=Image(path='1/q.e', size=(10, 15))),
                 ])
 
             def categories(self):

From 0a93b80c10c703dfb6f357718fbbe9e1d8d37eb3 Mon Sep 17 00:00:00 2001
From: Zhiltsov Max <zhiltsov.max35@gmail.com>
Date: Mon, 17 Feb 2020 18:34:27 +0300
Subject: [PATCH 16/25] Update VOC formar with image info

---
 datumaro/datumaro/plugins/voc_format/extractor.py | 7 +++----
 datumaro/tests/test_voc_format.py                 | 2 +-
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/datumaro/datumaro/plugins/voc_format/extractor.py b/datumaro/datumaro/plugins/voc_format/extractor.py
index cb8b5e4eb4f..66d83f0a5f6 100644
--- a/datumaro/datumaro/plugins/voc_format/extractor.py
+++ b/datumaro/datumaro/plugins/voc_format/extractor.py
@@ -86,6 +86,8 @@ def _load_det_annotations(self):
                 ann_file_data = f.read()
                 ann_file_root = ET.fromstring(ann_file_data)
                 item = ann_file_root.find('filename').text
+                if not item:
+                    item = ann_item
                 item = osp.splitext(item)[0]
                 det_annotations[item] = ann_file_data
 
@@ -130,11 +132,8 @@ def __iter__(self):
                 yield item
 
     def _get(self, item_id, subset_name):
-        image = None
-        image_path = osp.join(self._path, VocPath.IMAGES_DIR,
+        image = osp.join(self._path, VocPath.IMAGES_DIR,
             item_id + VocPath.IMAGE_EXT)
-        if osp.isfile(image_path):
-            image = lazy_image(image_path)
 
         annotations = self._get_annotations(item_id)
 
diff --git a/datumaro/tests/test_voc_format.py b/datumaro/tests/test_voc_format.py
index 30473e21f14..841e1b89853 100644
--- a/datumaro/tests/test_voc_format.py
+++ b/datumaro/tests/test_voc_format.py
@@ -696,7 +696,7 @@ def test_can_save_dataset_with_image_info(self):
         class TestExtractor(TestExtractorBase):
             def __iter__(self):
                 return iter([
-                    DatasetItem(id=1, image=Image(size=(10, 15))),
+                    DatasetItem(id=1, image=Image(path='1.jpg', size=(10, 15))),
                 ])
 
         with TestDir() as test_dir:

From f9e5c8ca16773700462c3826fc43ca71d93ec716 Mon Sep 17 00:00:00 2001
From: Zhiltsov Max <zhiltsov.max35@gmail.com>
Date: Mon, 17 Feb 2020 18:35:50 +0300
Subject: [PATCH 17/25] Update YOLO format with image info

---
 .../datumaro/plugins/yolo_format/converter.py | 21 +++++++-----
 .../datumaro/plugins/yolo_format/extractor.py | 33 ++++++++++++++-----
 .../datumaro/plugins/yolo_format/format.py    |  4 ++-
 .../datumaro/plugins/yolo_format/importer.py  | 10 +++---
 datumaro/tests/test_yolo_format.py            | 32 +++++++++++++++++-
 5 files changed, 77 insertions(+), 23 deletions(-)

diff --git a/datumaro/datumaro/plugins/yolo_format/converter.py b/datumaro/datumaro/plugins/yolo_format/converter.py
index 59a46a6260c..de30d7d785e 100644
--- a/datumaro/datumaro/plugins/yolo_format/converter.py
+++ b/datumaro/datumaro/plugins/yolo_format/converter.py
@@ -75,21 +75,24 @@ def __call__(self, extractor, save_dir):
             image_paths = OrderedDict()
 
             for item in subset:
-                image_name = '%s.jpg' % item.id
-                image_paths[item.id] = osp.join('data',
-                    osp.basename(subset_dir), image_name)
+                if not item.has_image:
+                    raise Exception("Failed to export item '%s': "
+                        "item has no image info" % item.id)
+                height, width = item.image.size
 
+                image_name = item.image.filename
+                item_name = osp.splitext(item.image.filename)[0]
                 if self._save_images:
                     if item.has_image and item.image.has_data:
+                        if not item_name:
+                            item_name = item.id
+                        image_name = item_name + '.jpg'
                         save_image(osp.join(subset_dir, image_name),
                             item.image.data)
                     else:
                         log.warning("Item '%s' has no image" % item.id)
-
-                if not item.has_image:
-                    raise Exception("Failed to export item '%s': "
-                        "item has no image info" % item.id)
-                height, width = item.image.size
+                image_paths[item.id] = osp.join('data',
+                    osp.basename(subset_dir), image_name)
 
                 yolo_annotation = ''
                 for bbox in item.annotations:
@@ -102,7 +105,7 @@ def __call__(self, extractor, save_dir):
                     yolo_bb = ' '.join('%.6f' % p for p in yolo_bb)
                     yolo_annotation += '%s %s\n' % (bbox.label, yolo_bb)
 
-                annotation_path = osp.join(subset_dir, '%s.txt' % item.id)
+                annotation_path = osp.join(subset_dir, '%s.txt' % item_name)
                 with open(annotation_path, 'w') as f:
                     f.write(yolo_annotation)
 
diff --git a/datumaro/datumaro/plugins/yolo_format/extractor.py b/datumaro/datumaro/plugins/yolo_format/extractor.py
index cb52327717c..7840b26c5ca 100644
--- a/datumaro/datumaro/plugins/yolo_format/extractor.py
+++ b/datumaro/datumaro/plugins/yolo_format/extractor.py
@@ -10,7 +10,7 @@
 from datumaro.components.extractor import (SourceExtractor, Extractor,
     DatasetItem, AnnotationType, Bbox, LabelCategories
 )
-from datumaro.util.image import lazy_image
+from datumaro.util.image import Image
 
 from .format import YoloPath
 
@@ -33,16 +33,31 @@ def __len__(self):
         def categories(self):
             return self._parent.categories()
 
-    def __init__(self, config_path):
+    def __init__(self, config_path, image_info=None):
         super().__init__()
 
         if not osp.isfile(config_path):
-            raise Exception("Can't read dataset descriptor file '%s'" % \
+            raise Exception("Can't read dataset descriptor file '%s'" %
                 config_path)
 
         rootpath = osp.dirname(config_path)
         self._path = rootpath
 
+        assert image_info is None or isinstance(image_info, (str, dict))
+        if image_info is None:
+            image_info = osp.join(rootpath, YoloPath.IMAGE_META_FILE)
+            if not osp.isfile(image_info):
+                image_info = {}
+        if isinstance(image_info, str):
+            if not osp.isfile(image_info):
+                raise Exception("Can't read image meta file '%s'" % image_info)
+            with open(image_info) as f:
+                image_info = {}
+                for line in f:
+                    image_name, h, w = line.strip().split()
+                    image_info[image_name] = (int(h), int(w))
+        self._image_info = image_info
+
         with open(config_path, 'r') as f:
             config_lines = f.readlines()
 
@@ -77,10 +92,10 @@ def __init__(self, config_path):
                 subset.items = OrderedDict(
                     (osp.splitext(osp.basename(p))[0], p.strip()) for p in f)
 
-            for image_path in subset.items.values():
+            for item_id, image_path in subset.items.items():
                 image_path = self._make_local_path(image_path)
-                if not osp.isfile(image_path):
-                    raise Exception("Can't find image '%s'" % image_path)
+                if not osp.isfile(image_path) and item_id not in image_info:
+                    raise Exception("Can't find image '%s'" % item_id)
 
             subsets[subset_name] = subset
 
@@ -103,8 +118,10 @@ def _get(self, item_id, subset_name):
 
         if isinstance(item, str):
             image_path = self._make_local_path(item)
-            image = lazy_image(image_path)
-            h, w, _ = image().shape
+            image_size = self._image_info.get(item_id)
+            image = Image(path=image_path, size=image_size)
+            h, w = image.size
+
             anno_path = osp.splitext(image_path)[0] + '.txt'
             annotations = self._parse_annotations(anno_path, w, h)
 
diff --git a/datumaro/datumaro/plugins/yolo_format/format.py b/datumaro/datumaro/plugins/yolo_format/format.py
index 8d44a9ba8fb..c88c99d442d 100644
--- a/datumaro/datumaro/plugins/yolo_format/format.py
+++ b/datumaro/datumaro/plugins/yolo_format/format.py
@@ -6,4 +6,6 @@
 
 class YoloPath:
     DEFAULT_SUBSET_NAME = 'train'
-    SUBSET_NAMES = ['train', 'valid']
\ No newline at end of file
+    SUBSET_NAMES = ['train', 'valid']
+
+    IMAGE_META_FILE = 'images.meta'
\ No newline at end of file
diff --git a/datumaro/datumaro/plugins/yolo_format/importer.py b/datumaro/datumaro/plugins/yolo_format/importer.py
index 26532d09128..fcee669dc6b 100644
--- a/datumaro/datumaro/plugins/yolo_format/importer.py
+++ b/datumaro/datumaro/plugins/yolo_format/importer.py
@@ -15,18 +15,20 @@ def __call__(self, path, **extra_params):
         from datumaro.components.project import Project # cyclic import
         project = Project()
 
-        if not osp.exists(path):
-            raise Exception("Failed to find 'yolo' dataset at '%s'" % path)
-
         if path.endswith('.data') and osp.isfile(path):
             config_paths = [path]
         else:
             config_paths = glob(osp.join(path, '*.data'))
 
+        if not osp.exists(path) or not config_paths:
+            raise Exception("Failed to find 'yolo' dataset at '%s'" % path)
+
         for config_path in config_paths:
             log.info("Found a dataset at '%s'" % config_path)
 
-            source_name = osp.splitext(osp.basename(config_path))[0]
+            source_name = '%s_%s' % (
+                osp.basename(osp.dirname(config_path)),
+                osp.splitext(osp.basename(config_path))[0])
             project.add_source(source_name, {
                 'url': config_path,
                 'format': 'yolo',
diff --git a/datumaro/tests/test_yolo_format.py b/datumaro/tests/test_yolo_format.py
index 3af08781c36..e9a95108a9e 100644
--- a/datumaro/tests/test_yolo_format.py
+++ b/datumaro/tests/test_yolo_format.py
@@ -59,7 +59,7 @@ class TestExtractor(Extractor):
             def __iter__(self):
                 return iter([
                     DatasetItem(id=1, subset='train',
-                        image=Image(size=(10, 15)),
+                        image=Image(path='1.jpg', size=(10, 15)),
                         annotations=[
                             Bbox(0, 2, 4, 2, label=2),
                             Bbox(3, 3, 2, 3, label=4),
@@ -84,3 +84,33 @@ def categories(self):
             parsed_dataset = YoloImporter()(test_dir).make_dataset()
 
             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([
+                    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,
+                }
+
+        with TestDir() as test_dir:
+            source_dataset = TestExtractor()
+
+            YoloConverter()(source_dataset, test_dir)
+
+            parsed_dataset = YoloImporter()(test_dir,
+                image_info={'1': (10, 15)}).make_dataset()
+
+            compare_datasets(self, source_dataset, parsed_dataset)

From bef057bbfb7f7631c19929ed4116618c2d5cc006 Mon Sep 17 00:00:00 2001
From: Zhiltsov Max <zhiltsov.max35@gmail.com>
Date: Mon, 17 Feb 2020 18:37:02 +0300
Subject: [PATCH 18/25] Update dataset manager bindings with image info

---
 cvat/apps/dataset_manager/bindings.py | 167 +++++++++++++++-----------
 1 file changed, 100 insertions(+), 67 deletions(-)

diff --git a/cvat/apps/dataset_manager/bindings.py b/cvat/apps/dataset_manager/bindings.py
index 09f0a698e99..0c984e330cf 100644
--- a/cvat/apps/dataset_manager/bindings.py
+++ b/cvat/apps/dataset_manager/bindings.py
@@ -14,7 +14,7 @@
 from cvat.apps.engine.models import Task, ShapeType, AttributeType
 
 import datumaro.components.extractor as datumaro
-from datumaro.util.image import lazy_image
+from datumaro.util.image import Image
 
 
 class CvatImagesDirExtractor(datumaro.Extractor):
@@ -29,8 +29,7 @@ def __init__(self, url):
                 path = osp.join(dirpath, name)
                 if self._is_image(path):
                     item_id = Task.get_image_frame(path)
-                    item = datumaro.DatasetItem(
-                        id=item_id, image=lazy_image(path))
+                    item = datumaro.DatasetItem(id=item_id, image=path)
                     items.append((item.id, item))
 
         items = sorted(items, key=lambda e: int(e[0]))
@@ -49,11 +48,6 @@ def __len__(self):
     def subsets(self):
         return self._subsets
 
-    def get(self, item_id, subset=None, path=None):
-        if path or subset:
-            raise KeyError()
-        return self._items[item_id]
-
     def _is_image(self, path):
         for ext in self._SUPPORTED_FORMATS:
             if osp.isfile(path) and path.endswith(ext):
@@ -61,29 +55,24 @@ def _is_image(self, path):
         return False
 
 
-class CvatTaskExtractor(datumaro.Extractor):
-    def __init__(self, url, db_task, user):
-        self._db_task = db_task
-        self._categories = self._load_categories()
-
-        cvat_annotations = TaskAnnotation(db_task.id, user)
-        with transaction.atomic():
-            cvat_annotations.init_from_db()
-        cvat_annotations = Annotation(cvat_annotations.ir_data, db_task)
+class CvatAnnotationsExtractor(datumaro.Extractor):
+    def __init__(self, url, cvat_annotations):
+        self._categories = self._load_categories(cvat_annotations)
 
         dm_annotations = []
 
-        for cvat_anno in cvat_annotations.group_by_frame():
-            dm_anno = self._read_cvat_anno(cvat_anno)
-            dm_item = datumaro.DatasetItem(
-                id=cvat_anno.frame, annotations=dm_anno)
+        for cvat_frame_anno in cvat_annotations.group_by_frame():
+            dm_anno = self._read_cvat_anno(cvat_frame_anno, cvat_annotations)
+            dm_image = Image(path=cvat_frame_anno.name, size=(
+                cvat_frame_anno.height, cvat_frame_anno.width)
+            )
+            dm_item = datumaro.DatasetItem(id=cvat_frame_anno.frame,
+                annotations=dm_anno, image=dm_image)
             dm_annotations.append((dm_item.id, dm_item))
 
         dm_annotations = sorted(dm_annotations, key=lambda e: int(e[0]))
         self._items = OrderedDict(dm_annotations)
 
-        self._subsets = None
-
     def __iter__(self):
         for item in self._items.values():
             yield item
@@ -91,70 +80,58 @@ def __iter__(self):
     def __len__(self):
         return len(self._items)
 
+    # pylint: disable=no-self-use
     def subsets(self):
-        return self._subsets
+        return []
+    # pylint: enable=no-self-use
 
-    def get(self, item_id, subset=None, path=None):
-        if path or subset:
-            raise KeyError()
-        return self._items[item_id]
+    def categories(self):
+        return self._categories
 
-    def _load_categories(self):
+    @staticmethod
+    def _load_categories(cvat_anno):
         categories = {}
         label_categories = datumaro.LabelCategories()
 
-        db_labels = self._db_task.label_set.all()
-        for db_label in db_labels:
-            db_attributes = db_label.attributespec_set.all()
-            label_categories.add(db_label.name)
-
-            for db_attr in db_attributes:
-                label_categories.attributes.add(db_attr.name)
+        for _, label in cvat_anno.meta['task']['labels']:
+            label_categories.add(label['name'])
+            for _, attr in label['attributes']:
+                label_categories.attributes.add(attr['name'])
 
         categories[datumaro.AnnotationType.label] = label_categories
 
         return categories
 
-    def categories(self):
-        return self._categories
-
-    def _read_cvat_anno(self, cvat_anno):
+    def _read_cvat_anno(self, cvat_frame_anno, cvat_task_anno):
         item_anno = []
 
         categories = self.categories()
         label_cat = categories[datumaro.AnnotationType.label]
-
-        label_map = {}
-        label_attrs = {}
-        db_labels = self._db_task.label_set.all()
-        for db_label in db_labels:
-            label_map[db_label.name] = label_cat.find(db_label.name)[0]
-
-            attrs = {}
-            db_attributes = db_label.attributespec_set.all()
-            for db_attr in db_attributes:
-                attrs[db_attr.name] = db_attr
-            label_attrs[db_label.name] = attrs
-        map_label = lambda label_db_name: label_map[label_db_name]
+        map_label = lambda name: label_cat.find(name)[0]
+        label_attrs = {
+            label['name']: label['attributes']
+            for _, label in cvat_task_anno.meta['task']['labels']
+        }
 
         def convert_attrs(label, cvat_attrs):
             cvat_attrs = {a.name: a.value for a in cvat_attrs}
             dm_attr = dict()
-            for attr_name, attr_spec in label_attrs[label].items():
-                attr_value = cvat_attrs.get(attr_name, attr_spec.default_value)
+            for _, a_desc in label_attrs[label]:
+                a_name = a_desc['name']
+                a_value = cvat_attrs.get(a_name, a_desc['default_value'])
                 try:
-                    if attr_spec.input_type == AttributeType.NUMBER:
-                        attr_value = float(attr_value)
-                    elif attr_spec.input_type == AttributeType.CHECKBOX:
-                        attr_value = attr_value.lower() == 'true'
-                    dm_attr[attr_name] = attr_value
+                    if a_desc['input_type'] == AttributeType.NUMBER:
+                        a_value = float(a_value)
+                    elif a_desc['input_type'] == AttributeType.CHECKBOX:
+                        a_value = (a_value.lower() == 'true')
+                    dm_attr[a_name] = a_value
                 except Exception as e:
-                    slogger.task[self._db_task.id].error(
-                        "Failed to convert attribute '%s'='%s': %s" % \
-                            (attr_name, attr_value, e))
+                    raise Exception(
+                        "Failed to convert attribute '%s'='%s': %s" %
+                        (a_name, a_value, e))
             return dm_attr
 
-        for tag_obj in cvat_anno.tags:
+        for tag_obj in cvat_frame_anno.tags:
             anno_group = tag_obj.group
             anno_label = map_label(tag_obj.label)
             anno_attr = convert_attrs(tag_obj.label, tag_obj.attributes)
@@ -163,7 +140,7 @@ def convert_attrs(label, cvat_attrs):
                 attributes=anno_attr, group=anno_group)
             item_anno.append(anno)
 
-        for shape_obj in cvat_anno.labeled_shapes:
+        for shape_obj in cvat_frame_anno.labeled_shapes:
             anno_group = shape_obj.group
             anno_label = map_label(shape_obj.label)
             anno_attr = convert_attrs(shape_obj.label, shape_obj.attributes)
@@ -183,8 +160,64 @@ def convert_attrs(label, cvat_attrs):
                 anno = datumaro.Bbox(x0, y0, x1 - x0, y1 - y0,
                     label=anno_label, attributes=anno_attr, group=anno_group)
             else:
-                raise Exception("Unknown shape type '%s'" % (shape_obj.type))
+                raise Exception("Unknown shape type '%s'" % shape_obj.type)
 
             item_anno.append(anno)
 
-        return item_anno
\ No newline at end of file
+        return item_anno
+
+
+class CvatTaskExtractor(CvatAnnotationsExtractor):
+    def __init__(self, url, db_task, user):
+        cvat_annotations = TaskAnnotation(db_task.id, user)
+        with transaction.atomic():
+            cvat_annotations.init_from_db()
+        cvat_annotations = Annotation(cvat_annotations.ir_data, db_task)
+        super().__init__(url, cvat_annotations)
+
+
+def match_frame(item, cvat_task_anno):
+    frame_number = None
+    if frame_number is None:
+        try:
+            frame_number = cvat_task_anno.match_frame(item.id)
+        except Exception:
+            pass
+    if frame_number is None and item.has_image:
+        try:
+            frame_number = cvat_task_anno.match_frame(item.image.filename)
+        except Exception:
+            pass
+    if frame_number is None:
+        try:
+            frame_number = int(item.id)
+        except Exception:
+            pass
+    if not frame_number in cvat_task_anno.frame_info:
+        raise Exception("Could not match item id: '%s' with any task frame" %
+            item.id)
+    return frame_number
+
+def import_dm_annotations(dm_dataset, cvat_task_anno):
+    shapes = {
+        datumaro.AnnotationType.bbox: ShapeType.RECTANGLE,
+        datumaro.AnnotationType.polygon: ShapeType.POLYGON,
+        datumaro.AnnotationType.polyline: ShapeType.POLYLINE,
+        datumaro.AnnotationType.points: ShapeType.POINTS,
+    }
+
+    label_cat = dm_dataset.categories()[datumaro.AnnotationType.label]
+
+    for item in dm_dataset:
+        frame_number = match_frame(item, cvat_task_anno)
+
+        for ann in item.annotations:
+            if ann.type in shapes:
+                cvat_task_anno.add_shape(cvat_task_anno.LabeledShape(
+                    type=shapes[ann.type],
+                    frame=frame_number,
+                    label=label_cat.items[ann.label].name,
+                    points=ann.points,
+                    occluded=False,
+                    attributes=[],
+                ))
\ No newline at end of file

From e00e5a976a6a4b6aa1e8d7ecd5da7872d8b90a36 Mon Sep 17 00:00:00 2001
From: Zhiltsov Max <zhiltsov.max35@gmail.com>
Date: Mon, 17 Feb 2020 18:38:20 +0300
Subject: [PATCH 19/25] Add image name to id transform

---
 datumaro/datumaro/plugins/transforms.py | 11 +++++++++--
 datumaro/tests/test_transforms.py       | 17 +++++++++++++++++
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/datumaro/datumaro/plugins/transforms.py b/datumaro/datumaro/plugins/transforms.py
index 44ec703b478..81c5ff50187 100644
--- a/datumaro/datumaro/plugins/transforms.py
+++ b/datumaro/datumaro/plugins/transforms.py
@@ -5,6 +5,7 @@
 
 from itertools import groupby
 import logging as log
+import os.path as osp
 
 import pycocotools.mask as mask_utils
 
@@ -273,7 +274,6 @@ def __iter__(self):
         for i, item in enumerate(self._extractor):
             yield self.wrap_item(item, id=i + self._start)
 
-
 class MapSubsets(Transform, CliPlugin):
     @staticmethod
     def _mapping_arg(s):
@@ -302,4 +302,11 @@ def __init__(self, extractor, mapping=None):
 
     def transform_item(self, item):
         return self.wrap_item(item,
-            subset=self._mapping.get(item.subset, item.subset))
\ No newline at end of file
+            subset=self._mapping.get(item.subset, item.subset))
+
+class IdFromImageName(Transform, CliPlugin):
+    def transform_item(self, item):
+        name = item.id
+        if item.has_image and item.image.filename:
+            name = osp.splitext(item.image.filename)[0]
+        return self.wrap_item(item, id=name)
diff --git a/datumaro/tests/test_transforms.py b/datumaro/tests/test_transforms.py
index 41daf173493..11e997b19d7 100644
--- a/datumaro/tests/test_transforms.py
+++ b/datumaro/tests/test_transforms.py
@@ -244,3 +244,20 @@ def __iter__(self):
         actual = transforms.ShapesToBoxes(SrcExtractor())
         compare_datasets(self, DstExtractor(), actual)
 
+    def test_id_from_image(self):
+        class SrcExtractor(Extractor):
+            def __iter__(self):
+                return iter([
+                    DatasetItem(id=1, image='path.jpg'),
+                    DatasetItem(id=2),
+                ])
+
+        class DstExtractor(Extractor):
+            def __iter__(self):
+                return iter([
+                    DatasetItem(id='path', image='path.jpg'),
+                    DatasetItem(id=2),
+                ])
+
+        actual = transforms.IdFromImageName(SrcExtractor())
+        compare_datasets(self, DstExtractor(), actual)

From addd22e354c605fcf6a23a5ced3fd3da8950fb21 Mon Sep 17 00:00:00 2001
From: Zhiltsov Max <zhiltsov.max35@gmail.com>
Date: Tue, 18 Feb 2020 17:29:19 +0300
Subject: [PATCH 20/25] Fix coco export

---
 datumaro/datumaro/plugins/coco_format/converter.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/datumaro/datumaro/plugins/coco_format/converter.py b/datumaro/datumaro/plugins/coco_format/converter.py
index 4af56b75ff6..7b91aef4f5d 100644
--- a/datumaro/datumaro/plugins/coco_format/converter.py
+++ b/datumaro/datumaro/plugins/coco_format/converter.py
@@ -547,7 +547,7 @@ def _save_image(self, item):
         filename += CocoPath.IMAGE_EXT
         path = osp.join(self._images_dir, filename)
         save_image(path, image)
-        return path
+        return filename
 
     def convert(self):
         self._make_dirs()

From 81e52d37a2bf79672761676e45c6a6de8c513efd Mon Sep 17 00:00:00 2001
From: Zhiltsov Max <zhiltsov.max35@gmail.com>
Date: Tue, 18 Feb 2020 17:24:50 +0300
Subject: [PATCH 21/25] Add masks support for tfrecord

---
 .../tf_detection_api_format/converter.py      | 253 +++++++++++-------
 .../tf_detection_api_format/extractor.py      |  33 ++-
 datumaro/tests/test_tfrecord_format.py        |  31 ++-
 3 files changed, 212 insertions(+), 105 deletions(-)

diff --git a/datumaro/datumaro/plugins/tf_detection_api_format/converter.py b/datumaro/datumaro/plugins/tf_detection_api_format/converter.py
index 340492638fe..ee87485c6c6 100644
--- a/datumaro/datumaro/plugins/tf_detection_api_format/converter.py
+++ b/datumaro/datumaro/plugins/tf_detection_api_format/converter.py
@@ -5,6 +5,7 @@
 
 import codecs
 from collections import OrderedDict
+from itertools import groupby
 import logging as log
 import os
 import os.path as osp
@@ -16,98 +17,32 @@
 from datumaro.components.converter import Converter
 from datumaro.components.cli_plugin import CliPlugin
 from datumaro.util.image import encode_image
+from datumaro.util.mask_tools import merge_masks
 from datumaro.util.tf_util import import_tf as _import_tf
 
 from .format import DetectionApiPath
 tf = _import_tf()
 
 
-# we need it to filter out non-ASCII characters, otherwise training will crash
+# filter out non-ASCII characters, otherwise training will crash
 _printable = set(string.printable)
 def _make_printable(s):
     return ''.join(filter(lambda x: x in _printable, s))
 
-def _make_tf_example(item, get_label_id, get_label, save_images=False):
-    def int64_feature(value):
-        return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
-
-    def int64_list_feature(value):
-        return tf.train.Feature(int64_list=tf.train.Int64List(value=value))
-
-    def bytes_feature(value):
-        return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
-
-    def bytes_list_feature(value):
-        return tf.train.Feature(bytes_list=tf.train.BytesList(value=value))
-
-    def float_list_feature(value):
-        return tf.train.Feature(float_list=tf.train.FloatList(value=value))
-
-
-    features = {
-        'image/source_id': bytes_feature(str(item.id).encode('utf-8')),
-        'image/filename': bytes_feature(
-            ('%s%s' % (item.id, DetectionApiPath.IMAGE_EXT)).encode('utf-8')),
-    }
-
-    if not item.has_image:
-        raise Exception("Failed to export dataset item '%s': "
-            "item has no image info" % item.id)
-    height, width = item.image.size
-
-    features.update({
-        'image/height': int64_feature(height),
-        'image/width': int64_feature(width),
-    })
-
-    features.update({
-        'image/encoded': bytes_feature(b''),
-        'image/format': bytes_feature(b'')
-    })
-    if save_images:
-        if item.has_image and item.image.has_data:
-            fmt = DetectionApiPath.IMAGE_FORMAT
-            buffer = encode_image(item.image.data, DetectionApiPath.IMAGE_EXT)
-
-            features.update({
-                'image/encoded': bytes_feature(buffer),
-                'image/format': bytes_feature(fmt.encode('utf-8')),
-            })
-        else:
-            log.warning("Item '%s' has no image" % item.id)
-
-    xmins = [] # List of normalized left x coordinates in bounding box (1 per box)
-    xmaxs = [] # List of normalized right x coordinates in bounding box (1 per box)
-    ymins = [] # List of normalized top y coordinates in bounding box (1 per box)
-    ymaxs = [] # List of normalized bottom y coordinates in bounding box (1 per box)
-    classes_text = [] # List of string class name of bounding box (1 per box)
-    classes = [] # List of integer class id of bounding box (1 per box)
-
-    boxes = [ann for ann in item.annotations if ann.type is AnnotationType.bbox]
-    for box in boxes:
-        box_label = _make_printable(get_label(box.label))
-
-        xmins.append(box.points[0] / width)
-        xmaxs.append(box.points[2] / width)
-        ymins.append(box.points[1] / height)
-        ymaxs.append(box.points[3] / height)
-        classes_text.append(box_label.encode('utf-8'))
-        classes.append(get_label_id(box.label))
-
-    if boxes:
-        features.update({
-            'image/object/bbox/xmin': float_list_feature(xmins),
-            'image/object/bbox/xmax': float_list_feature(xmaxs),
-            'image/object/bbox/ymin': float_list_feature(ymins),
-            'image/object/bbox/ymax': float_list_feature(ymaxs),
-            'image/object/class/text': bytes_list_feature(classes_text),
-            'image/object/class/label': int64_list_feature(classes),
-        })
+def int64_feature(value):
+    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
+
+def int64_list_feature(value):
+    return tf.train.Feature(int64_list=tf.train.Int64List(value=value))
+
+def bytes_feature(value):
+    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
 
-    tf_example = tf.train.Example(
-        features=tf.train.Features(feature=features))
+def bytes_list_feature(value):
+    return tf.train.Feature(bytes_list=tf.train.BytesList(value=value))
 
-    return tf_example
+def float_list_feature(value):
+    return tf.train.Feature(float_list=tf.train.FloatList(value=value))
 
 class TfDetectionApiConverter(Converter, CliPlugin):
     @classmethod
@@ -115,16 +50,29 @@ def build_cmdline_parser(cls, **kwargs):
         parser = super().build_cmdline_parser(**kwargs)
         parser.add_argument('--save-images', action='store_true',
             help="Save images (default: %(default)s)")
+        parser.add_argument('--save-masks', action='store_true',
+            help="Include instance masks (default: %(default)s)")
         return parser
 
-    def __init__(self, save_images=False):
+    def __init__(self, save_images=False, save_masks=False):
         super().__init__()
 
         self._save_images = save_images
+        self._save_masks = save_masks
 
     def __call__(self, extractor, save_dir):
         os.makedirs(save_dir, exist_ok=True)
 
+        label_categories = extractor.categories().get(AnnotationType.label,
+            LabelCategories())
+        get_label = lambda label_id: label_categories.items[label_id].name \
+            if label_id is not None else ''
+        label_ids = OrderedDict((label.name, 1 + idx)
+            for idx, label in enumerate(label_categories.items))
+        map_label_id = lambda label_id: label_ids.get(get_label(label_id), 0)
+        self._get_label = get_label
+        self._get_label_id = map_label_id
+
         subsets = extractor.subsets()
         if len(subsets) == 0:
             subsets = [ None ]
@@ -136,14 +84,6 @@ def __call__(self, extractor, save_dir):
                 subset_name = DEFAULT_SUBSET_NAME
                 subset = extractor
 
-            label_categories = subset.categories().get(AnnotationType.label,
-                LabelCategories())
-            get_label = lambda label_id: label_categories.items[label_id].name \
-                if label_id is not None else ''
-            label_ids = OrderedDict((label.name, 1 + idx)
-                for idx, label in enumerate(label_categories.items))
-            map_label_id = lambda label_id: label_ids.get(get_label(label_id), 0)
-
             labelmap_path = osp.join(save_dir, DetectionApiPath.LABELMAP_FILE)
             with codecs.open(labelmap_path, 'w', encoding='utf8') as f:
                 for label, idx in label_ids.items():
@@ -157,10 +97,133 @@ def __call__(self, extractor, save_dir):
             anno_path = osp.join(save_dir, '%s.tfrecord' % (subset_name))
             with tf.io.TFRecordWriter(anno_path) as writer:
                 for item in subset:
-                    tf_example = _make_tf_example(
-                        item,
-                        get_label=get_label,
-                        get_label_id=map_label_id,
-                        save_images=self._save_images,
-                    )
+                    tf_example = self._make_tf_example(item)
                     writer.write(tf_example.SerializeToString())
+
+    def _find_instance_parts(self, group, img_width, img_height):
+        boxes = [a for a in group if a.type == AnnotationType.bbox]
+        masks = [a for a in group if a.type == AnnotationType.mask]
+
+        anns = boxes + masks
+        leader = self.find_group_leader(anns)
+        bbox = self._compute_bbox(anns)
+
+        mask = None
+        if self._save_masks:
+            mask = merge_masks([m.image for m in masks])
+
+        return [leader, mask, bbox]
+
+    @staticmethod
+    def find_group_leader(group):
+        return max(group, key=lambda x: x.get_area())
+
+    @staticmethod
+    def _compute_bbox(annotations):
+        boxes = [ann.get_bbox() for ann in annotations]
+        x0 = min((b[0] for b in boxes), default=0)
+        y0 = min((b[1] for b in boxes), default=0)
+        x1 = max((b[0] + b[2] for b in boxes), default=0)
+        y1 = max((b[1] + b[3] for b in boxes), default=0)
+        return [x0, y0, x1 - x0, y1 - y0]
+
+    @staticmethod
+    def _find_instance_anns(annotations):
+        return [a for a in annotations
+            if a.type in { AnnotationType.bbox, AnnotationType.mask }
+        ]
+
+    @classmethod
+    def _find_instances(cls, annotations):
+        instance_anns = cls._find_instance_anns(annotations)
+
+        ann_groups = []
+        for g_id, group in groupby(instance_anns, lambda a: a.group):
+            if not g_id:
+                ann_groups.extend(([a] for a in group))
+            else:
+                ann_groups.append(list(group))
+
+        return ann_groups
+
+    def _export_instances(self, instances, width, height):
+        xmins = [] # List of normalized left x coordinates of bounding boxes (1 per box)
+        xmaxs = [] # List of normalized right x coordinates of bounding boxes (1 per box)
+        ymins = [] # List of normalized top y coordinates of bounding boxes (1 per box)
+        ymaxs = [] # List of normalized bottom y coordinates of bounding boxes (1 per box)
+        classes_text = [] # List of class names of bounding boxes (1 per box)
+        classes = [] # List of class ids of bounding boxes (1 per box)
+        masks = [] # List of PNG-encoded instance masks (1 per box)
+
+        for leader, mask, box in instances:
+            label = _make_printable(self._get_label(leader.label))
+            classes_text.append(label.encode('utf-8'))
+            classes.append(self._get_label_id(leader.label))
+
+            xmins.append(box[0] / width)
+            xmaxs.append((box[0] + box[2]) / width)
+            ymins.append(box[1] / height)
+            ymaxs.append((box[1] + box[3]) / height)
+
+            if self._save_masks:
+                if mask is not None:
+                    mask = encode_image(mask, '.png')
+                else:
+                    mask = b''
+                masks.append(mask)
+
+        result = {}
+        if classes:
+            result = {
+                'image/object/bbox/xmin': float_list_feature(xmins),
+                'image/object/bbox/xmax': float_list_feature(xmaxs),
+                'image/object/bbox/ymin': float_list_feature(ymins),
+                'image/object/bbox/ymax': float_list_feature(ymaxs),
+                'image/object/class/text': bytes_list_feature(classes_text),
+                'image/object/class/label': int64_list_feature(classes),
+            }
+            if masks:
+                result['image/object/mask'] = bytes_list_feature(masks)
+        return result
+
+    def _make_tf_example(self, item):
+        features = {
+            'image/source_id': bytes_feature(str(item.id).encode('utf-8')),
+            'image/filename': bytes_feature(
+                ('%s%s' % (item.id, DetectionApiPath.IMAGE_EXT)).encode('utf-8')),
+        }
+
+        if not item.has_image:
+            raise Exception("Failed to export dataset item '%s': "
+                "item has no image info" % item.id)
+        height, width = item.image.size
+
+        features.update({
+            'image/height': int64_feature(height),
+            'image/width': int64_feature(width),
+        })
+
+        features.update({
+            'image/encoded': bytes_feature(b''),
+            'image/format': bytes_feature(b'')
+        })
+        if self._save_images:
+            if item.has_image and item.image.has_data:
+                fmt = DetectionApiPath.IMAGE_FORMAT
+                buffer = encode_image(item.image.data, DetectionApiPath.IMAGE_EXT)
+
+                features.update({
+                    'image/encoded': bytes_feature(buffer),
+                    'image/format': bytes_feature(fmt.encode('utf-8')),
+                })
+            else:
+                log.warning("Item '%s' has no image" % item.id)
+
+        instances = self._find_instances(item.annotations)
+        instances = [self._find_instance_parts(i, width, height) for i in instances]
+        features.update(self._export_instances(instances, width, height))
+
+        tf_example = tf.train.Example(
+            features=tf.train.Features(feature=features))
+
+        return tf_example
diff --git a/datumaro/datumaro/plugins/tf_detection_api_format/extractor.py b/datumaro/datumaro/plugins/tf_detection_api_format/extractor.py
index 8974c65d805..31cf565d7ca 100644
--- a/datumaro/datumaro/plugins/tf_detection_api_format/extractor.py
+++ b/datumaro/datumaro/plugins/tf_detection_api_format/extractor.py
@@ -10,7 +10,7 @@
 
 from datumaro.components.extractor import (SourceExtractor,
     DEFAULT_SUBSET_NAME, DatasetItem,
-    AnnotationType, Bbox, LabelCategories
+    AnnotationType, Bbox, Mask, LabelCategories
 )
 from datumaro.util.image import Image, decode_image, lazy_image
 from datumaro.util.tf_util import import_tf as _import_tf
@@ -147,6 +147,8 @@ def _parse_tfrecord_file(cls, filepath, subset_name, images_dir):
             labels = tf.sparse.to_dense(
                 parsed_record['image/object/class/text'],
                 default_value=b'').numpy()
+            masks = tf.sparse.to_dense(
+                parsed_record['image/object/mask']).numpy()
 
             for label, label_id in zip(labels, label_ids):
                 label = label.decode('utf-8')
@@ -163,15 +165,28 @@ def _parse_tfrecord_file(cls, filepath, subset_name, images_dir):
                 item_id = osp.splitext(frame_filename)[0]
 
             annotations = []
-            for shape in np.dstack((labels, xmins, ymins, xmaxs, ymaxs))[0]:
+            for shape_id, shape in enumerate(
+                    np.dstack((labels, xmins, ymins, xmaxs, ymaxs))[0]):
                 label = shape[0].decode('utf-8')
-                x = clamp(shape[1] * frame_width, 0, frame_width)
-                y = clamp(shape[2] * frame_height, 0, frame_height)
-                w = clamp(shape[3] * frame_width, 0, frame_width) - x
-                h = clamp(shape[4] * frame_height, 0, frame_height) - y
-                annotations.append(Bbox(x, y, w, h,
-                    label=dataset_labels.get(label)
-                ))
+
+                mask = None
+                if masks:
+                    mask = masks[shape_id]
+
+                if mask is not None:
+                    if isinstance(mask, bytes):
+                        mask = lazy_image(mask, decode_image)
+                    annotations.append(Mask(image=mask,
+                        label=dataset_labels.get(label)
+                    ))
+                else:
+                    x = clamp(shape[1] * frame_width, 0, frame_width)
+                    y = clamp(shape[2] * frame_height, 0, frame_height)
+                    w = clamp(shape[3] * frame_width, 0, frame_width) - x
+                    h = clamp(shape[4] * frame_height, 0, frame_height) - y
+                    annotations.append(Bbox(x, y, w, h,
+                        label=dataset_labels.get(label)
+                    ))
 
             image_size = None
             if frame_height and frame_width:
diff --git a/datumaro/tests/test_tfrecord_format.py b/datumaro/tests/test_tfrecord_format.py
index efbef0fd2b8..0bd29ae4179 100644
--- a/datumaro/tests/test_tfrecord_format.py
+++ b/datumaro/tests/test_tfrecord_format.py
@@ -3,7 +3,7 @@
 from unittest import TestCase
 
 from datumaro.components.extractor import (Extractor, DatasetItem,
-    AnnotationType, Bbox, LabelCategories
+    AnnotationType, Bbox, Mask, LabelCategories
 )
 from datumaro.plugins.tf_detection_api_format.importer import TfDetectionApiImporter
 from datumaro.plugins.tf_detection_api_format.extractor import TfDetectionApiExtractor
@@ -65,6 +65,35 @@ def categories(self):
                 TestExtractor(), TfDetectionApiConverter(save_images=True),
                 test_dir)
 
+    def test_can_save_masks(self):
+        class TestExtractor(Extractor):
+            def __iter__(self):
+                return iter([
+                    DatasetItem(id=1, subset='train', image=np.ones((4, 5, 3)),
+                        annotations=[
+                            Mask(image=np.array([
+                                [1, 0, 0, 1],
+                                [0, 1, 1, 0],
+                                [0, 1, 1, 0],
+                                [1, 0, 0, 1],
+                            ]), label=1),
+                        ]
+                    ),
+                ])
+
+            def categories(self):
+                label_cat = LabelCategories()
+                for label in range(10):
+                    label_cat.add('label_' + str(label))
+                return {
+                    AnnotationType.label: label_cat,
+                }
+
+        with TestDir() as test_dir:
+            self._test_save_and_load(
+                TestExtractor(), TfDetectionApiConverter(save_masks=True),
+                test_dir)
+
     def test_can_save_dataset_with_no_subsets(self):
         class TestExtractor(Extractor):
             def __iter__(self):

From 2cbd7144f4194039fef08418886474c36ef3a1cb Mon Sep 17 00:00:00 2001
From: Zhiltsov Max <zhiltsov.max35@gmail.com>
Date: Tue, 18 Feb 2020 17:26:30 +0300
Subject: [PATCH 22/25] Refactor coco

---
 .../datumaro/plugins/coco_format/converter.py | 25 +++----------------
 1 file changed, 4 insertions(+), 21 deletions(-)

diff --git a/datumaro/datumaro/plugins/coco_format/converter.py b/datumaro/datumaro/plugins/coco_format/converter.py
index 7b91aef4f5d..5c927960f80 100644
--- a/datumaro/datumaro/plugins/coco_format/converter.py
+++ b/datumaro/datumaro/plugins/coco_format/converter.py
@@ -14,7 +14,7 @@
 
 from datumaro.components.converter import Converter
 from datumaro.components.extractor import (DEFAULT_SUBSET_NAME,
-    AnnotationType, Points, Mask
+    AnnotationType, Points
 )
 from datumaro.components.cli_plugin import CliPlugin
 from datumaro.util import find
@@ -194,7 +194,7 @@ def crop_segments(cls, instances, img_width, img_height):
             if inst[1]:
                 inst[1] = sum(new_segments, [])
             else:
-                mask = cls.merge_masks(new_segments)
+                mask = mask_tools.merge_masks(new_segments)
                 inst[2] = mask_tools.mask_to_rle(mask)
 
         return instances
@@ -228,14 +228,14 @@ def find_instance_parts(self, group, img_width, img_height):
             if masks:
                 if mask is not None:
                     masks += [mask]
-                mask = self.merge_masks(masks)
+                mask = mask_tools.merge_masks([m.image for m in masks])
 
             if mask is not None:
                 mask = mask_tools.mask_to_rle(mask)
             polygons = []
         else:
             if masks:
-                mask = self.merge_masks(masks)
+                mask = mask_tools.merge_masks([m.image for m in masks])
                 polygons += mask_tools.mask_to_polygons(mask)
             mask = None
 
@@ -245,23 +245,6 @@ def find_instance_parts(self, group, img_width, img_height):
     def find_group_leader(group):
         return max(group, key=lambda x: x.get_area())
 
-    @staticmethod
-    def merge_masks(masks):
-        if not masks:
-            return None
-
-        def get_mask(m):
-            if isinstance(m, Mask):
-                return m.image
-            else:
-                return m
-
-        binary_mask = get_mask(masks[0])
-        for m in masks[1:]:
-            binary_mask |= get_mask(m)
-
-        return binary_mask
-
     @staticmethod
     def compute_bbox(annotations):
         boxes = [ann.get_bbox() for ann in annotations]

From 7d58f4d42b5cb3f63b565e159027d003c290d881 Mon Sep 17 00:00:00 2001
From: Zhiltsov Max <zhiltsov.max35@gmail.com>
Date: Thu, 20 Feb 2020 13:27:49 +0300
Subject: [PATCH 23/25] Fix comparison

---
 datumaro/datumaro/plugins/tf_detection_api_format/extractor.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/datumaro/datumaro/plugins/tf_detection_api_format/extractor.py b/datumaro/datumaro/plugins/tf_detection_api_format/extractor.py
index 31cf565d7ca..09f47217f39 100644
--- a/datumaro/datumaro/plugins/tf_detection_api_format/extractor.py
+++ b/datumaro/datumaro/plugins/tf_detection_api_format/extractor.py
@@ -170,7 +170,7 @@ def _parse_tfrecord_file(cls, filepath, subset_name, images_dir):
                 label = shape[0].decode('utf-8')
 
                 mask = None
-                if masks:
+                if len(masks) != 0:
                     mask = masks[shape_id]
 
                 if mask is not None:

From 3b6fbf8d13be79b5597f1f6f36e1fc3519b8ca06 Mon Sep 17 00:00:00 2001
From: Zhiltsov Max <zhiltsov.max35@gmail.com>
Date: Thu, 20 Feb 2020 17:17:42 +0300
Subject: [PATCH 24/25] Remove dead code

---
 datumaro/datumaro/util/image.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/datumaro/datumaro/util/image.py b/datumaro/datumaro/util/image.py
index 712a4f789ea..2d465f71a4c 100644
--- a/datumaro/datumaro/util/image.py
+++ b/datumaro/datumaro/util/image.py
@@ -169,8 +169,6 @@ def __init__(self, data=None, path=None, loader=None, cache=None,
         if size is not None:
             assert len(size) == 2 and 0 < size[0] and 0 < size[1], size
             size = tuple(size)
-        else:
-            size = None
         self._size = size # (H, W)
 
         assert path is None or isinstance(path, str)

From 5e0dc3775d75b5b4d3d2fd3f76f27a28bbfdf445 Mon Sep 17 00:00:00 2001
From: Zhiltsov Max <zhiltsov.max35@gmail.com>
Date: Thu, 20 Feb 2020 20:23:52 +0300
Subject: [PATCH 25/25] Extract common code for instances

---
 .../datumaro/plugins/coco_format/converter.py | 33 +++-----------
 .../tf_detection_api_format/converter.py      | 44 ++++---------------
 datumaro/datumaro/plugins/transforms.py       | 44 +++----------------
 datumaro/datumaro/util/annotation_tools.py    | 28 ++++++++++++
 datumaro/tests/test_transforms.py             | 14 +++++-
 5 files changed, 62 insertions(+), 101 deletions(-)
 create mode 100644 datumaro/datumaro/util/annotation_tools.py

diff --git a/datumaro/datumaro/plugins/coco_format/converter.py b/datumaro/datumaro/plugins/coco_format/converter.py
index 5c927960f80..42d5497c471 100644
--- a/datumaro/datumaro/plugins/coco_format/converter.py
+++ b/datumaro/datumaro/plugins/coco_format/converter.py
@@ -20,6 +20,7 @@
 from datumaro.util import find
 from datumaro.util.image import save_image
 import datumaro.util.mask_tools as mask_tools
+import datumaro.util.annotation_tools as anno_tools
 
 from .format import CocoTask, CocoPath
 
@@ -205,8 +206,8 @@ def find_instance_parts(self, group, img_width, img_height):
         masks = [a for a in group if a.type == AnnotationType.mask]
 
         anns = boxes + polygons + masks
-        leader = self.find_group_leader(anns)
-        bbox = self.compute_bbox(anns)
+        leader = anno_tools.find_group_leader(anns)
+        bbox = anno_tools.compute_bbox(anns)
         mask = None
         polygons = [p.points for p in polygons]
 
@@ -241,38 +242,16 @@ def find_instance_parts(self, group, img_width, img_height):
 
         return [leader, polygons, mask, bbox]
 
-    @staticmethod
-    def find_group_leader(group):
-        return max(group, key=lambda x: x.get_area())
-
-    @staticmethod
-    def compute_bbox(annotations):
-        boxes = [ann.get_bbox() for ann in annotations]
-        x0 = min((b[0] for b in boxes), default=0)
-        y0 = min((b[1] for b in boxes), default=0)
-        x1 = max((b[0] + b[2] for b in boxes), default=0)
-        y1 = max((b[1] + b[3] for b in boxes), default=0)
-        return [x0, y0, x1 - x0, y1 - y0]
-
     @staticmethod
     def find_instance_anns(annotations):
         return [a for a in annotations
-            if a.type in { AnnotationType.bbox, AnnotationType.polygon } or \
-                a.type == AnnotationType.mask and a.label is not None
+            if a.type in { AnnotationType.bbox,
+                AnnotationType.polygon, AnnotationType.mask }
         ]
 
     @classmethod
     def find_instances(cls, annotations):
-        instance_anns = cls.find_instance_anns(annotations)
-
-        ann_groups = []
-        for g_id, group in groupby(instance_anns, lambda a: a.group):
-            if not g_id:
-                ann_groups.extend(([a] for a in group))
-            else:
-                ann_groups.append(list(group))
-
-        return ann_groups
+        return anno_tools.find_instances(cls.find_instance_anns(annotations))
 
     def save_annotations(self, item):
         instances = self.find_instances(item.annotations)
diff --git a/datumaro/datumaro/plugins/tf_detection_api_format/converter.py b/datumaro/datumaro/plugins/tf_detection_api_format/converter.py
index ee87485c6c6..2a32d4f151a 100644
--- a/datumaro/datumaro/plugins/tf_detection_api_format/converter.py
+++ b/datumaro/datumaro/plugins/tf_detection_api_format/converter.py
@@ -5,7 +5,6 @@
 
 import codecs
 from collections import OrderedDict
-from itertools import groupby
 import logging as log
 import os
 import os.path as osp
@@ -18,6 +17,8 @@
 from datumaro.components.cli_plugin import CliPlugin
 from datumaro.util.image import encode_image
 from datumaro.util.mask_tools import merge_masks
+from datumaro.util.annotation_tools import (compute_bbox,
+    find_group_leader, find_instances)
 from datumaro.util.tf_util import import_tf as _import_tf
 
 from .format import DetectionApiPath
@@ -100,13 +101,18 @@ def __call__(self, extractor, save_dir):
                     tf_example = self._make_tf_example(item)
                     writer.write(tf_example.SerializeToString())
 
+    @staticmethod
+    def _find_instances(annotations):
+        return find_instances(a for a in annotations
+            if a.type in { AnnotationType.bbox, AnnotationType.mask })
+
     def _find_instance_parts(self, group, img_width, img_height):
         boxes = [a for a in group if a.type == AnnotationType.bbox]
         masks = [a for a in group if a.type == AnnotationType.mask]
 
         anns = boxes + masks
-        leader = self.find_group_leader(anns)
-        bbox = self._compute_bbox(anns)
+        leader = find_group_leader(anns)
+        bbox = compute_bbox(anns)
 
         mask = None
         if self._save_masks:
@@ -114,38 +120,6 @@ def _find_instance_parts(self, group, img_width, img_height):
 
         return [leader, mask, bbox]
 
-    @staticmethod
-    def find_group_leader(group):
-        return max(group, key=lambda x: x.get_area())
-
-    @staticmethod
-    def _compute_bbox(annotations):
-        boxes = [ann.get_bbox() for ann in annotations]
-        x0 = min((b[0] for b in boxes), default=0)
-        y0 = min((b[1] for b in boxes), default=0)
-        x1 = max((b[0] + b[2] for b in boxes), default=0)
-        y1 = max((b[1] + b[3] for b in boxes), default=0)
-        return [x0, y0, x1 - x0, y1 - y0]
-
-    @staticmethod
-    def _find_instance_anns(annotations):
-        return [a for a in annotations
-            if a.type in { AnnotationType.bbox, AnnotationType.mask }
-        ]
-
-    @classmethod
-    def _find_instances(cls, annotations):
-        instance_anns = cls._find_instance_anns(annotations)
-
-        ann_groups = []
-        for g_id, group in groupby(instance_anns, lambda a: a.group):
-            if not g_id:
-                ann_groups.extend(([a] for a in group))
-            else:
-                ann_groups.append(list(group))
-
-        return ann_groups
-
     def _export_instances(self, instances, width, height):
         xmins = [] # List of normalized left x coordinates of bounding boxes (1 per box)
         xmaxs = [] # List of normalized right x coordinates of bounding boxes (1 per box)
diff --git a/datumaro/datumaro/plugins/transforms.py b/datumaro/datumaro/plugins/transforms.py
index 81c5ff50187..693edbc339c 100644
--- a/datumaro/datumaro/plugins/transforms.py
+++ b/datumaro/datumaro/plugins/transforms.py
@@ -3,16 +3,16 @@
 #
 # SPDX-License-Identifier: MIT
 
-from itertools import groupby
 import logging as log
 import os.path as osp
 
 import pycocotools.mask as mask_utils
 
 from datumaro.components.extractor import (Transform, AnnotationType,
-    Mask, RleMask, Polygon, Bbox)
+    RleMask, Polygon, Bbox)
 from datumaro.components.cli_plugin import CliPlugin
 import datumaro.util.mask_tools as mask_tools
+from datumaro.util.annotation_tools import find_group_leader, find_instances
 
 
 class CropCoveredSegments(Transform, CliPlugin):
@@ -125,7 +125,7 @@ def merge_segments(cls, instance, img_width, img_height,
         if not polygons and not masks:
             return []
 
-        leader = cls.find_group_leader(polygons + masks)
+        leader = find_group_leader(polygons + masks)
         instance = []
 
         # Build the resulting mask
@@ -138,9 +138,10 @@ def merge_segments(cls, instance, img_width, img_height,
             instance += polygons # keep unused polygons
 
         if masks:
+            masks = [m.image for m in masks]
             if mask is not None:
                 masks += [mask]
-            mask = cls.merge_masks(masks)
+            mask = mask_tools.merge_masks(masks)
 
         if mask is None:
             return instance
@@ -154,41 +155,10 @@ def merge_segments(cls, instance, img_width, img_height,
         )
         return instance
 
-    @staticmethod
-    def find_group_leader(group):
-        return max(group, key=lambda x: x.get_area())
-
-    @staticmethod
-    def merge_masks(masks):
-        if not masks:
-            return None
-
-        def get_mask(m):
-            if isinstance(m, Mask):
-                return m.image
-            else:
-                return m
-
-        binary_mask = get_mask(masks[0])
-        for m in masks[1:]:
-            binary_mask |= get_mask(m)
-
-        return binary_mask
-
     @staticmethod
     def find_instances(annotations):
-        segment_anns = (a for a in annotations
-            if a.type in {AnnotationType.polygon, AnnotationType.mask}
-        )
-
-        ann_groups = []
-        for g_id, group in groupby(segment_anns, lambda a: a.group):
-            if g_id is None:
-                ann_groups.extend(([a] for a in group))
-            else:
-                ann_groups.append(list(group))
-
-        return ann_groups
+        return find_instances(a for a in annotations
+            if a.type in {AnnotationType.polygon, AnnotationType.mask})
 
 class PolygonsToMasks(Transform, CliPlugin):
     def transform_item(self, item):
diff --git a/datumaro/datumaro/util/annotation_tools.py b/datumaro/datumaro/util/annotation_tools.py
new file mode 100644
index 00000000000..00871b157ec
--- /dev/null
+++ b/datumaro/datumaro/util/annotation_tools.py
@@ -0,0 +1,28 @@
+
+# Copyright (C) 2020 Intel Corporation
+#
+# SPDX-License-Identifier: MIT
+
+from itertools import groupby
+
+
+def find_instances(instance_anns):
+    ann_groups = []
+    for g_id, group in groupby(instance_anns, lambda a: a.group):
+        if not g_id:
+            ann_groups.extend(([a] for a in group))
+        else:
+            ann_groups.append(list(group))
+
+    return ann_groups
+
+def find_group_leader(group):
+    return max(group, key=lambda x: x.get_area())
+
+def compute_bbox(annotations):
+    boxes = [ann.get_bbox() for ann in annotations]
+    x0 = min((b[0] for b in boxes), default=0)
+    y0 = min((b[1] for b in boxes), default=0)
+    x1 = max((b[0] + b[2] for b in boxes), default=0)
+    y1 = max((b[1] + b[3] for b in boxes), default=0)
+    return [x0, y0, x1 - x0, y1 - y0]
\ No newline at end of file
diff --git a/datumaro/tests/test_transforms.py b/datumaro/tests/test_transforms.py
index 11e997b19d7..6260fe517fd 100644
--- a/datumaro/tests/test_transforms.py
+++ b/datumaro/tests/test_transforms.py
@@ -159,8 +159,10 @@ def __iter__(self):
                                     [1, 0, 0, 0, 0],
                                     [1, 1, 1, 0, 0]],
                                 ),
-                                z_order=0),
+                                z_order=0, group=1),
                             Polygon([1, 1, 4, 1, 4, 4, 1, 4],
+                                z_order=1, group=1),
+                            Polygon([0, 0, 0, 2, 2, 2, 2, 0],
                                 z_order=1),
                         ]
                     ),
@@ -178,7 +180,15 @@ def __iter__(self):
                                     [1, 1, 1, 1, 0],
                                     [1, 1, 1, 0, 0]],
                                 ),
-                                z_order=0),
+                                z_order=0, group=1),
+                            Mask(np.array([
+                                    [1, 1, 0, 0, 0],
+                                    [1, 1, 0, 0, 0],
+                                    [0, 0, 0, 0, 0],
+                                    [0, 0, 0, 0, 0],
+                                    [0, 0, 0, 0, 0]],
+                                ),
+                                z_order=1),
                         ]
                     ),
                 ])