[Datumaro] COCO 'merge instance polygons' option (#938)

* Add polygon merging option to coco converter
* Add test, refactor coco, add support for cli args
* Drop colormap application in datumaro format
* Add cli support in voc converter
* Add cli support in yolo converter
* Add converter cli options in project cli
* Add image data type conversion in image saving

Author: zhiltsov-max

datumaro/datumaro/cli/project/__init__.py

@@ -54,14 +54,14 @@ def build_import_parser(parser):
         help="Source project format (options: %s)" % (', '.join(importers_list)))
     parser.add_argument('-d', '--dest', default='.', dest='dst_dir',
         help="Directory to save the new project to (default: current dir)")
-    parser.add_argument('extra_args', nargs=argparse.REMAINDER,
-        help="Additional arguments for importer")
     parser.add_argument('-n', '--name', default=None,
         help="Name of the new project (default: same as project dir)")
     parser.add_argument('--overwrite', action='store_true',
         help="Overwrite existing files in the save directory")
     parser.add_argument('--copy', action='store_true',
         help="Make a deep copy instead of saving source links")
+    # parser.add_argument('extra_args', nargs=argparse.REMAINDER,
+    #     help="Additional arguments for importer (pass '-- -h' for help)")
     return parser
 
 def import_command(args):
@@ -111,8 +111,8 @@ def build_export_parser(parser):
         help="Output format")
     parser.add_argument('-p', '--project', dest='project_dir', default='.',
         help="Directory of the project to operate on (default: current dir)")
-    parser.add_argument('--save-images', action='store_true',
-        help="Save images")
+    parser.add_argument('extra_args', nargs=argparse.REMAINDER, default=None,
+        help="Additional arguments for converter (pass '-- -h' for help)")
     return parser
 
 def export_command(args):
@@ -125,7 +125,7 @@ def export_command(args):
         save_dir=dst_dir,
         output_format=args.output_format,
         filter_expr=args.filter,
-        save_images=args.save_images)
+        cmdline_args=args.extra_args)
     log.info("Project exported to '%s' as '%s'" % \
         (dst_dir, args.output_format))
 

datumaro/datumaro/components/converter.py

@@ -4,5 +4,16 @@
 # SPDX-License-Identifier: MIT
 
 class Converter:
+    def __init__(self, cmdline_args=None):
+        pass
+
     def __call__(self, extractor, save_dir):
         raise NotImplementedError()
+
+    def _parse_cmdline(self, cmdline):
+        parser = self.build_cmdline_parser()
+
+        if len(cmdline) != 0 and cmdline[0] == '--':
+            cmdline = cmdline[1:]
+        args = parser.parse_args(cmdline)
+        return vars(args)

datumaro/datumaro/components/converters/datumaro.py

@@ -19,7 +19,6 @@
 )
 from datumaro.components.formats.datumaro import DatumaroPath
 from datumaro.util.image import save_image
-from datumaro.util.mask_tools import apply_colormap
 
 
 def _cast(value, type_conv, default=None):
@@ -118,13 +117,6 @@ def _save_mask(self, mask):
         if mask is None:
             return mask_id
 
-        if self._converter._apply_colormap:
-            categories = self._converter._extractor.categories()
-            categories = categories[AnnotationType.mask]
-            colormap = categories.colormap
-
-            mask = apply_colormap(mask, colormap)
-
         mask_id = self._next_mask_id
         self._next_mask_id += 1
 
@@ -232,12 +224,10 @@ def _convert_points_categories(self, obj):
         return converted
 
 class _Converter:
-    def __init__(self, extractor, save_dir,
-            save_images=False, apply_colormap=False):
+    def __init__(self, extractor, save_dir, save_images=False,):
         self._extractor = extractor
         self._save_dir = save_dir
         self._save_images = save_images
-        self._apply_colormap = apply_colormap
 
     def convert(self):
         os.makedirs(self._save_dir, exist_ok=True)
@@ -282,13 +272,27 @@ def _save_image(self, item):
         save_image(image_path, image)
 
 class DatumaroConverter(Converter):
-    def __init__(self, save_images=False, apply_colormap=False):
+    def __init__(self, save_images=False, cmdline_args=None):
         super().__init__()
