PNG mask dumper (#633)

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)

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)

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)

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'),
 )