import math import numpy as np import PIL.Image import PIL.ImageDraw from labelme import logger def polygons_to_mask(img_shape, polygons, shape_type=None): logger.warning( "The 'polygons_to_mask' function is deprecated, " "use 'shape_to_mask' instead." ) return shape_to_mask(img_shape, points=polygons, shape_type=shape_type) def shape_to_mask(img_shape, points, shape_type=None): mask = np.zeros(img_shape[:2], dtype=np.uint8) mask = PIL.Image.fromarray(mask) draw = PIL.ImageDraw.Draw(mask) if shape_type == 'circle' and len(points) == 2: (cx, cy), (px, py) = points d = math.sqrt((cx - px) ** 2 + (cy - py) ** 2) draw.ellipse([cx - d, cy - d, cx + d, cy + d], outline=1, fill=1) elif shape_type == 'rectangle' and len(points) == 2: xy = [tuple(point) for point in points] draw.rectangle(xy, outline=1, fill=1) else: xy = [tuple(point) for point in points] draw.polygon(xy=xy, outline=1, fill=1) mask = np.array(mask, dtype=bool) return mask def shapes_to_label(img_shape, shapes, label_name_to_value, type='class'): assert type in ['class', 'instance'] cls = np.zeros(img_shape[:2], dtype=np.int32) if type == 'instance': ins = np.zeros(img_shape[:2], dtype=np.int32) instance_names = ['_background_'] for shape in shapes: points = shape['points'] label = shape['label'] shape_type = shape.get('shape_type', None) if type == 'class': cls_name = label elif type == 'instance': cls_name = label.split('-')[0] if label not in instance_names: instance_names.append(label) ins_id = len(instance_names) - 1 cls_id = label_name_to_value[cls_name] mask = shape_to_mask(img_shape[:2], points, shape_type) cls[mask] = cls_id if type == 'instance': ins[mask] = ins_id if type == 'instance': return cls, ins return cls def labelme_shapes_to_label(img_shape, shapes): logger.warn('labelme_shapes_to_label is deprecated, so please use ' 'shapes_to_label.') label_name_to_value = {'_background_': 0} for shape in shapes: label_name = shape['label'] if label_name in label_name_to_value: label_value = label_name_to_value[label_name] else: label_value = len(label_name_to_value) label_name_to_value[label_name] = label_value lbl = shapes_to_label(img_shape, shapes, label_name_to_value) return lbl, label_name_to_value