diff --git a/CHANGELOG.md b/CHANGELOG.md index e9b05595..9f578b59 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,7 +18,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Ability to create a custom extractors for unsupported media types - Added in PDF extractor - Added in a command line model manager tester -- Ability to dump/load annotations in several formats from UI (CVAT, Pascal VOC, YOLO, MS COCO) +- Ability to dump/load annotations in several formats from UI (CVAT, Pascal VOC, YOLO, MS COCO, png mask) ### Changed - Outside and keyframe buttons in the side panel for all interpolation shapes (they were only for boxes before) diff --git a/README.md b/README.md index 1a672e90..bf5d6a6b 100644 --- a/README.md +++ b/README.md @@ -38,6 +38,7 @@ Format selection is possible after clicking on the Upload annotation / Dump anno | [Pascal VOC](http://host.robots.ox.ac.uk/pascal/VOC/) | X | X | | [YOLO](https://pjreddie.com/darknet/yolo/) | X | X | | [MS COCO Object Detection](http://cocodataset.org/#format-data) | X | X | +| PNG mask | X | | ## Links - [Intel AI blog: New Computer Vision Tool Accelerates Annotation of Digital Images and Video](https://www.intel.ai/introducing-cvat) diff --git a/cvat/apps/annotation/mask.py b/cvat/apps/annotation/mask.py new file mode 100644 index 00000000..b1713fd0 --- /dev/null +++ b/cvat/apps/annotation/mask.py @@ -0,0 +1,80 @@ +# Copyright (C) 2019 Intel Corporation +# +# SPDX-License-Identifier: MIT + +format_spec = { + "name": "MASK", + "dumpers": [ + { + "display_name": "{name} {format} {version}", + "format": "ZIP", + "version": "1.0", + "handler": "dump" + }, + ], + "loaders": [ + ], +} + +def dump(file_object, annotations): + from zipfile import ZipFile + import numpy as np + import os + from pycocotools import mask as maskUtils + import matplotlib.image + import io + from collections import OrderedDict + + # RGB format, (0, 0, 0) used for background + def genearte_pascal_colormap(size=256): + 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 convert_box_to_polygon(points): + 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] + + colormap = genearte_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)) + + with ZipFile(file_object, "w") 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 = np.zeros((height, width, 3)) + buf = io.BytesIO() + for shape in shapes: + points = shape.points if shape.type != 'rectangle' else convert_box_to_polygon(shape.points) + rles = maskUtils.frPyObjects([points], height, width) + rle = maskUtils.merge(rles) + mask = maskUtils.decode(rle) + color = label_colors[shape.label] / 255 + idx = (mask > 0) + img[idx] = color + + matplotlib.image.imsave(buf, img, format='png') + output_zip.writestr(annotation_name, buf.getvalue()) + labels = '\n'.join('{}:{}'.format(label, ','.join(str(i) for i in color)) for label, color in label_colors.items()) + output_zip.writestr('colormap.txt', labels) diff --git a/cvat/apps/annotation/settings.py b/cvat/apps/annotation/settings.py index 5f685ce9..f9b28d4d 100644 --- a/cvat/apps/annotation/settings.py +++ b/cvat/apps/annotation/settings.py @@ -10,4 +10,5 @@ BUILTIN_FORMATS = ( os.path.join(path_prefix, 'pascal_voc.py'), os.path.join(path_prefix, 'yolo.py'), os.path.join(path_prefix, 'coco.py'), + os.path.join(path_prefix, 'mask.py'), )