[Datumaro] Update LabelMe format (#1296)

* Little refactoring

* Add LabelMe format

* Add usernames

* Update tests

* Add extractor test

datumaro/datumaro/components/project.py

@@ -367,6 +367,8 @@ class Dataset(Extractor):
     def get(self, item_id, subset=None, path=None):
         if path:
             raise KeyError("Requested dataset item path is not found")
+        if subset is None:
+            subset = ''
         return self._subsets[subset].items[item_id]
 
     def put(self, item, item_id=None, subset=None, path=None):

datumaro/datumaro/plugins/labelme_format.py

@@ -59,7 +59,9 @@ class LabelMeExtractor(SourceExtractor):
 
     def _parse(self, path):
         categories = {
-            AnnotationType.label: LabelCategories(attributes={'occluded'})
+            AnnotationType.label: LabelCategories(attributes={
+                'occluded', 'username'
+            })
         }
 
         items = []
@@ -136,10 +138,17 @@ class LabelMeExtractor(SourceExtractor):
             if deleted_elem is not None and deleted_elem.text:
                 deleted = bool(int(deleted_elem.text))
 
+            user = ''
+
             poly_elem = obj_elem.find('polygon')
             segm_elem = obj_elem.find('segm')
             type_elem = obj_elem.find('type') # the only value is 'bounding_box'
             if poly_elem is not None:
+                user_elem = poly_elem.find('username')
+                if user_elem is not None and user_elem.text:
+                    user = user_elem.text
+                attributes.append(('username', user))
+
                 points = []
                 for point_elem in poly_elem.iter('pt'):
                     x = float(point_elem.find('x').text)
@@ -153,20 +162,25 @@ class LabelMeExtractor(SourceExtractor):
                     ymin = min(points[1::2])
                     ymax = max(points[1::2])
                     ann_items.append(Bbox(xmin, ymin, xmax - xmin, ymax - ymin,
-                        label=label, attributes=attributes,
+                        label=label, attributes=attributes, id=obj_id,
                     ))
                 else:
                     ann_items.append(Polygon(points,
-                        label=label, attributes=attributes,
+                        label=label, attributes=attributes, id=obj_id,
                     ))
             elif segm_elem is not None:
+                user_elem = segm_elem.find('username')
+                if user_elem is not None and user_elem.text:
+                    user = user_elem.text
+                attributes.append(('username', user))
+
                 mask_path = osp.join(dataset_root, LabelMePath.MASKS_DIR,
                     segm_elem.find('mask').text)
                 if not osp.isfile(mask_path):
                     raise Exception("Can't find mask at '%s'" % mask_path)
                 mask = load_mask(mask_path)
                 mask = np.any(mask, axis=2)
-                ann_items.append(Mask(image=mask, label=label,
+                ann_items.append(Mask(image=mask, label=label, id=obj_id,
                     attributes=attributes))
 
             if not deleted:
@@ -368,7 +382,7 @@ class LabelMeConverter(Converter, CliPlugin):
             ET.SubElement(obj_elem, 'deleted').text = '0'
             ET.SubElement(obj_elem, 'verified').text = '0'
             ET.SubElement(obj_elem, 'occluded').text = \
-                'yes' if ann.attributes.get('occluded') == True else 'no'
+                'yes' if ann.attributes.pop('occluded', '') == True else 'no'
             ET.SubElement(obj_elem, 'date').text = ''
             ET.SubElement(obj_elem, 'id').text = str(obj_id)
 
@@ -390,7 +404,8 @@ class LabelMeConverter(Converter, CliPlugin):
                     ET.SubElement(point_elem, 'x').text = '%.2f' % x
                     ET.SubElement(point_elem, 'y').text = '%.2f' % y
 
-                ET.SubElement(poly_elem, 'username').text = ''
+                ET.SubElement(poly_elem, 'username').text = \
+                    str(ann.attributes.pop('username', ''))
             elif ann.type == AnnotationType.polygon:
                 poly_elem = ET.SubElement(obj_elem, 'polygon')
                 for x, y in zip(ann.points[::2], ann.points[1::2]):
@@ -398,7 +413,8 @@ class LabelMeConverter(Converter, CliPlugin):
                     ET.SubElement(point_elem, 'x').text = '%.2f' % x
                     ET.SubElement(point_elem, 'y').text = '%.2f' % y
 
-                ET.SubElement(poly_elem, 'username').text = ''
+                ET.SubElement(poly_elem, 'username').text = \
+                    str(ann.attributes.pop('username', ''))
             elif ann.type == AnnotationType.mask:
                 mask_filename = '%s_mask_%s.png' % (item.id, obj_id)
                 save_image(osp.join(subset_dir, LabelMePath.MASKS_DIR,
@@ -416,13 +432,14 @@ class LabelMeConverter(Converter, CliPlugin):
                     '%.2f' % (bbox[0] + bbox[2])
                 ET.SubElement(box_elem, 'ymax').text = \
                     '%.2f' % (bbox[1] + bbox[3])
+
+                ET.SubElement(segm_elem, 'username').text = \
+                    str(ann.attributes.pop('username', ''))
             else:
                 raise NotImplementedError("Unknown shape type '%s'" % ann.type)
 
             attrs = []
             for k, v in ann.attributes.items():
-                if k == 'occluded':
-                    continue
                 if isinstance(v, bool):
                     attrs.append(k)
                 else:

datumaro/tests/assets/labelme_dataset/img1.xml

@@ -0,0 +1 @@
+<annotation><filename>img1.png</filename><folder>example_folder</folder><source><sourceImage>The MIT-CSAIL database of objects and scenes</sourceImage><sourceAnnotation>LabelMe Webtool</sourceAnnotation></source><object><name>window</name><deleted>0</deleted><verified>0</verified><date>25-May-2012 00:09:48</date><id>0</id><polygon><username>admin</username><pt><x>43</x><y>34</y></pt><pt><x>45</x><y>34</y></pt><pt><x>45</x><y>37</y></pt><pt><x>43</x><y>37</y></pt></polygon><parts><hasparts/><ispartof/></parts></object><imagesize><nrows>77</nrows><ncols>102</ncols></imagesize><object><name>license plate</name><deleted>0</deleted><verified>0</verified><occluded>no</occluded><attributes/><parts><hasparts/><ispartof/></parts><date>27-Jul-2014 02:58:50</date><id>1</id><segm><username>brussell</username><box><xmin>58</xmin><ymin>66</ymin><xmax>62</xmax><ymax>68</ymax></box><mask>img1_mask_1.png</mask><scribbles><xmin>58</xmin><ymin>66</ymin><xmax>62</xmax><ymax>68</ymax><scribble_name>img1_scribble_1.png</scribble_name></scribbles></segm></object><object><name>o1</name><deleted>0</deleted><verified>0</verified><occluded>yes</occluded><attributes>a1</attributes><parts><hasparts>3,4</hasparts><ispartof></ispartof></parts><date>15-Nov-2019 14:38:51</date><id>2</id><polygon><username>anonymous</username><pt><x>30</x><y>12</y></pt><pt><x>42</x><y>21</y></pt><pt><x>24</x><y>26</y></pt><pt><x>15</x><y>22</y></pt><pt><x>18</x><y>14</y></pt><pt><x>22</x><y>12</y></pt><pt><x>27</x><y>12</y></pt></polygon></object><object><name>q1</name><deleted>0</deleted><verified>0</verified><occluded>no</occluded><attributes>kj</attributes><parts><hasparts></hasparts><ispartof>2</ispartof></parts><date>15-Nov-2019 14:39:00</date><id>3</id><polygon><username>anonymous</username><pt><x>35</x><y>21</y></pt><pt><x>43</x><y>22</y></pt><pt><x>40</x><y>28</y></pt><pt><x>28</x><y>31</y></pt><pt><x>31</x><y>22</y></pt><pt><x>32</x><y>25</y></pt></polygon></object><object><name>b1</name><deleted>0</deleted><verified>0</verified><occluded>yes</occluded><attributes>hg</attributes><parts><hasparts></hasparts><ispartof>2</ispartof></parts><date>15-Nov-2019 14:39:09</date><id>4</id><type>bounding_box</type><polygon><username>anonymous</username><pt><x>13</x><y>19</y></pt><pt><x>23</x><y>19</y></pt><pt><x>23</x><y>30</y></pt><pt><x>13</x><y>30</y></pt></polygon></object><object><name>m1</name><deleted>0</deleted><verified>0</verified><occluded>no</occluded><attributes>d</attributes><parts><hasparts>6</hasparts><ispartof></ispartof></parts><date>15-Nov-2019 14:39:30</date><id>5</id><type>bounding_box</type><segm><username>anonymous</username><box><xmin>56</xmin><ymin>14</ymin><xmax>70</xmax><ymax>23</ymax></box><mask>img1_mask_5.png</mask><scribbles><xmin>55</xmin><ymin>13</ymin><xmax>70</xmax><ymax>23</ymax><scribble_name>img1_scribble_5.png</scribble_name></scribbles></segm></object><object><name>hg</name><deleted>0</deleted><verified>0</verified><occluded>no</occluded><attributes>gfd lkj lkj hi</attributes><parts><hasparts></hasparts><ispartof>5</ispartof></parts><date>15-Nov-2019 14:41:57</date><id>6</id><polygon><username>anonymous</username><pt><x>64</x><y>21</y></pt><pt><x>74</x><y>24</y></pt><pt><x>72</x><y>32</y></pt><pt><x>62</x><y>34</y></pt><pt><x>60</x><y>27</y></pt><pt><x>62</x><y>22</y></pt></polygon></object></annotation>

datumaro/tests/test_labelme_format.py

@@ -1,11 +1,14 @@
 import numpy as np
+import os.path as osp
 
 from unittest import TestCase
 
 from datumaro.components.extractor import (Extractor, DatasetItem,
     AnnotationType, Bbox, Mask, Polygon, LabelCategories
 )
-from datumaro.plugins.labelme_format import LabelMeImporter, LabelMeConverter
+from datumaro.components.project import Dataset
+from datumaro.plugins.labelme_format import LabelMeExtractor, LabelMeImporter, \
+    LabelMeConverter
 from datumaro.util.test_utils import TestDir, compare_datasets
 
 
@@ -35,7 +38,8 @@ class LabelMeConverterTest(TestCase):
                             Polygon([0, 4, 4, 4, 5, 6], label=3, attributes={
                                 'occluded': True
                             }),
-                            Mask(np.array([[0, 1], [1, 0], [1, 1]]), group=2),
+                            Mask(np.array([[0, 1], [1, 0], [1, 1]]), group=2,
+                                attributes={ 'username': 'test' }),
                             Bbox(1, 2, 3, 4, group=3),
                             Mask(np.array([[0, 0], [0, 0], [1, 1]]), group=3,
                                 attributes={ 'occluded': True }
@@ -58,20 +62,28 @@ class LabelMeConverterTest(TestCase):
                     DatasetItem(id=1, subset='train',
                         image=np.ones((16, 16, 3)),
                         annotations=[
-                            Bbox(0, 4, 4, 8, label=0, group=2, attributes={
+                            Bbox(0, 4, 4, 8, label=0, group=2, id=0,
-                                'occluded': False
+                                attributes={
-                            }),
+                                    'occluded': False, 'username': '',
-                            Polygon([0, 4, 4, 4, 5, 6], label=1, attributes={
+                                }
-                                'occluded': True
+                            ),
-                            }),
+                            Polygon([0, 4, 4, 4, 5, 6], label=1, id=1,
+                                attributes={
+                                    'occluded': True, 'username': '',
+                                }
+                            ),
                             Mask(np.array([[0, 1], [1, 0], [1, 1]]), group=2,
-                                attributes={ 'occluded': False }
+                                id=2, attributes={
+                                    'occluded': False, 'username': 'test'
+                                }
                             ),
-                            Bbox(1, 2, 3, 4, group=1, attributes={
+                            Bbox(1, 2, 3, 4, group=1, id=3, attributes={
-                                'occluded': False
+                                'occluded': False, 'username': '',
                             }),
                             Mask(np.array([[0, 0], [0, 0], [1, 1]]), group=1,
-                                attributes={ 'occluded': True }
+                                id=4, attributes={
+                                    'occluded': True, 'username': ''
+                                }
                             ),
                         ]
                     ),
@@ -90,31 +102,113 @@ class LabelMeConverterTest(TestCase):
                 SrcExtractor(), LabelMeConverter(save_images=True),
                 test_dir, target_dataset=DstExtractor())
 
-class LabelMeImporterTest(TestCase):
+
-    def test_can_detect(self):
+DUMMY_DATASET_DIR = osp.join(osp.dirname(__file__), 'assets', 'labelme_dataset')
-        class TestExtractor(Extractor):
+
+class LabelMeExtractorTest(TestCase):
+    def test_can_load(self):
+        class DstExtractor(Extractor):
             def __iter__(self):
+                img1 = np.ones((77, 102, 3)) * 255
+                img1[6:32, 7:41] = 0
+
+                mask1 = np.zeros((77, 102), dtype=int)
+                mask1[67:69, 58:63] = 1
+
+                mask2 = np.zeros((77, 102), dtype=int)
+                mask2[13:25, 54:71] = [
+                    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                    [0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0],
+                    [0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0],
+                    [0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0],
+                    [0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0],
+                    [0, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 1, 1, 0, 0],
+                    [0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 0, 0, 0],
+                    [0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
+                    [0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
+                    [0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
+                    [0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
+                    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+                ]
+
                 return iter([
-                    DatasetItem(id=1, subset='train',
+                    DatasetItem(id='img1', image=img1,
-                        image=np.ones((16, 16, 3)),
                         annotations=[
-                            Bbox(0, 4, 4, 8, label=2),
+                            Polygon([43, 34, 45, 34, 45, 37, 43, 37],
+                                label=0, id=0,
+                                attributes={
+                                    'occluded': False,
+                                    'username': 'admin'
+                                }
+                            ),
+                            Mask(mask1, label=1, id=1,
+                                attributes={
+                                    'occluded': False,
+                                    'username': 'brussell'
+                                }
+                            ),
+                            Polygon([30, 12, 42, 21, 24, 26, 15, 22, 18, 14, 22, 12, 27, 12],
+                                label=2, group=2, id=2,
+                                attributes={
+                                    'a1': '1',
+                                    'occluded': True,
+                                    'username': 'anonymous'
+                                }
+                            ),
+                            Polygon([35, 21, 43, 22, 40, 28, 28, 31, 31, 22, 32, 25],
+                                label=3, group=2, id=3,
+                                attributes={
+                                    'kj': '1',
+                                    'occluded': False,
+                                    'username': 'anonymous'
+                                }
+                            ),
+                            Bbox(13, 19, 10, 11, label=4, group=2, id=4,
+                                attributes={
+                                    'hg': '1',
+                                    'occluded': True,
+                                    'username': 'anonymous'
+                                }
+                            ),
+                            Mask(mask2, label=5, group=1, id=5,
+                                attributes={
+                                    'd': '1',
+                                    'occluded': False,
+                                    'username': 'anonymous'
+                                }
+                            ),
+                            Polygon([64, 21, 74, 24, 72, 32, 62, 34, 60, 27, 62, 22],
+                                label=6, group=1, id=6,
+                                attributes={
+                                    'gfd lkj lkj hi': '1',
+                                    'occluded': False,
+                                    'username': 'anonymous'
+                                }
+                            ),
                         ]
                     ),
                 ])
 
             def categories(self):
                 label_cat = LabelCategories()
-                for label in range(10):
+                label_cat.add('window')
-                    label_cat.add('label_' + str(label))
+                label_cat.add('license plate')
+                label_cat.add('o1')
+                label_cat.add('q1')
+                label_cat.add('b1')
+                label_cat.add('m1')
+                label_cat.add('hg')
                 return {
                     AnnotationType.label: label_cat,
                 }
 
-        def generate_dummy(path):
+        parsed = Dataset.from_extractors(LabelMeExtractor(DUMMY_DATASET_DIR))
-            LabelMeConverter()(TestExtractor(), save_dir=path)
+        compare_datasets(self, expected=DstExtractor(), actual=parsed)
 
-        with TestDir() as test_dir:
+class LabelMeImporterTest(TestCase):
-            generate_dummy(test_dir)
+    def test_can_detect(self):
+        self.assertTrue(LabelMeImporter.detect(DUMMY_DATASET_DIR))
 
-            self.assertTrue(LabelMeImporter.detect(test_dir))
+    def test_can_import(self):
+        parsed = LabelMeImporter()(DUMMY_DATASET_DIR).make_dataset()
+        self.assertEqual(1, len(parsed))