from math import pow from math import sqrt import numpy as np import PIL.Image import PIL.ImageDraw from labelme import logger def polygons_to_mask(img_shape, polygons, 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(polygons) == 2: ((cx, cy), (px, py)) = polygons d = sqrt(pow(cx - px, 2) + pow(cy - py, 2)) draw.ellipse([cx - d, cy - d, cx + d, cy + d], outline=1, fill=1) else: xy = list(map(tuple, polygons)) 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: polygons = 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 = polygons_to_mask(img_shape[:2], polygons, 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