stick.py
1.69 KB
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
from PIL import Image, ImageDraw, ImageColor
import numpy as np
from keypoint import OPENPOSE_18
COLORS = [(255, 0, 0, 0),
(0, 255, 0, 0),
(0, 0, 255, 0)]
def kp2stick(kps, size=[256, 256], kp_model=OPENPOSE_18):
# Create canvas
im = Image.fromarray(np.zeros(size + [3], dtype='uint8'))
draw = ImageDraw.Draw(im)
# Draw Body Polygon
body = []
for idx in kp_model.CENTER_BODY:
#point = kps[idx].tolist()
point = kps[idx]
if point[0] <= 0 and point[1] <= 0:
continue
body += [tuple(point)]
draw.polygon(body, fill=COLORS[1])
# Draw Lines
all_lines = [
kp_model.LEFT_LINES,
kp_model.CENTER_LINES,
kp_model.RIGHT_LINES
]
for channel, lines in enumerate(all_lines):
for p1idx, p2idx in lines:
point1 = tuple(list(kps[p1idx]))
point2 = tuple(list(kps[p2idx]))
if (point1[0] <= 0 and point1[1] <= 0) \
or (point2[0] <= 0 and point2[1] <= 0):
continue
draw.line([point1, point2], fill=COLORS[channel], width=4)
# Draw Points
point_size = np.array([3,3])
all_points = [
kp_model.LEFT_POINTS,
kp_model.CENTER_POINTS,
kp_model.RIGHT_POINTS,
]
for channel, points in enumerate(all_points):
for pidx in points:
point = np.array(kps[pidx])
if (point[0] <= 0 and point[1] <= 0):
continue
box = list(point - point_size) + list(point + point_size)
draw.ellipse(box, fill=COLORS[channel])
del draw
return np.array(im)