From ed56ecbcb08b6a8c19aaab8a9ac88201afe59464 Mon Sep 17 00:00:00 2001 From: Zhiltsov Max Date: Wed, 19 Feb 2020 15:09:15 +0300 Subject: [PATCH] Replace mask format support --- cvat/apps/annotation/mask.py | 163 +++++++----------------- cvat/apps/engine/tests/test_rest_api.py | 10 +- 2 files changed, 51 insertions(+), 122 deletions(-) diff --git a/cvat/apps/annotation/mask.py b/cvat/apps/annotation/mask.py index ad1b1a6c8089..77b9f47774b0 100644 --- a/cvat/apps/annotation/mask.py +++ b/cvat/apps/annotation/mask.py @@ -6,130 +6,57 @@ "name": "MASK", "dumpers": [ { - "display_name": "{name} (by class) {format} {version}", + "display_name": "{name} {format} {version}", "format": "ZIP", - "version": "1.0", - "handler": "dump_by_class" + "version": "1.1", + "handler": "dump", }, + ], + "loaders": [ { - "display_name": "{name} (by instance) {format} {version}", + "display_name": "{name} {format} {version}", "format": "ZIP", - "version": "1.0", - "handler": "dump_by_instance" + "version": "1.1", + "handler": "load", }, ], - "loaders": [ - ], } -MASK_BY_CLASS = 0 -MASK_BY_INSTANCE = 1 - -def convert_box_to_polygon(shape): - xtl = shape.points[0] - ytl = shape.points[1] - xbr = shape.points[2] - ybr = shape.points[3] - - return [xtl, ytl, xbr, ytl, xbr, ybr, xtl, ybr] - -def create_mask_colorizer(annotations, colorize_type): - import numpy as np - from collections import OrderedDict - - class MaskColorizer: - - def __init__(self, annotations, colorize_type): - - if colorize_type == MASK_BY_CLASS: - self.colors = self.gen_class_mask_colors(annotations) - elif colorize_type == MASK_BY_INSTANCE: - self.colors = self.gen_instance_mask_colors() - - def generate_pascal_colormap(self, size=256): - # RGB format, (0, 0, 0) used for background - colormap = np.zeros((size, 3), dtype=int) - ind = np.arange(size, dtype=int) - - for shift in reversed(range(8)): - for channel in range(3): - colormap[:, channel] |= ((ind >> channel) & 1) << shift - ind >>= 3 - - return colormap - - def gen_class_mask_colors(self, annotations): - colormap = self.generate_pascal_colormap() - labels = [label[1]["name"] for label in annotations.meta["task"]["labels"] if label[1]["name"] != 'background'] - labels.insert(0, 'background') - label_colors = OrderedDict((label, colormap[idx]) for idx, label in enumerate(labels)) - - return label_colors - - def gen_instance_mask_colors(self): - colormap = self.generate_pascal_colormap() - # The first color is black - instance_colors = OrderedDict((idx, colormap[idx]) for idx in range(len(colormap))) - - return instance_colors - - return MaskColorizer(annotations, colorize_type) - -def dump(file_object, annotations, colorize_type): - - from zipfile import ZipFile, ZIP_STORED - import numpy as np - import os - from pycocotools import mask as maskUtils - import matplotlib.image - import io - - colorizer = create_mask_colorizer(annotations, colorize_type=colorize_type) - if colorize_type == MASK_BY_CLASS: - save_dir = "SegmentationClass" - elif colorize_type == MASK_BY_INSTANCE: - save_dir = "SegmentationObject" - - with ZipFile(file_object, "w", ZIP_STORED) as output_zip: - for frame_annotation in annotations.group_by_frame(): - image_name = frame_annotation.name - annotation_name = "{}.png".format(os.path.splitext(os.path.basename(image_name))[0]) - width = frame_annotation.width - height = frame_annotation.height - - shapes = frame_annotation.labeled_shapes - # convert to mask only rectangles and polygons - shapes = [shape for shape in shapes if shape.type == 'rectangle' or shape.type == 'polygon'] - if not shapes: - continue - shapes = sorted(shapes, key=lambda x: int(x.z_order)) - img_mask = np.zeros((height, width, 3)) - buf_mask = io.BytesIO() - for shape_index, shape in enumerate(shapes): - points = shape.points if shape.type != 'rectangle' else convert_box_to_polygon(shape) - rles = maskUtils.frPyObjects([points], height, width) - rle = maskUtils.merge(rles) - mask = maskUtils.decode(rle) - idx = (mask > 0) - # get corresponding color - if colorize_type == MASK_BY_CLASS: - color = colorizer.colors[shape.label] / 255 - elif colorize_type == MASK_BY_INSTANCE: - color = colorizer.colors[shape_index+1] / 255 - - img_mask[idx] = color - - # write mask - matplotlib.image.imsave(buf_mask, img_mask, format='png') - output_zip.writestr(os.path.join(save_dir, annotation_name), buf_mask.getvalue()) - # Store color map for each class - labels = '\n'.join('{}:{}'.format(label, ','.join(str(i) for i in color)) for label, color in colorizer.colors.items()) - output_zip.writestr('colormap.txt', labels) - -def dump_by_class(file_object, annotations): - - return dump(file_object, annotations, MASK_BY_CLASS) - -def dump_by_instance(file_object, annotations): - return dump(file_object, annotations, MASK_BY_INSTANCE) +def dump(file_object, annotations): + from cvat.apps.dataset_manager.bindings import CvatAnnotationsExtractor + from cvat.apps.dataset_manager.util import make_zip_archive + from datumaro.components.project import Environment + from tempfile import TemporaryDirectory + + env = Environment() + polygons_to_masks = env.transforms.get('polygons_to_masks') + boxes_to_masks = env.transforms.get('boxes_to_masks') + id_from_image = env.transforms.get('id_from_image_name') + + extractor = CvatAnnotationsExtractor('', annotations) + extractor = extractor.transform(polygons_to_masks) + extractor = extractor.transform(boxes_to_masks) + extractor = extractor.transform(id_from_image) + converter = env.make_converter('voc_segmentation', + apply_colormap=True, label_map='source') + with TemporaryDirectory() as temp_dir: + converter(extractor, save_dir=temp_dir) + make_zip_archive(temp_dir, file_object) + +def load(file_object, annotations): + from pyunpack import Archive + from tempfile import TemporaryDirectory + from datumaro.plugins.voc_format.importer import VocImporter + from datumaro.components.project import Environment + from cvat.apps.dataset_manager.bindings import import_dm_annotations + + archive_file = file_object if isinstance(file_object, str) else getattr(file_object, "name") + with TemporaryDirectory() as tmp_dir: + Archive(archive_file).extractall(tmp_dir) + + dm_project = VocImporter()(tmp_dir) + dm_dataset = dm_project.make_dataset() + masks_to_polygons = Environment().transforms.get('masks_to_polygons') + dm_dataset = dm_dataset.transform(masks_to_polygons) + import_dm_annotations(dm_dataset, annotations) diff --git a/cvat/apps/engine/tests/test_rest_api.py b/cvat/apps/engine/tests/test_rest_api.py index 15bfea7abad8..278038a3ca90 100644 --- a/cvat/apps/engine/tests/test_rest_api.py +++ b/cvat/apps/engine/tests/test_rest_api.py @@ -1674,7 +1674,7 @@ def _patch_api_v1_jobs_id_data(self, jid, user, action, data): def _check_response(self, response, data): if not response.status_code in [ status.HTTP_403_FORBIDDEN, status.HTTP_401_UNAUTHORIZED]: - compare_objects(self, data, response.data, ignore_keys=["id"]) + compare_objects(self, data, response.data, ignore_keys=["id"]) def _run_api_v1_jobs_id_annotations(self, owner, assignee, annotator): task, jobs = self._create_task(owner, assignee) @@ -2658,9 +2658,9 @@ def _get_initial_annotation(annotation_format): elif annotation_format == "COCO JSON 1.0": annotations["shapes"] = polygon_shapes_wo_attrs - elif annotation_format == "MASK ZIP 1.0": - annotations["shapes"] = rectangle_shapes_with_attrs + rectangle_shapes_wo_attrs + polygon_shapes_wo_attrs - annotations["tracks"] = rectangle_tracks_with_attrs + rectangle_tracks_wo_attrs + elif annotation_format == "MASK ZIP 1.1": + annotations["shapes"] = rectangle_shapes_wo_attrs + polygon_shapes_wo_attrs + annotations["tracks"] = rectangle_tracks_wo_attrs elif annotation_format == "MOT CSV 1.0": annotations["tracks"] = rectangle_tracks_wo_attrs @@ -2730,6 +2730,8 @@ def _get_initial_annotation(annotation_format): } for loader in annotation_format["loaders"]: + if loader["display_name"] == "MASK ZIP 1.1": + continue # can't really predict the result and check response = self._upload_api_v1_tasks_id_annotations(task["id"], annotator, uploaded_data, "format={}".format(loader["display_name"])) self.assertEqual(response.status_code, HTTP_202_ACCEPTED)