# Copyright (C) 2020-2021 Intel Corporation
#
# SPDX-License-Identifier: MIT

from enum import Enum, auto
from itertools import chain, groupby
import json
import logging as log
import os
import os.path as osp

import pycocotools.mask as mask_utils

from datumaro.components.converter import Converter
from datumaro.components.dataset import ItemStatus
from datumaro.components.extractor import (
    _COORDINATE_ROUNDING_DIGITS, AnnotationType, DatasetItem, Points,
)
from datumaro.util import cast, find, str_to_bool
from datumaro.util.image import save_image
import datumaro.util.annotation_util as anno_tools
import datumaro.util.mask_tools as mask_tools

from .format import CocoPath, CocoTask


class SegmentationMode(Enum):
    guess = auto()
    polygons = auto()
    mask = auto()

class _TaskConverter:
    def __init__(self, context):
        self._min_ann_id = 1
        self._context = context

        data = {
            'licenses': [],
            'info': {},
            'categories': [],
            'images': [],
            'annotations': []
            }

        data['licenses'].append({
            'name': '',
            'id': 0,
            'url': ''
        })

        data['info'] = {
            'contributor': '',
            'date_created': '',
            'description': '',
            'url': '',
            'version': '',
            'year': ''
        }
        self._data = data

    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):
        h = 0
        w = 0
        if item.has_image and item.image.size:
            h, w = item.image.size

        self._data['images'].append({
            'id': self._get_image_id(item),
            'width': int(w),
            'height': int(h),
            'file_name': cast(filename, str, ''),
            'license': 0,
            'flickr_url': '',
            'coco_url': '',
            'date_captured': 0,
        })

    def save_categories(self, dataset):
        raise NotImplementedError()

    def save_annotations(self, item):
        raise NotImplementedError()

    def write(self, path):
        next_id = self._min_ann_id
        for ann in self.annotations:
            if not ann['id']:
                ann['id'] = next_id
                next_id += 1

        with open(path, 'w', encoding='utf-8') as outfile:
            json.dump(self._data, outfile, ensure_ascii=False)

    @property
    def annotations(self):
        return self._data['annotations']

    @property
    def categories(self):
        return self._data['categories']

    def _get_ann_id(self, annotation):
        ann_id = 0 if self._context._reindex else annotation.id
        if ann_id:
            self._min_ann_id = max(ann_id, self._min_ann_id)
        return ann_id

    @staticmethod
    def _convert_attributes(ann):
        return { k: v for k, v in ann.attributes.items()
            if k not in {'is_crowd', 'score'}
        }

class _ImageInfoConverter(_TaskConverter):
    def is_empty(self):
        return len(self._data['images']) == 0

    def save_categories(self, dataset):
        pass

    def save_annotations(self, item):
        pass

class _CaptionsConverter(_TaskConverter):
    def save_categories(self, dataset):
        pass

    def save_annotations(self, item):
        for ann_idx, ann in enumerate(item.annotations):
            if ann.type != AnnotationType.caption:
                continue

            elem = {
                'id': self._get_ann_id(ann),
                'image_id': self._get_image_id(item),
                'category_id': 0, # NOTE: workaround for a bug in cocoapi
                'caption': ann.caption,
            }
            if 'score' in ann.attributes:
                try:
                    elem['score'] = float(ann.attributes['score'])
                except Exception as e:
                    log.warning("Item '%s', ann #%s: failed to convert "
                        "attribute 'score': %e" % (item.id, ann_idx, e))
            if self._context._allow_attributes:
                attrs = self._convert_attributes(ann)
                if attrs:
                    elem['attributes'] = attrs

            self.annotations.append(elem)

