-
Notifications
You must be signed in to change notification settings - Fork 39
/
robot.py
111 lines (84 loc) · 3.57 KB
/
robot.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
110
111
import sys
sys.path.append('pingpong')
from pingpong.pingpongthread import PingPongThread
import cv2
import mediapipe as mp
import numpy as np
from tensorflow.keras.models import load_model
PingPongThreadInstance = PingPongThread(number=2)
PingPongThreadInstance.start()
PingPongThreadInstance.wait_until_full_connect()
actions = ['come', 'away', 'spin']
seq_length = 30
model = load_model('models/model2_1.0.h5')
# MediaPipe hands model
mp_hands = mp.solutions.hands
mp_drawing = mp.solutions.drawing_utils
hands = mp_hands.Hands(
max_num_hands=1,
min_detection_confidence=0.5,
min_tracking_confidence=0.5)
cap = cv2.VideoCapture(0)
seq = []
action_seq = []
last_action = None
while cap.isOpened():
ret, img = cap.read()
if not ret:
break
img = cv2.flip(img, 1)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
result = hands.process(img)
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
if result.multi_hand_landmarks is not None:
for res in result.multi_hand_landmarks:
joint = np.zeros((21, 4))
for j, lm in enumerate(res.landmark):
joint[j] = [lm.x, lm.y, lm.z, lm.visibility]
# Compute angles between joints
v1 = joint[[0,1,2,3,0,5,6,7,0,9,10,11,0,13,14,15,0,17,18,19], :3] # Parent joint
v2 = joint[[1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20], :3] # Child joint
v = v2 - v1 # [20, 3]
# Normalize v
v = v / np.linalg.norm(v, axis=1)[:, np.newaxis]
# Get angle using arcos of dot product
angle = np.arccos(np.einsum('nt,nt->n',
v[[0,1,2,4,5,6,8,9,10,12,13,14,16,17,18],:],
v[[1,2,3,5,6,7,9,10,11,13,14,15,17,18,19],:])) # [15,]
angle = np.degrees(angle) # Convert radian to degree
d = np.concatenate([joint.flatten(), angle])
seq.append(d)
mp_drawing.draw_landmarks(img, res, mp_hands.HAND_CONNECTIONS)
if len(seq) < seq_length:
continue
input_data = np.expand_dims(np.array(seq[-seq_length:], dtype=np.float32), axis=0)
y_pred = model.predict(input_data).squeeze()
i_pred = int(np.argmax(y_pred))
conf = y_pred[i_pred]
if conf < 0.9:
continue
action = actions[i_pred]
action_seq.append(action)
if len(action_seq) < 3:
continue
this_action = '?'
if action_seq[-1] == action_seq[-2] == action_seq[-3]:
this_action = action
if last_action != this_action:
if this_action == 'come':
PingPongThreadInstance.run_motor(1, 5)
PingPongThreadInstance.run_motor(2, -5)
elif this_action == 'away':
PingPongThreadInstance.run_motor(1, -5)
PingPongThreadInstance.run_motor(2, 5)
elif this_action == 'spin':
PingPongThreadInstance.run_motor(1, -5)
PingPongThreadInstance.run_motor(2, -5)
last_action = this_action
cv2.putText(img, f'{this_action.upper()}', org=(int(res.landmark[0].x * img.shape[1]), int(res.landmark[0].y * img.shape[0] + 20)), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, color=(255, 255, 255), thickness=2)
cv2.imshow('img', img)
if cv2.waitKey(1) == ord('q'):
break
PingPongThreadInstance.run_motor(1, 'stop')
PingPongThreadInstance.run_motor(2, 'stop')
PingPongThreadInstance.end()