-        self._save_images = save_images
-        self._apply_colormap = apply_colormap
+
+        self._options = {
+            'save_images': save_images,
+        }
+
+        if cmdline_args is not None:
+            self._options.update(self._parse_cmdline(cmdline_args))
+
+    @classmethod
+    def build_cmdline_parser(cls, parser=None):
+        import argparse
+        if not parser:
+            parser = argparse.ArgumentParser()
+
+        parser.add_argument('--save-images', action='store_true',
+            help="Save images (default: %(default)s)")
+
+        return parser
 
     def __call__(self, extractor, save_dir):
-        converter = _Converter(extractor, save_dir,
-            apply_colormap=self._apply_colormap,
-            save_images=self._save_images)
+        converter = _Converter(extractor, save_dir, **self._options)
         converter.convert()

datumaro/datumaro/components/converters/ms_coco.py

@@ -14,7 +14,7 @@
 from datumaro.components.extractor import (
     DEFAULT_SUBSET_NAME, AnnotationType, PointsObject, BboxObject
 )
-from datumaro.components.formats.ms_coco import CocoAnnotationType, CocoPath
+from datumaro.components.formats.ms_coco import CocoTask, CocoPath
 from datumaro.util import find
 from datumaro.util.image import save_image
 import datumaro.util.mask_tools as mask_tools
@@ -29,8 +29,9 @@ def _cast(value, type_conv, default=None):
         return default
 
 class _TaskConverter:
-    def __init__(self):
+    def __init__(self, context):
         self._min_ann_id = 1
+        self._context = context
 
         data = {
             'licenses': [],
@@ -191,6 +192,13 @@ def save_annotations(self, item):
                     rle = mask_utils.merge(rles)
                     area = mask_utils.area(rle)
 
+                    if self._context._merge_polygons:
+                        binary_mask = mask_utils.decode(rle).astype(np.bool)
+                        binary_mask = np.asfortranarray(binary_mask, dtype=np.uint8)
+                        segmentation = mask_tools.convert_mask_to_rle(binary_mask)
+                        is_crowd = True
+                        bbox = [int(i) for i in mask_utils.toBbox(rle)]
+
             if ann.group is not None:
                 # Mark the group as visited to prevent repeats
                 for a in annotations[:]:
@@ -201,6 +209,18 @@ def save_annotations(self, item):
                 is_crowd = False
                 segmentation = [ann.get_polygon()]
                 area = ann.area()
+
+                if self._context._merge_polygons:
+                    h, w, _ = item.image.shape
+                    rles = mask_utils.frPyObjects(segmentation, h, w)
+                    rle = mask_utils.merge(rles)
+                    area = mask_utils.area(rle)
+                    binary_mask = mask_utils.decode(rle).astype(np.bool)
+                    binary_mask = np.asfortranarray(binary_mask, dtype=np.uint8)
+                    segmentation = mask_tools.convert_mask_to_rle(binary_mask)
+                    is_crowd = True
+                    bbox = [int(i) for i in mask_utils.toBbox(rle)]
+
             if bbox is None:
                 bbox = ann.get_bbox()
 
@@ -340,22 +360,30 @@ def save_annotations(self, item):
 
 class _Converter:
     _TASK_CONVERTER = {
-        CocoAnnotationType.image_info: _ImageInfoConverter,
-        CocoAnnotationType.instances: _InstancesConverter,
-        CocoAnnotationType.person_keypoints: _KeypointsConverter,
-        CocoAnnotationType.captions: _CaptionsConverter,
-        CocoAnnotationType.labels: _LabelsConverter,
+        CocoTask.image_info: _ImageInfoConverter,
+        CocoTask.instances: _InstancesConverter,
+        CocoTask.person_keypoints: _KeypointsConverter,
+        CocoTask.captions: _CaptionsConverter,
+        CocoTask.labels: _LabelsConverter,
     }
 
-    def __init__(self, extractor, save_dir, save_images=False, task=None):
-        if not task:
-            task = list(self._TASK_CONVERTER.keys())
-        elif task in CocoAnnotationType:
-            task = [task]
-        self._task = task
+    def __init__(self, extractor, save_dir,
+            tasks=None, save_images=False, merge_polygons=False):
+        assert tasks is None or isinstance(tasks, (CocoTask, list))
+        if tasks is None:
+            tasks = list(self._TASK_CONVERTER)
+        elif isinstance(tasks, CocoTask):
+            tasks = [tasks]
+        else:
+            for t in tasks:
+                assert t in CocoTask
+        self._tasks = tasks
+
         self._extractor = extractor
         self._save_dir = save_dir
+
         self._save_images = save_images
+        self._merge_polygons = merge_polygons
 
     def make_dirs(self):
         self._images_dir = osp.join(self._save_dir, CocoPath.IMAGES_DIR)
@@ -365,11 +393,13 @@ def make_dirs(self):
         os.makedirs(self._ann_dir, exist_ok=True)
 
     def make_task_converter(self, task):
-        return self._TASK_CONVERTER[task]()
+        if task not in self._TASK_CONVERTER:
+            raise NotImplementedError()
+        return self._TASK_CONVERTER[task](self)
 
     def make_task_converters(self):
         return {
-            task: self.make_task_converter(task) for task in self._task
+            task: self.make_task_converter(task) for task in self._tasks
         }
 
     def save_image(self, item, filename):
@@ -411,32 +441,56 @@ def convert(self):
                         '%s_%s.json' % (task.name, subset_name)))
 
 class CocoConverter(Converter):
