Adding dump for VOC instance mask. (cvat-ai#859)

* Add mask instance dumper
* Fix bug
* Merge mask instance into mask
* Merge the change into mask
* Create MaskColorizer
* Add dump method

Author: gachiemchiep

README.md

@@ -42,6 +42,7 @@ Format selection is possible after clicking on the Upload annotation / Dump anno
 | [YOLO](https://pjreddie.com/darknet/yolo/)                                         | X      | X      |
 | [MS COCO Object Detection](http://cocodataset.org/#format-data)                    | X      | X      |
 | PNG mask                                                                           | X      |        |
+| PNG instance mask                                                                  | X      |        |
 | [TFrecord](https://www.tensorflow.org/tutorials/load_data/tf_records)              | X      | X      |
 | [MOT](https://motchallenge.net/)                                                   | X      | X      |
 | [LabelMe](http://labelme.csail.mit.edu/Release3.0)                                 | X      | X      |

cvat/apps/annotation/mask.py

@@ -6,51 +6,91 @@
     "name": "MASK",
     "dumpers": [
         {
-            "display_name": "{name} {format} {version}",
+            "display_name": "{name} (by class) {format} {version}",
             "format": "ZIP",
             "version": "1.0",
-            "handler": "dump"
+            "handler": "dump_by_class"
+        },
+        {
+            "display_name": "{name} (by instance) {format} {version}",
+            "format": "ZIP",
+            "version": "1.0",
+            "handler": "dump_by_instance"
         },
     ],
     "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
+MASK_BY_CLASS = 0
+MASK_BY_INSTANCE = 1
 
-    def convert_box_to_polygon(points):
+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]
 
-    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))
+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
 
-    with ZipFile(file_object, "w") as output_zip:
+            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])
@@ -63,18 +103,33 @@ def convert_box_to_polygon(points):
             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)
+            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)
-                color = label_colors[shape.label] / 255
                 idx = (mask > 0)
-                img[idx] = color
+                # 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
 
-            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())
+            # 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)