diff --git a/README.md b/README.md index b3d0480a..931105bc 100644 --- a/README.md +++ b/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 | diff --git a/cvat/apps/annotation/mask.py b/cvat/apps/annotation/mask.py index b1713fd0..ad1b1a6c 100644 --- a/cvat/apps/annotation/mask.py +++ b/cvat/apps/annotation/mask.py @@ -6,38 +6,26 @@ format_spec = { "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] @@ -45,12 +33,64 @@ def dump(file_object, annotations): 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 dump(file_object, annotations): 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)