-    def __init__(self, task=None, save_images=False):
+    def __init__(self,
+            tasks=None, save_images=False, merge_polygons=False,
+            cmdline_args=None):
         super().__init__()
-        self._task = task
-        self._save_images = save_images
+
+        self._options = {
+            'tasks': tasks,
+            'save_images': save_images,
+            'merge_polygons': merge_polygons,
+        }
+
+        if cmdline_args is not None:
+            self._options.update(self._parse_cmdline(cmdline_args))
+
+    @staticmethod
+    def _split_tasks_string(s):
+        return [CocoTask[i.strip()] for i in s.split(',')]
+
+    @classmethod
+    def build_cmdline_parser(cls, parser=None):
+        import argparse
+        if not parser:
+            parser = argparse.ArgumentParser()
+
+        parser.add_argument('--save-images', action='store_true',
+            help="Save images (default: %(default)s)")
+        parser.add_argument('--merge-polygons', action='store_true',
+            help="Merge instance polygons into a mask (default: %(default)s)")
+        parser.add_argument('--tasks', type=cls._split_tasks_string,
+            default=None,
+            help="COCO task filter, comma-separated list of {%s} "
+                "(default: all)" % ', '.join([t.name for t in CocoTask]))
+
+        return parser
 
     def __call__(self, extractor, save_dir):
-        converter = _Converter(extractor, save_dir,
-            save_images=self._save_images, task=self._task)
+        converter = _Converter(extractor, save_dir, **self._options)
         converter.convert()
 
-def CocoInstancesConverter(save_images=False):
-    return CocoConverter(CocoAnnotationType.instances,
-        save_images=save_images)
+def CocoInstancesConverter(**kwargs):
+    return CocoConverter(CocoTask.instances, **kwargs)
 
-def CocoImageInfoConverter(save_images=False):
-    return CocoConverter(CocoAnnotationType.image_info,
-        save_images=save_images)
+def CocoImageInfoConverter(**kwargs):
+    return CocoConverter(CocoTask.image_info, **kwargs)
 
-def CocoPersonKeypointsConverter(save_images=False):
-    return CocoConverter(CocoAnnotationType.person_keypoints,
-        save_images=save_images)
+def CocoPersonKeypointsConverter(**kwargs):
+    return CocoConverter(CocoTask.person_keypoints, **kwargs)
 
-def CocoCaptionsConverter(save_images=False):
-    return CocoConverter(CocoAnnotationType.captions,
-        save_images=save_images)
+def CocoCaptionsConverter(**kwargs):
+    return CocoConverter(CocoTask.captions, **kwargs)
 
-def CocoLabelsConverter(save_images=False):
-    return CocoConverter(CocoAnnotationType.labels,
-        save_images=save_images)
+def CocoLabelsConverter(**kwargs):
+    return CocoConverter(CocoTask.labels, **kwargs)