wangchunlin
/
cvat
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Adding dump for VOC instance mask. (#859)

Browse Source 
* Add mask instance dumper
* Fix bug
* Merge mask instance into mask
* Merge the change into mask
* Create MaskColorizer
* Add dump method
main
vugia truong 6 years ago committed by Nikita Manovich
parent f57586a03c
commit 458bd23222
2 changed files with 93 additions and 37 deletions
Show Stats Download Patch File Download Diff File

1
README.md
Unescape Escape View File

@ -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 |

129
cvat/apps/annotation/mask.py
Unescape Escape View File

@ -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)

Write Preview
Loading…
Cancel
Save