-
Notifications
You must be signed in to change notification settings - Fork 1
/
get_data.py
83 lines (73 loc) · 2.79 KB
/
get_data.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
import cv2
import sys
import os
import shutil
import subprocess
SEND_IMAGES = False
RENDERED_DATA = True
TARGET_FPS = 10
SERVER = '[email protected]'
def get_data_alphapose(path, results_dir_name, is_vid, get_rendered):
'''
Run AlphaPose remotely on data and get back the joint data
path: path to video file or directory of images
results_dir_name: name of directory where results are to be stored
is_vid: True if video file, False if directory of images
get_rendered: if rendered data is to be generated
'''
# SCP data to the server
subprocess.run(['scp', '-r', path, SERVER+':~/AlphaPose/'])
# Run alphapose
end_path = os.path.split(path)[-1]
alph_cmd = 'echo ; source ~/miniconda3/bin/activate alphapose ; cd ~/AlphaPose ;'
alph_cmd += 'python scripts/demo_inference.py --cfg configs/coco/resnet/256x192_res50_lr1e-3_1x.yaml --checkpoint pretrained_models/fast_res50_256x192.pth'
if is_vid:
alph_cmd += f' --video {end_path} --outdir {results_dir_name}'
if get_rendered:
alph_cmd += ' --save_video'
else:
alph_cmd += f' --indir {end_path} --outdir {results_dir_name}'
if get_rendered:
alph_cmd += ' --save_img'
subprocess.run(['ssh', SERVER] + alph_cmd.split())
# SCP results back
subprocess.run(['scp', '-r', SERVER+':~/AlphaPose/'+results_dir_name, '.'])
# Delete data on server
subprocess.run(['ssh', SERVER, 'rm', '-r', '~/AlphaPose/' + end_path, '~/AlphaPose/' + results_dir_name])
def get_images_from_video(vid_path, target_fps):
cap = cv2.VideoCapture(vid_path)
original_fps = cap.get(cv2.CAP_PROP_FPS)
frames_skip = 0
if original_fps > target_fps:
frames_skip = original_fps // target_fps
print('original fps:', original_fps)
print('frames skip:', frames_skip)
if not os.path.exists('./frames'):
os.mkdir('frames')
i, retained = 0, 0
while True:
ret, frame = cap.read()
if not ret:
break
if i % int(frames_skip) == 0:
cv2.imwrite('frames/'+str(i)+'.jpg', frame)
retained += 1
i += 1
print(f'{retained} out of {i} frames retained')
cap.release()
cv2.destroyAllWindows()
def get_data(vid_path):
vid_name = vid_path.split('/')[-1].split('.')[0]
if SEND_IMAGES:
get_images_from_video(vid_path, TARGET_FPS)
get_data_alphapose('frames', 'results-'+vid_name, False, RENDERED_DATA)
shutil.rmtree('frames')
else:
get_data_alphapose(vid_path, 'results-'+vid_name, True, RENDERED_DATA)
if __name__ == '__main__':
if len(sys.argv) != 2:
print('Usage: python3 get_data.py [path to file]')
print('e.g. python3 get_data.py awesome.mp4')
exit()
vid_path = sys.argv[1]
get_data(vid_path)