forked from MrEliptik/HandPose
-
Notifications
You must be signed in to change notification settings - Fork 0
/
AddPose.py
157 lines (125 loc) · 5.35 KB
/
AddPose.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
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
import numpy as np
import cv2
import os
from utils import detector_utils as detector_utils
import tensorflow as tf
def main():
currentPath = ''
currentExample = ''
print('Do you want to : \n 1 - Add new pose \
\n 2 - Add examples to existing pose \
\n 3 - Add garbage examples')
menu_choice = input()
while(menu_choice != '1' and menu_choice != '2' and menu_choice != '3'):
print('Please enter 1 or 2')
menu_choice = input()
if(menu_choice == '1'):
print('Enter a name for the pose you want to add :')
name_pose = input()
# Create folder for pose
if not os.path.exists('Poses/' + name_pose):
os.makedirs('Poses/' + name_pose + '/' + name_pose + '_1/')
currentPath = 'Poses/' + name_pose + '/' + name_pose + '_1/'
currentExample = name_pose + '_1_'
elif(menu_choice == '2'):
# Display current poses
dirs = os.listdir('Poses/')
dirs_choice = ''
possible_choices = []
i = 1
for _dir in dirs:
dirs_choice += str(i) + ' - ' + str(_dir) + ' / '
possible_choices.append(str(i))
i+=1
# Ask user to choose to which pose to add examples
print('Choose one of the following pose:')
print(dirs_choice)
choice = input()
while(not choice in possible_choices and dirs[int(choice)-1]=='garbage'):
print('Please enter one of the following (not garbage): ' + str(possible_choices))
choice = input()
# Count number of files to increment new example directory
subdirs = os.listdir('Poses/' + dirs[int(choice)-1] + '/')
index = len(subdirs) + 1
# Create new example directory
if not os.path.exists('Poses/' + dirs[int(choice)-1] + '/' + dirs[int(choice)-1] + '_' + str(index) + '/'):
os.makedirs('Poses/' + dirs[int(choice)-1] + '/' + dirs[int(choice)-1] + '_' + str(index) + '/')
#Update current path
currentPath = 'Poses/' + dirs[int(choice)-1] + '/' + dirs[int(choice)-1] + '_' + str(index) + '/'
currentExample = dirs[int(choice)-1] + '_' + str(index) + '_'
name_pose = dirs[int(choice) - 1]
elif(menu_choice == '3'):
# Create folder for pose
if not os.path.exists('Poses/Garbage/'):
os.makedirs('Poses/Garbage/Garbage_1/')
currentPath = 'Poses/Garbage/Garbage_1/'
currentExample = 'Garbage_1_'
name_pose = 'Garbage'
else:
# Count number of files to increment new example directory
subdirs = os.listdir('Poses/Garbage/')
index = len(subdirs) + 1
# Create new example directory
if not os.path.exists('Poses/Garbage/Garbage_' + str(index) + '/'):
os.makedirs('Poses/Garbage/Garbage_' + str(index) + '/')
#Update current path
currentPath = 'Poses/Garbage/Garbage_' + str(index) + '/'
currentExample ='Garbage_' + str(index) + '_'
name_pose = 'Garbage'
print('You\'ll now be prompted to record the pose you want to add. \n \
Please place your hand beforehand facing the camera, and press any key when ready. \n \
When finished press \'q\'.')
input()
# Begin capturing
cap = cv2.VideoCapture(0)
# Define the codec and create VideoWriter object
fourcc = cv2.VideoWriter_fourcc(*'XVID')
out = cv2.VideoWriter(currentPath + name_pose + '.avi', fourcc, 25.0, (640, 480))
while(cap.isOpened()):
ret, frame = cap.read()
if ret:
# write the frame
out.write(frame)
cv2.imshow('frame', frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
else:
break
# Release everything if job is finished
cap.release()
out.release()
cv2.destroyAllWindows()
vid = cv2.VideoCapture(currentPath + name_pose + '.avi')
# Check if the video
if (not vid.isOpened()):
print('Error opening video stream or file')
return
print('>> loading frozen model..')
detection_graph, sess = detector_utils.load_inference_graph()
sess = tf.Session(graph=detection_graph)
print('>> model loaded!')
_iter = 1
# Read until video is completed
while(vid.isOpened()):
# Capture frame-by-frame
ret, frame = vid.read()
if ret:
print(' Processing frame: ' + str(_iter))
# Resize and convert to RGB for NN to work with
frame = cv2.resize(frame, (320, 180), interpolation=cv2.INTER_AREA)
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
# Detect object
boxes, scores = detector_utils.detect_objects(frame, detection_graph, sess)
# get region of interest
res = detector_utils.get_box_image(1, 0.2, scores, boxes, 320, 180, frame)
# Save cropped image
if(res is not None):
cv2.imwrite(currentPath + currentExample + str(_iter) + '.png', cv2.cvtColor(res, cv2.COLOR_RGB2BGR))
_iter += 1
# Break the loop
else:
break
print(' Processed ' + str(_iter) + ' frames!')
vid.release()
if __name__ == '__main__':
main()