class _InstancesConverter(_TaskConverter):
    def save_categories(self, dataset):
        label_categories = dataset.categories().get(AnnotationType.label)
        if label_categories is None:
            return

        for idx, cat in enumerate(label_categories.items):
            self.categories.append({
                'id': 1 + idx,
                'name': cast(cat.name, str, ''),
                'supercategory': cast(cat.parent, str, ''),
            })

    @classmethod
    def crop_segments(cls, instances, img_width, img_height):
        instances = sorted(instances, key=lambda x: x[0].z_order)

        segment_map = []
        segments = []
        for inst_idx, (_, polygons, mask, _) in enumerate(instances):
            if polygons:
                segment_map.extend(inst_idx for p in polygons)
                segments.extend(polygons)
            elif mask is not None:
                segment_map.append(inst_idx)
                segments.append(mask)

        segments = mask_tools.crop_covered_segments(
            segments, img_width, img_height)

        for inst_idx, inst in enumerate(instances):
            new_segments = [s for si_id, s in zip(segment_map, segments)
                if si_id == inst_idx]

            if not new_segments:
                inst[1] = []
                inst[2] = None
                continue

            if inst[1]:
                inst[1] = sum(new_segments, [])
            else:
                mask = mask_tools.merge_masks(new_segments)
                inst[2] = mask_tools.mask_to_rle(mask)

        return instances

    def find_instance_parts(self, group, img_width, img_height):
        boxes = [a for a in group if a.type == AnnotationType.bbox]
        polygons = [a for a in group if a.type == AnnotationType.polygon]
        masks = [a for a in group if a.type == AnnotationType.mask]

        anns = boxes + polygons + masks
        leader = anno_tools.find_group_leader(anns)
        bbox = anno_tools.max_bbox(anns)
        mask = None
        polygons = [p.points for p in polygons]

        if self._context._segmentation_mode == SegmentationMode.guess:
            use_masks = True == leader.attributes.get('is_crowd',
                find(masks, lambda x: x.label == leader.label) is not None)
        elif self._context._segmentation_mode == SegmentationMode.polygons:
            use_masks = False
        elif self._context._segmentation_mode == SegmentationMode.mask:
            use_masks = True
        else:
            raise NotImplementedError("Unexpected segmentation mode '%s'" % \
                self._context._segmentation_mode)

        if use_masks:
            if polygons:
                mask = mask_tools.rles_to_mask(polygons, img_width, img_height)

            if masks:
                masks = (m.image for m in masks)
                if mask is not None:
                    masks += chain(masks, [mask])
                mask = mask_tools.merge_masks(masks)

            if mask is not None:
                mask = mask_tools.mask_to_rle(mask)
            polygons = []
        else:
            if masks:
                mask = mask_tools.merge_masks(m.image for m in masks)
                polygons += mask_tools.mask_to_polygons(mask)
            mask = None

        return [leader, polygons, mask, bbox]

    @staticmethod
    def find_instance_anns(annotations):
        return [a for a in annotations
            if a.type in { AnnotationType.bbox,
                AnnotationType.polygon, AnnotationType.mask }
        ]

    @classmethod
    def find_instances(cls, annotations):
        return anno_tools.find_instances(cls.find_instance_anns(annotations))

    def save_annotations(self, item):
        instances = self.find_instances(item.annotations)
        if not instances:
            return

        if not item.has_image or not item.image.size:
            log.warning("Item '%s': skipping writing instances "
                "since no image info available" % item.id)
            return
        h, w = item.image.size
        instances = [self.find_instance_parts(i, w, h) for i in instances]

        if self._context._crop_covered:
            instances = self.crop_segments(instances, w, h)

        for instance in instances:
            elem = self.convert_instance(instance, item)
            if elem:
                self.annotations.append(elem)

    def convert_instance(self, instance, item):
        ann, polygons, mask, bbox = instance

        is_crowd = mask is not None
        if is_crowd:
            segmentation = {
                'counts': list(int(c) for c in mask['counts']),
                'size': list(int(c) for c in mask['size'])
            }
        else:
            segmentation = [list(map(float, p)) for p in polygons]

        area = 0
        if segmentation:
            if item.has_image and item.image.size:
                h, w = item.image.size
            else:
                # Here we can guess the image size as
                # it is only needed for the area computation
                w = bbox[0] + bbox[2]
                h = bbox[1] + bbox[3]

            rles = mask_utils.frPyObjects(segmentation, h, w)
            if is_crowd:
                rles = [rles]
            else:
                rles = mask_utils.merge(rles)
            area = mask_utils.area(rles)
        else:
            _, _, w, h = bbox
            segmentation = []
            area = w * h

        elem = {
            'id': self._get_ann_id(ann),
            'image_id': self._get_image_id(item),
            'category_id': cast(ann.label, int, -1) + 1,
            'segmentation': segmentation,
            'area': float(area),
            'bbox': [round(float(n), _COORDINATE_ROUNDING_DIGITS) for n in bbox],
            'iscrowd': int(is_crowd),
        }
        if 'score' in ann.attributes:
            try:
                elem['score'] = float(ann.attributes['score'])
            except Exception as e:
                log.warning("Item '%s': failed to convert attribute "
                    "'score': %e" % (item.id, e))
        if self._context._allow_attributes:
                attrs = self._convert_attributes(ann)
                if attrs:
                    elem['attributes'] = attrs

        return elem

