forked from jinman-kim/Capstone
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathcreate_dataset.py
73 lines (59 loc) · 2.67 KB
/
create_dataset.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
import cv2
import mediapipe as mp
import numpy as np
import time, os
actions = ['rewind','advance','stop']
secs_for_action = 10
# 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)
created_time = int(time.time())
os.makedirs('dataset', exist_ok=True)
while cap.isOpened():
data = []
for idx, action in enumerate(actions):
ret, img = cap.read()
img = cv2.flip(img, 1)
cv2.putText(img, f'Waiting for collecting {action.upper()} action...', org=(10, 30), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, color=(255, 255, 255), thickness=2)
cv2.imshow('img', img)
cv2.waitKey(1500)
start_time = time.time()
while time.time() - start_time < secs_for_action:
ret, img = cap.read()
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
angle_label = np.array([angle], dtype=np.float32)
angle_label = np.append(angle_label, idx)
data.append(angle_label)
mp_drawing.draw_landmarks(img, res, mp_hands.HAND_CONNECTIONS)
cv2.resize(img, (1000, 750))
cv2.imshow('img', img)
if cv2.waitKey(1) == ord('q'):
break
data = np.array(data)
with open(f'dataset/data_{created_time}.csv', 'w', encoding='UTF-8', newline='') as file:
np.savetxt(file,data, delimiter=',')
break