-
Notifications
You must be signed in to change notification settings - Fork 254
/
main.py
187 lines (153 loc) · 7.23 KB
/
main.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
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
import os
import cv2
import time
import torch
import argparse
import numpy as np
from Detection.Utils import ResizePadding
from CameraLoader import CamLoader, CamLoader_Q
from DetectorLoader import TinyYOLOv3_onecls
from PoseEstimateLoader import SPPE_FastPose
from fn import draw_single
from Track.Tracker import Detection, Tracker
from ActionsEstLoader import TSSTG
#source = '../Data/test_video/test7.mp4'
#source = '../Data/falldata/Home/Videos/video (2).avi' # hard detect
source = '../Data/falldata/Home/Videos/video (1).avi'
#source = 2
def preproc(image):
"""preprocess function for CameraLoader.
"""
image = resize_fn(image)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
return image
def kpt2bbox(kpt, ex=20):
"""Get bbox that hold on all of the keypoints (x,y)
kpt: array of shape `(N, 2)`,
ex: (int) expand bounding box,
"""
return np.array((kpt[:, 0].min() - ex, kpt[:, 1].min() - ex,
kpt[:, 0].max() + ex, kpt[:, 1].max() + ex))
if __name__ == '__main__':
par = argparse.ArgumentParser(description='Human Fall Detection Demo.')
par.add_argument('-C', '--camera', default=source, # required=True, # default=2,
help='Source of camera or video file path.')
par.add_argument('--detection_input_size', type=int, default=384,
help='Size of input in detection model in square must be divisible by 32 (int).')
par.add_argument('--pose_input_size', type=str, default='224x160',
help='Size of input in pose model must be divisible by 32 (h, w)')
par.add_argument('--pose_backbone', type=str, default='resnet50',
help='Backbone model for SPPE FastPose model.')
par.add_argument('--show_detected', default=False, action='store_true',
help='Show all bounding box from detection.')
par.add_argument('--show_skeleton', default=True, action='store_true',
help='Show skeleton pose.')
par.add_argument('--save_out', type=str, default='',
help='Save display to video file.')
par.add_argument('--device', type=str, default='cuda',
help='Device to run model on cpu or cuda.')
args = par.parse_args()
device = args.device
# DETECTION MODEL.
inp_dets = args.detection_input_size
detect_model = TinyYOLOv3_onecls(inp_dets, device=device)
# POSE MODEL.
inp_pose = args.pose_input_size.split('x')
inp_pose = (int(inp_pose[0]), int(inp_pose[1]))
pose_model = SPPE_FastPose(args.pose_backbone, inp_pose[0], inp_pose[1], device=device)
# Tracker.
max_age = 30
tracker = Tracker(max_age=max_age, n_init=3)
# Actions Estimate.
action_model = TSSTG()
resize_fn = ResizePadding(inp_dets, inp_dets)
cam_source = args.camera
if type(cam_source) is str and os.path.isfile(cam_source):
# Use loader thread with Q for video file.
cam = CamLoader_Q(cam_source, queue_size=1000, preprocess=preproc).start()
else:
# Use normal thread loader for webcam.
cam = CamLoader(int(cam_source) if cam_source.isdigit() else cam_source,
preprocess=preproc).start()
#frame_size = cam.frame_size
#scf = torch.min(inp_size / torch.FloatTensor([frame_size]), 1)[0]
outvid = False
if args.save_out != '':
outvid = True
codec = cv2.VideoWriter_fourcc(*'MJPG')
writer = cv2.VideoWriter(args.save_out, codec, 30, (inp_dets * 2, inp_dets * 2))
fps_time = 0
f = 0
while cam.grabbed():
f += 1
frame = cam.getitem()
image = frame.copy()
# Detect humans bbox in the frame with detector model.
detected = detect_model.detect(frame, need_resize=False, expand_bb=10)
# Predict each tracks bbox of current frame from previous frames information with Kalman filter.
tracker.predict()
# Merge two source of predicted bbox together.
for track in tracker.tracks:
det = torch.tensor([track.to_tlbr().tolist() + [0.5, 1.0, 0.0]], dtype=torch.float32)
detected = torch.cat([detected, det], dim=0) if detected is not None else det
detections = [] # List of Detections object for tracking.
if detected is not None:
#detected = non_max_suppression(detected[None, :], 0.45, 0.2)[0]
# Predict skeleton pose of each bboxs.
poses = pose_model.predict(frame, detected[:, 0:4], detected[:, 4])
# Create Detections object.
detections = [Detection(kpt2bbox(ps['keypoints'].numpy()),
np.concatenate((ps['keypoints'].numpy(),
ps['kp_score'].numpy()), axis=1),
ps['kp_score'].mean().numpy()) for ps in poses]
# VISUALIZE.
if args.show_detected:
for bb in detected[:, 0:5]:
frame = cv2.rectangle(frame, (bb[0], bb[1]), (bb[2], bb[3]), (0, 0, 255), 1)
# Update tracks by matching each track information of current and previous frame or
# create a new track if no matched.
tracker.update(detections)
# Predict Actions of each track.
for i, track in enumerate(tracker.tracks):
if not track.is_confirmed():
continue
track_id = track.track_id
bbox = track.to_tlbr().astype(int)
center = track.get_center().astype(int)
action = 'pending..'
clr = (0, 255, 0)
# Use 30 frames time-steps to prediction.
if len(track.keypoints_list) == 30:
pts = np.array(track.keypoints_list, dtype=np.float32)
out = action_model.predict(pts, frame.shape[:2])
action_name = action_model.class_names[out[0].argmax()]
action = '{}: {:.2f}%'.format(action_name, out[0].max() * 100)
if action_name == 'Fall Down':
clr = (255, 0, 0)
elif action_name == 'Lying Down':
clr = (255, 200, 0)
# VISUALIZE.
if track.time_since_update == 0:
if args.show_skeleton:
frame = draw_single(frame, track.keypoints_list[-1])
frame = cv2.rectangle(frame, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 1)
frame = cv2.putText(frame, str(track_id), (center[0], center[1]), cv2.FONT_HERSHEY_COMPLEX,
0.4, (255, 0, 0), 2)
frame = cv2.putText(frame, action, (bbox[0] + 5, bbox[1] + 15), cv2.FONT_HERSHEY_COMPLEX,
0.4, clr, 1)
# Show Frame.
frame = cv2.resize(frame, (0, 0), fx=2., fy=2.)
frame = cv2.putText(frame, '%d, FPS: %f' % (f, 1.0 / (time.time() - fps_time)),
(10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
frame = frame[:, :, ::-1]
fps_time = time.time()
if outvid:
writer.write(frame)
cv2.imshow('frame', frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
# Clear resource.
cam.stop()
if outvid:
writer.release()
cv2.destroyAllWindows()