class _KeypointsConverter(_InstancesConverter):
    def save_categories(self, dataset):
        label_categories = dataset.categories().get(AnnotationType.label)
        if label_categories is None:
            return
        point_categories = dataset.categories().get(AnnotationType.points)

        for idx, label_cat in enumerate(label_categories.items):
            cat = {
                'id': 1 + idx,
                'name': cast(label_cat.name, str, ''),
                'supercategory': cast(label_cat.parent, str, ''),
                'keypoints': [],
                'skeleton': [],
            }

            if point_categories is not None:
                kp_cat = point_categories.items.get(idx)
                if kp_cat is not None:
                    cat.update({
                        'keypoints': [str(l) for l in kp_cat.labels],
                        'skeleton': [list(map(int, j)) for j in kp_cat.joints],
                    })
            self.categories.append(cat)

    def save_annotations(self, item):
        point_annotations = [a for a in item.annotations
            if a.type == AnnotationType.points]
        if not point_annotations:
            return

        # Create annotations for solitary keypoints annotations
        for points in self.find_solitary_points(item.annotations):
            instance = [points, [], None, points.get_bbox()]
            elem = super().convert_instance(instance, item)
            elem.update(self.convert_points_object(points))
            self.annotations.append(elem)

        # Create annotations for complete instance + keypoints annotations
        super().save_annotations(item)

    @classmethod
    def find_solitary_points(cls, annotations):
        annotations = sorted(annotations, key=lambda a: a.group)
        solitary_points = []

        for g_id, group in groupby(annotations, lambda a: a.group):
            if not g_id or g_id and not cls.find_instance_anns(group):
                group = [a for a in group if a.type == AnnotationType.points]
                solitary_points.extend(group)

        return solitary_points

    @staticmethod
    def convert_points_object(ann):
        keypoints = []
        points = ann.points
        visibility = ann.visibility
        for index in range(0, len(points), 2):
            kp = points[index : index + 2]
            state = visibility[index // 2].value
            keypoints.extend([*kp, state])

        num_annotated = len([v for v in visibility \
            if v != Points.Visibility.absent])

        return {
            'keypoints': keypoints,
            'num_keypoints': num_annotated,
        }

    def convert_instance(self, instance, item):
        points_ann = find(item.annotations, lambda x: \
            x.type == AnnotationType.points and \
            instance[0].group and x.group == instance[0].group)
        if not points_ann:
            return None

        elem = super().convert_instance(instance, item)
        elem.update(self.convert_points_object(points_ann))

        return elem

class _LabelsConverter(_TaskConverter):
    def save_categories(self, dataset):
        label_categories = dataset.categories().get(AnnotationType.label)
        if label_categories is None:
            return

        for idx, cat in enumerate(label_categories.items):
            self.categories.append({
                'id': 1 + idx,
                'name': cast(cat.name, str, ''),
                'supercategory': cast(cat.parent, str, ''),
            })

    def save_annotations(self, item):
        for ann in item.annotations:
            if ann.type != AnnotationType.label:
                continue

            elem = {
                'id': self._get_ann_id(ann),
                'image_id': self._get_image_id(item),
                'category_id': int(ann.label) + 1,
            }
            if 'score' in ann.attributes:
                try:
                    elem['score'] = float(ann.attributes['score'])
                except Exception as e:
                    log.warning("Item '%s': failed to convert attribute "
                        "'score': %e" % (item.id, e))
            if self._context._allow_attributes:
                attrs = self._convert_attributes(ann)
                if attrs:
                    elem['attributes'] = attrs

            self.annotations.append(elem)

class _StuffConverter(_InstancesConverter):
    pass

class _PanopticConverter(_TaskConverter):
    def write(self, path):
        with open(path, 'w', encoding='utf-8') as outfile:
            json.dump(self._data, outfile, ensure_ascii=False)

    def save_categories(self, dataset):
        label_categories = dataset.categories().get(AnnotationType.label)
        if label_categories is None:
            return

        for idx, cat in enumerate(label_categories.items):
            self.categories.append({
                'id': 1 + idx,
                'name': cast(cat.name, str, ''),
                'supercategory': cast(cat.parent, str, ''),
                'isthing': 0, # TODO: can't represent this information yet
            })

    def save_annotations(self, item):
        if not item.has_image:
            return

        ann_filename = item.id + CocoPath.PANOPTIC_EXT

        segments_info = list()
        masks = []
        next_id = self._min_ann_id
        for ann in item.annotations:
            if ann.type != AnnotationType.mask:
                continue

            if not ann.id:
                ann.id = next_id
                next_id += 1

            segment_info = {}
            segment_info['id'] = ann.id
            segment_info['category_id'] = cast(ann.label, int, -1) + 1
            segment_info['area'] = float(ann.get_area())
            segment_info['bbox'] = [float(p) for p in ann.get_bbox()]
            segment_info['iscrowd'] = cast(ann.attributes.get("is_crowd"), int, 0)
            segments_info.append(segment_info)
            masks.append(ann)

        if not masks:
            return

        pan_format = mask_tools.merge_masks((m.image, m.id) for m in masks)
        save_image(osp.join(self._context._segmentation_dir, ann_filename),
            mask_tools.index2bgr(pan_format), create_dir=True)

        elem = {
            'image_id': self._get_image_id(item),
            'file_name': ann_filename,
            'segments_info': segments_info
        }
        self.annotations.append(elem)

class CocoConverter(Converter):
    @staticmethod
    def _split_tasks_string(s):
        return [CocoTask[i.strip()] for i in s.split(',')]

    @classmethod
    def build_cmdline_parser(cls, **kwargs):
        parser = super().build_cmdline_parser(**kwargs)
        parser.add_argument('--segmentation-mode',
            choices=[m.name for m in SegmentationMode],
            default=SegmentationMode.guess.name,
            help="""
                Save mode for instance segmentation:|n
                - '{sm.guess.name}': guess the mode for each instance,|n
                |s|suse 'is_crowd' attribute as hint|n
                - '{sm.polygons.name}': save polygons,|n
                |s|smerge and convert masks, prefer polygons|n
                - '{sm.mask.name}': save masks,|n
                |s|smerge and convert polygons, prefer masks|n
                Default: %(default)s.
                """.format(sm=SegmentationMode))
        parser.add_argument('--crop-covered', action='store_true',
            help="Crop covered segments so that background objects' "
                "segmentation was more accurate (default: %(default)s)")
        parser.add_argument('--allow-attributes',
            type=str_to_bool, default=True,
            help="Allow export of attributes (default: %(default)s)")
        parser.add_argument('--reindex', type=str_to_bool, default=False,
            help="Assign new indices to images and annotations, "
                "useful to avoid merge conflicts (default: %(default)s)")
        parser.add_argument('--merge-images', type=str_to_bool, default=False,
            help="Save all images into a single "
                "directory (default: %(default)s)")
        parser.add_argument('--tasks', type=cls._split_tasks_string,
            help="COCO task filter, comma-separated list of {%s} "
                "(default: all)" % ', '.join(t.name for t in CocoTask))
        return parser

    DEFAULT_IMAGE_EXT = CocoPath.IMAGE_EXT

    _TASK_CONVERTER = {
        CocoTask.image_info: _ImageInfoConverter,
        CocoTask.instances: _InstancesConverter,
        CocoTask.person_keypoints: _KeypointsConverter,
        CocoTask.captions: _CaptionsConverter,
        CocoTask.labels: _LabelsConverter,
        CocoTask.panoptic: _PanopticConverter,
        CocoTask.stuff: _StuffConverter,
    }

    def __init__(self, extractor, save_dir,
            tasks=None, segmentation_mode=None, crop_covered=False,
            allow_attributes=True, reindex=False, merge_images=False,
            **kwargs):
        super().__init__(extractor, save_dir, **kwargs)

        assert tasks is None or isinstance(tasks, (CocoTask, list, str))
        if isinstance(tasks, CocoTask):
            tasks = [tasks]
        elif isinstance(tasks, str):
            tasks = [CocoTask[tasks]]
        elif tasks:
            for i, t in enumerate(tasks):
                if isinstance(t, str):
                    tasks[i] = CocoTask[t]
                else:
                    assert t in CocoTask, t
        else:
            tasks = set()
        self._tasks = tasks

        assert segmentation_mode is None or \
            isinstance(segmentation_mode, str) or \
            segmentation_mode in SegmentationMode
        if segmentation_mode is None:
            segmentation_mode = SegmentationMode.guess
        if isinstance(segmentation_mode, str):
            segmentation_mode = SegmentationMode[segmentation_mode]
        self._segmentation_mode = segmentation_mode

        self._crop_covered = crop_covered
        self._allow_attributes = allow_attributes
        self._reindex = reindex
        self._merge_images = merge_images

        self._image_ids = {}

        self._patch = None

    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_segmentation_dir(self, subset_name):
        self._segmentation_dir = osp.join(self._save_dir,
            CocoPath.ANNOTATIONS_DIR, 'panoptic_'+ subset_name)
        os.makedirs(self._segmentation_dir, exist_ok=True)

    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):
        return { task: self._make_task_converter(task)
            for task in (self._tasks or self._TASK_CONVERTER) }

    def _get_image_id(self, item):
        image_id = self._image_ids.get(item.id)
        if image_id is None:
            if not self._reindex:
                image_id = cast(item.attributes.get('id'), int,
                    len(self._image_ids) + 1)
            else:
                image_id = len(self._image_ids) + 1
            self._image_ids[item.id] = image_id
        return image_id

    def apply(self):
        self._make_dirs()

        for subset_name, subset in self._extractor.subsets().items():
            task_converters = self._make_task_converters()
            for task_conv in task_converters.values():
                task_conv.save_categories(subset)
            if CocoTask.panoptic in task_converters:
                self._make_segmentation_dir(subset_name)

            for item in subset:
                if self._save_images:
                    if item.has_image:
                        self._save_image(item, subdir=osp.join(self._images_dir,
                            '' if self._merge_images else subset_name))
                    else:
                        log.debug("Item '%s' has no image info", item.id)
                for task_conv in task_converters.values():
                    task_conv.save_image_info(item,
                        self._make_image_filename(item))
                    task_conv.save_annotations(item)

            for task, task_conv in task_converters.items():
                ann_file = osp.join(self._ann_dir,
                    '%s_%s.json' % (task.name, subset_name))

                if task_conv.is_empty() and (not self._tasks or self._patch):
                    if task == CocoTask.panoptic:
                        os.rmdir(self._segmentation_dir)
                    if self._patch:
                        if osp.isfile(ann_file):
                            # Remove subsets that became empty
                            os.remove(ann_file)
                    continue

                task_conv.write(ann_file)

    @classmethod
    def patch(cls, dataset, patch, save_dir, **kwargs):
        for subset in patch.updated_subsets:
            conv = cls(dataset.get_subset(subset), save_dir=save_dir, **kwargs)
            conv._patch = patch
            conv.apply()

        conv = cls(dataset, save_dir=save_dir, **kwargs)
        images_dir = osp.join(save_dir, CocoPath.IMAGES_DIR)
        for (item_id, subset), status in patch.updated_items.items():
            if status != ItemStatus.removed:
                item = patch.data.get(item_id, subset)
            else:
                item = DatasetItem(item_id, subset=subset)

            if not (status == ItemStatus.removed or not item.has_image):
                continue

            # Converter supports saving in separate dirs and common image dir

            image_path = osp.join(images_dir, conv._make_image_filename(item))
            if osp.isfile(image_path):
                os.unlink(image_path)

            image_path = osp.join(images_dir, subset,
                conv._make_image_filename(item))
            if osp.isfile(image_path):
                os.unlink(image_path)


