[Datumaro] COCO 'merge instance polygons' option

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)