diff --git a/deploy/python/det_keypoint_unite_infer.py b/deploy/python/det_keypoint_unite_infer.py index 951864b2a1a86b111f9bb829617133a3b4ea5f98..62d3b82f4512488a10319ef7cc3d145b316d7524 100644 --- a/deploy/python/det_keypoint_unite_infer.py +++ b/deploy/python/det_keypoint_unite_infer.py @@ -143,8 +143,8 @@ def topdown_unite_predict_video(detector, writer = cv2.VideoWriter(out_path, fourcc, fps, (width, height)) index = 0 store_res = [] - previous_keypoints = None - keypoint_smoothing = KeypointSmoothing(width, height, filter_type=FLAGS.filter_type, alpha=0.8, beta=1) + keypoint_smoothing = KeypointSmoothing( + width, height, filter_type=FLAGS.filter_type, beta=0.05) while (1): ret, frame = capture.read() @@ -164,11 +164,11 @@ def topdown_unite_predict_video(detector, keypoint_res = predict_with_given_det( frame2, results, topdown_keypoint_detector, keypoint_batch_size, FLAGS.run_benchmark) - + if FLAGS.smooth: current_keypoints = np.array(keypoint_res['keypoint'][0][0]) - smooth_keypoints = keypoint_smoothing.smooth_process(previous_keypoints, current_keypoints) - previous_keypoints = smooth_keypoints + smooth_keypoints = keypoint_smoothing.smooth_process( + current_keypoints) keypoint_res['keypoint'][0][0] = smooth_keypoints.tolist() @@ -205,13 +205,24 @@ def topdown_unite_predict_video(detector, class KeypointSmoothing(object): # The following code are modified from: - # https://github.com/610265158/Peppa_Pig_Face_Engine/blob/7bb1066ad3fbb12697924ba7f9287bf198c15232/lib/core/LK/lk.py - - def __init__(self, width, height, filter_type, alpha=0.5, fc_d=1, fc_min=1, beta=0): + # https://github.com/jaantollander/OneEuroFilter + + def __init__(self, + width, + height, + filter_type, + alpha=0.5, + fc_d=0.1, + fc_min=0.1, + beta=0.1, + thres_mult=0.2): super(KeypointSmoothing, self).__init__() self.image_width = width self.image_height = height - self.threshold = [0.005, 0.005, 0.005, 0.005, 0.005, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01] + self.threshold = np.array([ + 0.005, 0.005, 0.005, 0.005, 0.005, 0.01, 0.01, 0.01, 0.01, 0.01, + 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01 + ]) * thres_mult self.filter_type = filter_type self.alpha = alpha self.dx_prev_hat = None @@ -219,7 +230,7 @@ class KeypointSmoothing(object): self.fc_d = fc_d self.fc_min = fc_min self.beta = beta - + if self.filter_type == 'one_euro': self.smooth_func = self.one_euro_filter elif self.filter_type == 'ema': @@ -227,51 +238,50 @@ class KeypointSmoothing(object): else: raise ValueError('filter type must be one_euro or ema') - def smooth_process(self, previous_keypoints, current_keypoints): - if previous_keypoints is None: - previous_keypoints = current_keypoints - result = current_keypoints + def smooth_process(self, current_keypoints): + if self.x_prev_hat is None: + self.x_prev_hat = current_keypoints[:, :2] + self.dx_prev_hat = np.zeros(current_keypoints[:, :2].shape) + return current_keypoints else: - result = [] + result = current_keypoints num_keypoints = len(current_keypoints) for i in range(num_keypoints): - result.append(self.smooth(previous_keypoints[i], current_keypoints[i], self.threshold[i])) - return np.array(result) - - - def smooth(self, previous_keypoint, current_keypoint, threshold): - distance = np.sqrt(np.square((current_keypoint[0] - previous_keypoint[0]) / self.image_width) + np.square((current_keypoint[1] - previous_keypoint[1]) / self.image_height)) + result[i, :2] = self.smooth(current_keypoints[i, :2], + self.threshold[i], i) + return result + + def smooth(self, current_keypoint, threshold, index): + distance = np.sqrt( + np.square((current_keypoint[0] - self.x_prev_hat[index][0]) / + self.image_width) + np.square((current_keypoint[ + 1] - self.x_prev_hat[index][1]) / self.image_height)) if distance < threshold: - result = previous_keypoint + result = self.x_prev_hat[index] else: - result = self.smooth_func(previous_keypoint, current_keypoint) + result = self.smooth_func(current_keypoint, self.x_prev_hat[index], + index) return result - - def one_euro_filter(self, x_prev, x_cur): + def one_euro_filter(self, x_cur, x_pre, index): te = 1 self.alpha = self.smoothing_factor(te, self.fc_d) - if self.x_prev_hat is None: - self.x_prev_hat = x_prev - dx_cur = (x_cur - self.x_prev_hat) / te - if self.dx_prev_hat is None: - self.dx_prev_hat = 0 - dx_cur_hat = self.exponential_smoothing(self.dx_prev_hat, dx_cur) - + dx_cur = (x_cur - x_pre) / te + dx_cur_hat = self.exponential_smoothing(dx_cur, self.dx_prev_hat[index]) + fc = self.fc_min + self.beta * np.abs(dx_cur_hat) self.alpha = self.smoothing_factor(te, fc) - x_cur_hat = self.exponential_smoothing(self.x_prev_hat, x_cur) - self.dx_prev_hat = dx_cur_hat - self.x_prev_hat = x_cur_hat + x_cur_hat = self.exponential_smoothing(x_cur, x_pre) + self.dx_prev_hat[index] = dx_cur_hat + self.x_prev_hat[index] = x_cur_hat return x_cur_hat - def smoothing_factor(self, te, fc): r = 2 * math.pi * fc * te return r / (r + 1) - - def exponential_smoothing(self, x_prev, x_cur): - return self.alpha * x_cur + (1 - self.alpha) * x_prev + + def exponential_smoothing(self, x_cur, x_pre, index=0): + return self.alpha * x_cur + (1 - self.alpha) * x_pre def main():