class CocoInstancesConverter(CocoConverter):
    def __init__(self, *args, **kwargs):
        kwargs['tasks'] = CocoTask.instances
        super().__init__(*args, **kwargs)

class CocoImageInfoConverter(CocoConverter):
    def __init__(self, *args, **kwargs):
        kwargs['tasks'] = CocoTask.image_info
        super().__init__(*args, **kwargs)

class CocoPersonKeypointsConverter(CocoConverter):
    def __init__(self, *args, **kwargs):
        kwargs['tasks'] = CocoTask.person_keypoints
        super().__init__(*args, **kwargs)

class CocoCaptionsConverter(CocoConverter):
    def __init__(self, *args, **kwargs):
        kwargs['tasks'] = CocoTask.captions
        super().__init__(*args, **kwargs)

class CocoLabelsConverter(CocoConverter):
    def __init__(self, *args, **kwargs):
        kwargs['tasks'] = CocoTask.labels
        super().__init__(*args, **kwargs)

class CocoPanopticConverter(CocoConverter):
    def __init__(self, *args, **kwargs):
        kwargs['tasks'] = CocoTask.panoptic
        super().__init__(*args, **kwargs)

class CocoStuffConverter(CocoConverter):
    def __init__(self, *args, **kwargs):
        kwargs['tasks'] = CocoTask.stuff
        kwargs['segmentation_mode'] = SegmentationMode.mask
        super().__init__(*args, **kwargs)

datumaro/plugins/coco_format/converter.py
# Copyright (C) 2020-2021 Intel Corporation
#
# SPDX-License-Identifier: MIT

from enum import Enum, auto
from itertools import chain, groupby
import json
import logging as log
import os
import os.path as osp

import pycocotools.mask as mask_utils

from datumaro.components.converter import Converter
from datumaro.components.dataset import ItemStatus
from datumaro.components.extractor import (
    _COORDINATE_ROUNDING_DIGITS, AnnotationType, DatasetItem, Points,
)
from datumaro.util import cast, find, str_to_bool
from datumaro.util.image import save_image
import datumaro.util.annotation_util as anno_tools
import datumaro.util.mask_tools as mask_tools

from .format import CocoPath, CocoTask


class SegmentationMode(Enum):
    guess = auto()
    polygons = auto()
    mask = auto()

class _TaskConverter:
    def __init__(self, context):
        self._min_ann_id = 1
        self._context = context

        data = {
            'licenses': [],
            'info': {},
            'categories': [],
            'images': [],
            'annotations': []
            }

        data['licenses'].append({
            'name': '',
            'id': 0,
            'url': ''
        })

        data['info'] = {
            'contributor': '',
            'date_created': '',
            'description': '',
            'url': '',
            'version': '',
            'year': ''
        }
        self._data = data

    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):
        h = 0
        w = 0
        if item.has_image and item.image.size:
            h, w = item.image.size

        self._data['images'].append({
            'id': self._get_image_id(item),
            'width': int(w),
            'height': int(h),
            'file_name': cast(filename, str, ''),
            'license': 0,
            'flickr_url': '',
            'coco_url': '',
            'date_captured': 0,
        })

    def save_categories(self, dataset):
        raise NotImplementedError()

    def save_annotations(self, item):
        raise NotImplementedError()

    def write(self, path):
        next_id = self._min_ann_id
        for ann in self.annotations:
            if not ann['id']:
                ann['id'] = next_id
                next_id += 1

        with open(path, 'w', encoding='utf-8') as outfile:
            json.dump(self._data, outfile, ensure_ascii=False)

    @property
    def annotations(self):
        return self._data['annotations']

    @property
    def categories(self):
        return self._data['categories']

    def _get_ann_id(self, annotation):
        ann_id = 0 if self._context._reindex else annotation.id
        if ann_id:
            self._min_ann_id = max(ann_id, self._min_ann_id)
        return ann_id

    @staticmethod
    def _convert_attributes(ann):
        return { k: v for k, v in ann.attributes.items()
            if k not in {'is_crowd', 'score'}
        }

