-
Notifications
You must be signed in to change notification settings - Fork 237
/
demo.py
128 lines (109 loc) · 4.17 KB
/
demo.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
import argparse
import sys
import numpy as np
import norfair
from norfair import Detection, Tracker, Video
from norfair.distances import create_keypoints_voting_distance
# Import openpose
openpose_install_path = (
"/openpose" # Insert the path to your openpose instalation folder here
)
try:
sys.path.append(openpose_install_path + "/build/python")
from openpose import pyopenpose as op
except ImportError as e:
print(
"Error: OpenPose library could not be found. Did you enable `BUILD_PYTHON` in CMake and have this Python script in the right folder?"
)
raise e
# Define constants
DETECTION_THRESHOLD = 0.1
DISTANCE_THRESHOLD = 0.4
INITIALIZATION_DELAY = 4
HIT_COUNTER_MAX = 30
POINTWISE_HIT_COUNTER_MAX = 10
# Wrapper implementation for OpenPose detector
class OpenposeDetector:
def __init__(self, num_gpu_start=None):
# Set OpenPose flags
config = {}
config["model_folder"] = openpose_install_path + "/models/"
config["model_pose"] = "BODY_25"
config["logging_level"] = 3
config["output_resolution"] = "-1x-1"
config["net_resolution"] = "-1x768"
config["num_gpu"] = 1
config["alpha_pose"] = 0.6
config["render_threshold"] = 0.05
config["scale_number"] = 1
config["scale_gap"] = 0.3
config["disable_blending"] = False
# If GPU version is built, and multiple GPUs are available,
# you can change the ID using the num_gpu_start parameter
if num_gpu_start is not None:
config["num_gpu_start"] = num_gpu_start
# Starting OpenPose
self.detector = op.WrapperPython()
self.detector.configure(config)
self.detector.start()
def __call__(self, image):
return self.detector.emplaceAndPop(image)
if __name__ == "__main__":
# CLI configuration
parser = argparse.ArgumentParser(description="Track human poses in a video.")
parser.add_argument("files", type=str, nargs="+", help="Video files to process")
parser.add_argument(
"--skip-frame", dest="skip_frame", type=int, default=1, help="Frame skip period"
)
parser.add_argument(
"--select-gpu",
dest="select_gpu",
help="Number of the gpu that you want to use",
default=None,
type=int,
)
args = parser.parse_args()
# Process Videos
detector = OpenposeDetector(args.select_gpu)
datum = op.Datum()
for input_path in args.files:
print(f"Video: {input_path}")
video = Video(input_path=input_path)
KEYPOINT_DIST_THRESHOLD = video.input_height / 40
tracker = Tracker(
distance_function=create_keypoints_voting_distance(
keypoint_distance_threshold=KEYPOINT_DIST_THRESHOLD,
detection_threshold=DETECTION_THRESHOLD,
),
distance_threshold=DISTANCE_THRESHOLD,
detection_threshold=DETECTION_THRESHOLD,
initialization_delay=INITIALIZATION_DELAY,
hit_counter_max=HIT_COUNTER_MAX,
pointwise_hit_counter_max=POINTWISE_HIT_COUNTER_MAX,
)
for i, frame in enumerate(video):
if i % args.skip_frame == 0:
datum.cvInputData = frame
detector(op.VectorDatum([datum]))
detected_poses = datum.poseKeypoints
if detected_poses is None:
tracked_objects = tracker.update(period=args.skip_frame)
continue
detections = (
[]
if not detected_poses.any()
else [
Detection(p, scores=s)
for (p, s) in zip(
detected_poses[:, :, :2], detected_poses[:, :, 2]
)
]
)
tracked_objects = tracker.update(
detections=detections, period=args.skip_frame
)
norfair.draw_points(frame, detections)
else:
tracked_objects = tracker.update(period=args.skip_frame)
norfair.draw_tracked_objects(frame, tracked_objects)
video.write(frame)