def load_models(self): self.detect_model = TinyYOLOv3_onecls(self.inp_dets, device=self.device) self.pose_model = SPPE_FastPose(self.pose_backbone, self.inp_pose[0], self.inp_pose[1], device=self.device) self.tracker = Tracker(30, n_init=3) self.action_model = TSSTG(device=self.device)
default='cpu', help='Device to run model on cpu or cuda.') args = par.parse_args() device = args.device # device= torch.args.device("cpu") # 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,
class Models: def __init__(self): self.inp_dets = 416 self.inp_pose = (256, 192) self.pose_backbone = 'resnet50' self.show_detected = True self.show_skeleton = True self.device = 'cuda' self.load_models() def load_models(self): self.detect_model = TinyYOLOv3_onecls(self.inp_dets, device=self.device) self.pose_model = SPPE_FastPose(self.pose_backbone, self.inp_pose[0], self.inp_pose[1], device=self.device) self.tracker = Tracker(30, n_init=3) self.action_model = TSSTG(device=self.device) def kpt2bbox(self, kpt, ex=20): return np.array((kpt[:, 0].min() - ex, kpt[:, 1].min() - ex, kpt[:, 0].max() + ex, kpt[:, 1].max() + ex)) def process_frame(self, frame): detected = self.detect_model.detect(frame, need_resize=False, expand_bb=10) self.tracker.predict() for track in self.tracker.tracks: det = torch.tensor([track.to_tlbr().tolist() + [1.0, 1.0, 0.0]], dtype=torch.float32) detected = torch.cat([detected, det], dim=0) if detected is not None else det detections = [] if detected is not None: poses = self.pose_model.predict(frame, detected[:, 0:4], detected[:, 4]) detections = [ Detection( self.kpt2bbox(ps['keypoints'].numpy()), np.concatenate( (ps['keypoints'].numpy(), ps['kp_score'].numpy()), axis=1), ps['kp_score'].mean().numpy()) for ps in poses ] if self.show_detected: for bb in detected[:, 0:5]: frame = cv2.rectangle(frame, (bb[0], bb[1]), (bb[2], bb[3]), (0, 0, 255), 1) self.tracker.update(detections) for i, track in enumerate(self.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) if len(track.keypoints_list) == 30: pts = np.array(track.keypoints_list, dtype=np.float32) out = self.action_model.predict(pts, frame.shape[:2]) action_name = self.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) track.actions = out if track.time_since_update == 0: if self.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_DUPLEX, 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) return frame
from PoseEstimateLoader import SPPE_FastPose from fn import vis_frame_fast save_path = '../../Data/Home_new-pose+score.csv' annot_file = '../../Data/Home_new.csv' # from create_dataset_1.py video_folder = '../Data/falldata/Home/Videos' annot_folder = '../Data/falldata/Home/Annotation_files' # bounding box annotation for each frame. # DETECTION MODEL. detector = TinyYOLOv3_onecls() # POSE MODEL. inp_h = 320 inp_w = 256 pose_estimator = SPPE_FastPose(inp_h, inp_w) # with score. columns = [ 'video', 'frame', 'Nose_x', 'Nose_y', 'Nose_s', 'LShoulder_x', 'LShoulder_y', 'LShoulder_s', 'RShoulder_x', 'RShoulder_y', 'RShoulder_s', 'LElbow_x', 'LElbow_y', 'LElbow_s', 'RElbow_x', 'RElbow_y', 'RElbow_s', 'LWrist_x', 'LWrist_y', 'LWrist_s', 'RWrist_x', 'RWrist_y', 'RWrist_s', 'LHip_x', 'LHip_y', 'LHip_s', 'RHip_x', 'RHip_y', 'RHip_s', 'LKnee_x', 'LKnee_y', 'LKnee_s', 'RKnee_x', 'RKnee_y', 'RKnee_s', 'LAnkle_x', 'LAnkle_y', 'LAnkle_s', 'RAnkle_x', 'RAnkle_y', 'RAnkle_s', 'label' ] def normalize_points_with_size(points_xy, width, height, flip=False): points_xy[:, 0] /= width