class _ImageInfoConverter(_TaskConverter):
    def is_empty(self):
        return len(self._data['images']) == 0

    def save_categories(self, dataset):
        pass

    def save_annotations(self, item):
        pass

class _CaptionsConverter(_TaskConverter):
    def save_categories(self, dataset):
        pass

    def save_annotations(self, item):
        for ann_idx, ann in enumerate(item.annotations):
            if ann.type != AnnotationType.caption:
                continue

            elem = {
                'id': self._get_ann_id(ann),
                'image_id': self._get_image_id(item),
                'category_id': 0, # NOTE: workaround for a bug in cocoapi
                'caption': ann.caption,
            }
            if 'score' in ann.attributes:
                try:
                    elem['score'] = float(ann.attributes['score'])
                except Exception as e:
                    log.warning("Item '%s', ann #%s: failed to convert "
                        "attribute 'score': %e" % (item.id, ann_idx, e))
            if self._context._allow_attributes:
                attrs = self._convert_attributes(ann)
                if attrs:
                    elem['attributes'] = attrs

            self.annotations.append(elem)

class _InstancesConverter(_TaskConverter):
    def save_categories(self, dataset):
        label_categories = dataset.categories().get(AnnotationType.label)
        if label_categories is None:
            return

        for idx, cat in enumerate(label_categories.items):
            self.categories.append({
                'id': 1 + idx,
                'name': cast(cat.name, str, ''),
                'supercategory': cast(cat.parent, str, ''),
            })

    @classmethod
    def crop_segments(cls, instances, img_width, img_height):
        instances = sorted(instances, key=lambda x: x[0].z_order)

        segment_map = []
        segments = []
        for inst_idx, (_, polygons, mask, _) in enumerate(instances):
            if polygons:
                segment_map.extend(inst_idx for p in polygons)
                segments.extend(polygons)
            elif mask is not None:
                segment_map.append(inst_idx)
                segments.append(mask)

        segments = mask_tools.crop_covered_segments(
            segments, img_width, img_height)

        for inst_idx, inst in enumerate(instances):
            new_segments = [s for si_id, s in zip(segment_map, segments)
                if si_id == inst_idx]

            if not new_segments:
                inst[1] = []
                inst[2] = None
                continue

            if inst[1]:
                inst[1] = sum(new_segments, [])
            else:
                mask = mask_tools.merge_masks(new_segments)
                inst[2] = mask_tools.mask_to_rle(mask)

        return instances

    def find_instance_parts(self, group, img_width, img_height):
        boxes = [a for a in group if a.type == AnnotationType.bbox]
        polygons = [a for a in group if a.type == AnnotationType.polygon]
        masks = [a for a in group if a.type == AnnotationType.mask]

        anns = boxes + polygons + masks
        leader = anno_tools.find_group_leader(anns)
        bbox = anno_tools.max_bbox(anns)
        mask = None
        polygons = [p.points for p in polygons]

        if self._context._segmentation_mode == SegmentationMode.guess:
            use_masks = True == leader.attributes.get('is_crowd',
                find(masks, lambda x: x.label == leader.label) is not None)
        elif self._context._segmentation_mode == SegmentationMode.polygons:
            use_masks = False
        elif self._context._segmentation_mode == SegmentationMode.mask:
            use_masks = True
        else:
            raise NotImplementedError("Unexpected segmentation mode '%s'" % \
                self._context._segmentation_mode)

        if use_masks:
            if polygons:
                mask = mask_tools.rles_to_mask(polygons, img_width, img_height)

            if masks:
                masks = (m.image for m in masks)
                if mask is not None:
                    masks = chain(masks, [mask])
                mask = mask_tools.merge_masks(masks)

            if mask is not None:
                mask = mask_tools.mask_to_rle(mask)
            polygons = []
        else:
            if masks:
                mask = mask_tools.merge_masks(m.image for m in masks)
                polygons += mask_tools.mask_to_polygons(mask)
            mask = None

        return [leader, polygons, mask, bbox]

    @staticmethod
    def find_instance_anns(annotations):
        return [a for a in annotations
            if a.type in { AnnotationType.bbox,
                AnnotationType.polygon, AnnotationType.mask }
        ]

    @classmethod
    def find_instances(cls, annotations):
        return anno_tools.find_instances(cls.find_instance_anns(annotations))

    def save_annotations(self, item):
        instances = self.find_instances(item.annotations)
        if not instances:
            return

        if not item.has_image or not item.image.size:
            log.warning("Item '%s': skipping writing instances "
                "since no image info available" % item.id)
            return
        h, w = item.image.size
        instances = [self.find_instance_parts(i, w, h) for i in instances]

        if self._context._crop_covered:
            instances = self.crop_segments(instances, w, h)

        for instance in instances:
            elem = self.convert_instance(instance, item)
            if elem:
                self.annotations.append(elem)

    def convert_instance(self, instance, item):
        ann, polygons, mask, bbox = instance

        is_crowd = mask is not None
        if is_crowd:
            segmentation = {
                'counts': list(int(c) for c in mask['counts']),
                'size': list(int(c) for c in mask['size'])
            }
        else:
            segmentation = [list(map(float, p)) for p in polygons]

        area = 0
        if segmentation:
            if item.has_image and item.image.size:
                h, w = item.image.size
            else:
                # Here we can guess the image size as
                # it is only needed for the area computation
                w = bbox[0] + bbox[2]
                h = bbox[1] + bbox[3]

            rles = mask_utils.frPyObjects(segmentation, h, w)
            if is_crowd:
                rles = [rles]
            else:
                rles = mask_utils.merge(rles)
            area = mask_utils.area(rles)
        else:
            _, _, w, h = bbox
            segmentation = []
            area = w * h

        elem = {
            'id': self._get_ann_id(ann),
            'image_id': self._get_image_id(item),
            'category_id': cast(ann.label, int, -1) + 1,
            'segmentation': segmentation,
            'area': float(area),
            'bbox': [round(float(n), _COORDINATE_ROUNDING_DIGITS) for n in bbox],
            'iscrowd': int(is_crowd),
        }
        if 'score' in ann.attributes:
            try:
                elem['score'] = float(ann.attributes['score'])
            except Exception as e:
                log.warning("Item '%s': failed to convert attribute "
                    "'score': %e" % (item.id, e))
        if self._context._allow_attributes:
                attrs = self._convert_attributes(ann)
                if attrs:
                    elem['attributes'] = attrs

        return elem

