/
generate_tracks.py
65 lines (51 loc) · 2.87 KB
/
generate_tracks.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
import argparse
import logging
import os
from action_recognition.tracker import Tracker
from action_recognition.detector import TFOpenpose, CaffeOpenpose
def main(args):
if args.tf_openpose:
detector = TFOpenpose()
else:
detector = CaffeOpenpose(model_path=args.model_path)
videos = parse_path(args.video, args.out_directory, args.allowed_video_formats)
for video, out_dir in videos:
tracker = Tracker(detector=detector, out_dir=out_dir)
tracker.video(video, args.draw_frames)
def parse_path(video, out_directory, allowed_video_formats):
videos = []
if os.path.isfile(video):
videos.append((video, out_directory))
elif os.path.isdir(video):
for dirpath, dirnames, filenames in os.walk(video):
for file in filenames:
if any(file.endswith(format_) for format_ in allowed_video_formats):
out_dir = out_directory + "/".join(dirpath.split("/")[1:])
video_file = os.path.join(dirpath, file)
videos.append((video_file, out_dir))
return videos
if __name__ == '__main__':
parser = argparse.ArgumentParser(
description=('Generates tracks of people using OpenPose. '
'Each track is a [n_frames, n_keypoints, 3] numpy.ndarray which is predicted '
'as being a single person through several frames, making the final '
'outputted array of shape [n_tracks, n_frames, n_keypoints, 3], '
'where the values are (x, y, confidence), saved to a .npz file. '
'Note, however, that it will be an array with dtype=object, '
'since the n_frames per track will differ.'))
parser.add_argument('--video', type=str, default='media/video.avi',
help=('The video/folder to run tracking on. If folder, maintains the '
'hierarchy within that folder in the output.'))
parser.add_argument('--tf-openpose', action='store_true',
help='Use to make the program use the tensorflow implementation of OpenPose.')
parser.add_argument('--model-path', type=str, default='../openpose/models/',
help='The model path for the caffe implementation.')
parser.add_argument('--out-directory', type=str, default='output',
help='Root directory to where the annotated video is saved.')
parser.add_argument('--draw-frames', action='store_true',
help='Flag for if the frames with identified frames should be drawn or not.')
parser.add_argument('--allowed-video-formats', type=str, nargs='+', default=['.mp4', '.avi'],
help='Used for filtering of videos if the parameter video is a directory.')
args = parser.parse_args()
logging.basicConfig(level=logging.INFO)
main(args)