-
-
Notifications
You must be signed in to change notification settings - Fork 3.3k
/
shape.py
109 lines (90 loc) · 3.56 KB
/
shape.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
# MIT License
# Copyright (c) Kentaro Wada
import math
import uuid
import numpy as np
import PIL.Image
import PIL.ImageDraw
from labelme.logger 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, line_width=10, point_size=5):
mask = np.zeros(img_shape[:2], dtype=np.uint8)
mask = PIL.Image.fromarray(mask)
draw = PIL.ImageDraw.Draw(mask)
xy = [tuple(point) for point in points]
if shape_type == "circle":
assert len(xy) == 2, "Shape of shape_type=circle must have 2 points"
(cx, cy), (px, py) = xy
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":
assert len(xy) == 2, "Shape of shape_type=rectangle must have 2 points"
draw.rectangle(xy, outline=1, fill=1)
elif shape_type == "line":
assert len(xy) == 2, "Shape of shape_type=line must have 2 points"
draw.line(xy=xy, fill=1, width=line_width)
elif shape_type == "linestrip":
draw.line(xy=xy, fill=1, width=line_width)
elif shape_type == "point":
assert len(xy) == 1, "Shape of shape_type=point must have 1 points"
cx, cy = xy[0]
r = point_size
draw.ellipse([cx - r, cy - r, cx + r, cy + r], outline=1, fill=1)
else:
assert len(xy) > 2, "Polygon must have points more than 2"
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):
cls = np.zeros(img_shape[:2], dtype=np.int32)
ins = np.zeros_like(cls)
instances = []
for shape in shapes:
points = shape["points"]
label = shape["label"]
group_id = shape.get("group_id")
if group_id is None:
group_id = uuid.uuid1()
shape_type = shape.get("shape_type", None)
cls_name = label
instance = (cls_name, group_id)
if instance not in instances:
instances.append(instance)
ins_id = instances.index(instance) + 1
cls_id = label_name_to_value[cls_name]
mask = shape_to_mask(img_shape[:2], points, shape_type)
cls[mask] = cls_id
ins[mask] = ins_id
return cls, ins
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
def masks_to_bboxes(masks):
if masks.ndim != 3:
raise ValueError("masks.ndim must be 3, but it is {}".format(masks.ndim))
if masks.dtype != bool:
raise ValueError(
"masks.dtype must be bool type, but it is {}".format(masks.dtype)
)
bboxes = []
for mask in masks:
where = np.argwhere(mask)
(y1, x1), (y2, x2) = where.min(0), where.max(0) + 1
bboxes.append((y1, x1, y2, x2))
bboxes = np.asarray(bboxes, dtype=np.float32)
return bboxes