class _KeypointsConverter(_InstancesConverter):
    def save_categories(self, dataset):
        label_categories = dataset.categories().get(AnnotationType.label)
        if label_categories is None:
            return
        point_categories = dataset.categories().get(AnnotationType.points)

        for idx, label_cat in enumerate(label_categories.items):
            cat = {
                'id': 1 + idx,
                'name': cast(label_cat.name, str, ''),
                'supercategory': cast(label_cat.parent, str, ''),
                'keypoints': [],
                'skeleton': [],
            }

            if point_categories is not None:
                kp_cat = point_categories.items.get(idx)
                if kp_cat is not None:
                    cat.update({
                        'keypoints': [str(l) for l in kp_cat.labels],
                        'skeleton': [list(map(int, j)) for j in kp_cat.joints],
                    })
            self.categories.append(cat)

    def save_annotations(self, item):
        point_annotations = [a for a in item.annotations
            if a.type == AnnotationType.points]
        if not point_annotations:
            return

        # Create annotations for solitary keypoints annotations
        for points in self.find_solitary_points(item.annotations):
            instance = [points, [], None, points.get_bbox()]
            elem = super().convert_instance(instance, item)
            elem.update(self.convert_points_object(points))
            self.annotations.append(elem)

        # Create annotations for complete instance + keypoints annotations
        super().save_annotations(item)

    @classmethod
    def find_solitary_points(cls, annotations):
        annotations = sorted(annotations, key=lambda a: a.group)
        solitary_points = []

        for g_id, group in groupby(annotations, lambda a: a.group):
            if not g_id or g_id and not cls.find_instance_anns(group):
                group = [a for a in group if a.type == AnnotationType.points]
                solitary_points.extend(group)

        return solitary_points

    @staticmethod
    def convert_points_object(ann):
        keypoints = []
        points = ann.points
        visibility = ann.visibility
        for index in range(0, len(points), 2):
            kp = points[index : index + 2]
            state = visibility[index // 2].value
            keypoints.extend([*kp, state])

        num_annotated = len([v for v in visibility \
            if v != Points.Visibility.absent])

        return {
            'keypoints': keypoints,
            'num_keypoints': num_annotated,
        }

    def convert_instance(self, instance, item):
        points_ann = find(item.annotations, lambda x: \
            x.type == AnnotationType.points and \
            instance[0].group and x.group == instance[0].group)
        if not points_ann:
            return None

        elem = super().convert_instance(instance, item)
        elem.update(self.convert_points_object(points_ann))

        return elem

class _LabelsConverter(_TaskConverter):
    def save_categories(self, dataset):
        label_categories = dataset.categories().get(AnnotationType.label)
        if label_categories is None:
            return

        for idx, cat in enumerate(label_categories.items):
            self.categories.append({
                'id': 1 + idx,
                'name': cast(cat.name, str, ''),
                'supercategory': cast(cat.parent, str, ''),
            })

    def save_annotations(self, item):
        for ann in item.annotations:
            if ann.type != AnnotationType.label:
                continue

            elem = {
                'id': self._get_ann_id(ann),
                'image_id': self._get_image_id(item),
                'category_id': int(ann.label) + 1,
            }
            if 'score' in ann.attributes:
                try:
                    elem['score'] = float(ann.attributes['score'])
                except Exception as e:
                    log.warning("Item '%s': failed to convert attribute "
                        "'score': %e" % (item.id, e))
            if self._context._allow_attributes:
                attrs = self._convert_attributes(ann)
                if attrs:
                    elem['attributes'] = attrs

            self.annotations.append(elem)

class _StuffConverter(_InstancesConverter):
    pass

class _PanopticConverter(_TaskConverter):
    def write(self, path):
        with open(path, 'w', encoding='utf-8') as outfile:
            json.dump(self._data, outfile, ensure_ascii=False)

    def save_categories(self, dataset):
        label_categories = dataset.categories().get(AnnotationType.label)
        if label_categories is None:
            return

        for idx, cat in enumerate(label_categories.items):
            self.categories.append({
                'id': 1 + idx,
                'name': cast(cat.name, str, ''),
                'supercategory': cast(cat.parent, str, ''),
                'isthing': 0, # TODO: can't represent this information yet
            })

    def save_annotations(self, item):
        if not item.has_image:
            return

        ann_filename = item.id + CocoPath.PANOPTIC_EXT

        segments_info = list()
        masks = []
        next_id = self._min_ann_id
        for ann in item.annotations:
            if ann.type != AnnotationType.mask:
                continue

            if not ann.id:
                ann.id = next_id
                next_id += 1

            segment_info = {}
            segment_info['id'] = ann.id
            segment_info['category_id'] = cast(ann.label, int, -1) + 1
            segment_info['area'] = float(ann.get_area())
            segment_info['bbox'] = [float(p) for p in ann.get_bbox()]
            segment_info['iscrowd'] = cast(ann.attributes.get("is_crowd"), int, 0)
            segments_info.append(segment_info)
            masks.append(ann)

        if not masks:
            return

        pan_format = mask_tools.merge_masks((m.image, m.id) for m in masks)
        save_image(osp.join(self._context._segmentation_dir, ann_filename),
            mask_tools.index2bgr(pan_format), create_dir=True)

        elem = {
            'image_id': self._get_image_id(item),
            'file_name': ann_filename,
            'segments_info': segments_info
        }
        self.annotations.append(elem)

class CocoConverter(Converter):
    @staticmethod
    def _split_tasks_string(s):
        return [CocoTask[i.strip()] for i in s.split(',')]

    @classmethod
    def build_cmdline_parser(cls, **kwargs):
        parser = super().build_cmdline_parser(**kwargs)
        parser.add_argument('--segmentation-mode',
            choices=[m.name for m in SegmentationMode],
            default=SegmentationMode.guess.name,
            help="""
                Save mode for instance segmentation:|n
                - '{sm.guess.name}': guess the mode for each instance,|n
                |s|suse 'is_crowd' attribute as hint|n
                - '{sm.polygons.name}': save polygons,|n
                |s|smerge and convert masks, prefer polygons|n
                - '{sm.mask.name}': save masks,|n
                |s|smerge and convert polygons, prefer masks|n
                Default: %(default)s.
                """.format(sm=SegmentationMode))
        parser.add_argument('--crop-covered', action='store_true',
            help="Crop covered segments so that background objects' "
                "segmentation was more accurate (default: %(default)s)")
        parser.add_argument('--allow-attributes',
            type=str_to_bool, default=True,
            help="Allow export of attributes (default: %(default)s)")
        parser.add_argument('--reindex', type=str_to_bool, default=False,
            help="Assign new indices to images and annotations, "
                "useful to avoid merge conflicts (default: %(default)s)")
        parser.add_argument('--merge-images', type=str_to_bool, default=False,
            help="Save all images into a single "
                "directory (default: %(default)s)")
        parser.add_argument('--tasks', type=cls._split_tasks_string,
            help="COCO task filter, comma-separated list of {%s} "
                "(default: all)" % ', '.join(t.name for t in CocoTask))
        return parser

    DEFAULT_IMAGE_EXT = CocoPath.IMAGE_EXT

    _TASK_CONVERTER = {
        CocoTask.image_info: _ImageInfoConverter,
        CocoTask.instances: _InstancesConverter,
        CocoTask.person_keypoints: _KeypointsConverter,
        CocoTask.captions: _CaptionsConverter,
        CocoTask.labels: _LabelsConverter,
        CocoTask.panoptic: _PanopticConverter,
        CocoTask.stuff: _StuffConverter,
    }

    def __init__(self, extractor, save_dir,
            tasks=None, segmentation_mode=None, crop_covered=False,
            allow_attributes=True, reindex=False, merge_images=False,
            **kwargs):
        super().__init__(extractor, save_dir, **kwargs)

        assert tasks is None or isinstance(tasks, (CocoTask, list, str))
        if isinstance(tasks, CocoTask):
            tasks = [tasks]
        elif isinstance(tasks, str):
            tasks = [CocoTask[tasks]]
        elif tasks:
            for i, t in enumerate(tasks):
                if isinstance(t, str):
                    tasks[i] = CocoTask[t]
                else:
                    assert t in CocoTask, t
        else:
            tasks = set()
        self._tasks = tasks

        assert segmentation_mode is None or \
            isinstance(segmentation_mode, str) or \
            segmentation_mode in SegmentationMode
        if segmentation_mode is None:
            segmentation_mode = SegmentationMode.guess
        if isinstance(segmentation_mode, str):
            segmentation_mode = SegmentationMode[segmentation_mode]
        self._segmentation_mode = segmentation_mode

        self._crop_covered = crop_covered
        self._allow_attributes = allow_attributes
        self._reindex = reindex
        self._merge_images = merge_images

        self._image_ids = {}

        self._patch = None

    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_segmentation_dir(self, subset_name):
        self._segmentation_dir = osp.join(self._save_dir,
            CocoPath.ANNOTATIONS_DIR, 'panoptic_'+ subset_name)
        os.makedirs(self._segmentation_dir, exist_ok=True)

    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):
        return { task: self._make_task_converter(task)
            for task in (self._tasks or self._TASK_CONVERTER) }

    def _get_image_id(self, item):
        image_id = self._image_ids.get(item.id)
        if image_id is None:
            if not self._reindex:
                image_id = cast(item.attributes.get('id'), int,
                    len(self._image_ids) + 1)
            else:
                image_id = len(self._image_ids) + 1
            self._image_ids[item.id] = image_id
        return image_id

    def apply(self):
        self._make_dirs()

        for subset_name, subset in self._extractor.subsets().items():
            task_converters = self._make_task_converters()
            for task_conv in task_converters.values():
                task_conv.save_categories(subset)
            if CocoTask.panoptic in task_converters:
                self._make_segmentation_dir(subset_name)

            for item in subset:
                if self._save_images:
                    if item.has_image:
                        self._save_image(item, subdir=osp.join(self._images_dir,
                            '' if self._merge_images else subset_name))
                    else:
                        log.debug("Item '%s' has no image info", item.id)
                for task_conv in task_converters.values():
                    task_conv.save_image_info(item,
                        self._make_image_filename(item))
                    task_conv.save_annotations(item)

            for task, task_conv in task_converters.items():
                ann_file = osp.join(self._ann_dir,
                    '%s_%s.json' % (task.name, subset_name))

                if task_conv.is_empty() and (not self._tasks or self._patch):
                    if task == CocoTask.panoptic:
                        os.rmdir(self._segmentation_dir)
                    if self._patch:
                        if osp.isfile(ann_file):
                            # Remove subsets that became empty
                            os.remove(ann_file)
                    continue

                task_conv.write(ann_file)

    @classmethod
    def patch(cls, dataset, patch, save_dir, **kwargs):
        for subset in patch.updated_subsets:
            conv = cls(dataset.get_subset(subset), save_dir=save_dir, **kwargs)
            conv._patch = patch
            conv.apply()

        conv = cls(dataset, save_dir=save_dir, **kwargs)
        images_dir = osp.join(save_dir, CocoPath.IMAGES_DIR)
        for (item_id, subset), status in patch.updated_items.items():
            if status != ItemStatus.removed:
                item = patch.data.get(item_id, subset)
            else:
                item = DatasetItem(item_id, subset=subset)

            if not (status == ItemStatus.removed or not item.has_image):
                continue

            # Converter supports saving in separate dirs and common image dir

            image_path = osp.join(images_dir, conv._make_image_filename(item))
            if osp.isfile(image_path):
                os.unlink(image_path)

            image_path = osp.join(images_dir, subset,
                conv._make_image_filename(item))
            if osp.isfile(image_path):
                os.unlink(image_path)


class CocoInstancesConverter(CocoConverter):
    def __init__(self, *args, **kwargs):
        kwargs['tasks'] = CocoTask.instances
        super().__init__(*args, **kwargs)

class CocoImageInfoConverter(CocoConverter):
    def __init__(self, *args, **kwargs):
        kwargs['tasks'] = CocoTask.image_info
        super().__init__(*args, **kwargs)

class CocoPersonKeypointsConverter(CocoConverter):
    def __init__(self, *args, **kwargs):
        kwargs['tasks'] = CocoTask.person_keypoints
        super().__init__(*args, **kwargs)

class CocoCaptionsConverter(CocoConverter):
    def __init__(self, *args, **kwargs):
        kwargs['tasks'] = CocoTask.captions
        super().__init__(*args, **kwargs)

class CocoLabelsConverter(CocoConverter):
    def __init__(self, *args, **kwargs):
        kwargs['tasks'] = CocoTask.labels
        super().__init__(*args, **kwargs)

class CocoPanopticConverter(CocoConverter):
    def __init__(self, *args, **kwargs):
        kwargs['tasks'] = CocoTask.panoptic
        super().__init__(*args, **kwargs)

class CocoStuffConverter(CocoConverter):
    def __init__(self, *args, **kwargs):
        kwargs['tasks'] = CocoTask.stuff
        kwargs['segmentation_mode'] = SegmentationMode.mask
        super().__init__(*args, **kwargs)

datumaro/plugins/coco_format/converter.py
diff --git a/datumaro/plugins/coco_format/converter.py b/datumaro/plugins/coco_format/converter.py
index 65a7dfc007..5a76d877db 100644
--- a/datumaro/plugins/coco_format/converter.py
+++ b/datumaro/plugins/coco_format/converter.py
@@ -231,7 +231,7 @@ def find_instance_parts(self, group, img_width, img_height):
             if masks:
                 masks = (m.image for m in masks)
                 if mask is not None:
-                    masks += chain(masks, [mask])
+                    masks = chain(masks, [mask])
                 mask = mask_tools.merge_masks(masks)
 
             if mask is not None:
diff --git a/tests/requirements.py b/tests/requirements.py
index fded451187..ada1e0cf27 100644
--- a/tests/requirements.py
+++ b/tests/requirements.py
@@ -34,6 +34,7 @@ class Requirements:
     DATUM_BUG_314 = "Unsuccessful remap_labels"
     DATUM_BUG_402 = "Troubles running 'remap_labels' on ProjectDataset"
     DATUM_BUG_404 = "custom importer/extractor not loading"
+    DATUM_BUG_425 = "Bug: concatenation for the different types in COCO format"
 
 
 class SkipMessages:
diff --git a/tests/test_coco_format.py b/tests/test_coco_format.py
index 68005b52f8..438083f14d 100644
--- a/tests/test_coco_format.py
+++ b/tests/test_coco_format.py
@@ -1017,3 +1017,44 @@ def test_inplace_save_writes_only_updated_data(self):
             self.assertFalse(osp.isfile(osp.join(path, 'images', 'c', '3.jpg')))
             compare_datasets(self, expected, Dataset.import_from(path, 'coco'),
                 require_images=True, ignored_attrs={'id'})
+
+    @mark_requirement(Requirements. DATUM_BUG_425)
+    def test_can_save_and_load_grouped_masks_and_polygons(self):
+        source_dataset = Dataset.from_iterable([
+            DatasetItem(id=1, image=np.zeros((5, 5, 3)),
+                annotations=[
+                    Mask(np.array([
+                            [0, 0, 0, 0, 0],
+                            [0, 1, 1, 1, 0],
+                            [0, 1, 1, 1, 0],
+                            [0, 1, 1, 1, 0],
+                            [0, 0, 0, 0, 0]],
+                        ),
+                        label=0, id=0, z_order=0, group=1),
+                    Polygon([1, 1, 1, 3, 3, 3, 3, 1],
+                        label=0, id=1, z_order=0, group=1),
+                ]
+            ),
+        ], categories=['label_1'])
+
+        target_dataset = Dataset.from_iterable([
+            DatasetItem(id=1, image=np.zeros((5, 5, 3)),
+                annotations=[
+                    Mask(np.array([
+                            [0, 0, 0, 0, 0],
+                            [0, 1, 1, 1, 0],
+                            [0, 1, 1, 1, 0],
+                            [0, 1, 1, 1, 0],
+                            [0, 0, 0, 0, 0]],
+                        ),
+                        attributes={ 'is_crowd': True },
+                        label=0, id=0, group=1),
+                ], attributes={'id': 1}
+            ),
+        ], categories=['label_1'])
+
+        with TestDir() as test_dir:
+            self._test_save_and_load(source_dataset,
+                 partial(CocoInstancesConverter.convert),
+                 test_dir, target_dataset=target_dataset)
+

Bug: concatenation for the different types in COCO format 
https://github.com/openvinotoolkit/datumaro/blob/935fb0ede3c70a68582c7d13b5f5450e51f81235/datumaro/plugins/coco_format/converter.py#L232-L234
masks has type generator, but then we use += between generator and itertools.chain Python doesn't support such concatenation.
