diff --git a/estimate_head_pose.py b/estimate_head_pose.py index 65a6547..bcc3f2a 100644 --- a/estimate_head_pose.py +++ b/estimate_head_pose.py @@ -22,6 +22,8 @@ def main(): # Introduce point stabilizer. stabilizers = [point_stabilizer.Stabilizer( cov_process=0.01, cov_measure=0.1) for _ in range(68)] + target_latest_state = 0 # 1: moving; 0: still. + target_current_state = 1 # Introduce an optical flow tracker to help to decide how kalman filter # should be configured. Alos keep one frame for optical flow tracker. @@ -42,7 +44,13 @@ def main(): # Optical flow tracker should work before kalman filter. frame_opt_flw = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) if len(tracker.tracks) > 0: + if tracker.get_average_track_length() > 2: + target_current_state = 1 + else: + target_current_state = 0 + tracker.update_tracks(frame_prev, frame_opt_flw) + frame_prev = frame_opt_flw # CNN benchmark. @@ -78,16 +86,16 @@ def main(): # in the same state. An optical flow tracker also proved to be helpfull to # tell which state the target is currently in. stabile_marks = [] - # TODO: use optical flow to determine how to set stabilizer. - if True: - cov_process = 0.0001 - cov_measure = 0.1 - for stabilizer in stabilizers: - stabilizer.set_q_r(cov_process=cov_process, - cov_measure=cov_measure) - else: - cov_process = 0.1 - cov_measure = 0.001 + if target_current_state != target_latest_state: + if target_current_state == 1: + cov_process = 0.1 + cov_measure = 0.001 + else: + cov_process = 0.001 + cov_measure = 0.1 + + target_latest_state = target_current_state + for stabilizer in stabilizers: stabilizer.set_q_r(cov_process=cov_process, cov_measure=cov_measure) diff --git a/optical_flow_tracker.py b/optical_flow_tracker.py index 25fa151..41f9d28 100644 --- a/optical_flow_tracker.py +++ b/optical_flow_tracker.py @@ -3,6 +3,8 @@ Lucas-Kanade sparse optical flow tracker. Uses goodFeaturesToTrack for track initialization and back-tracking for match verification between frames. ''' +from math import sqrt + import numpy as np import cv2 @@ -12,7 +14,7 @@ class Tracker: """Lucas-Kanade sparse optical flow tracker""" def __init__(self): - self.track_len = 10 + self.track_len = 5 self.tracks = [] self.lk_params = dict(winSize=(15, 15), maxLevel=2, @@ -74,6 +76,19 @@ class Tracker: for x, y in np.float32(feature_points).reshape(-1, 2): self.tracks.append([(x, y)]) + def get_average_track_length(self): + """Get the average track length""" + length = 0 + tracks = np.array(self.tracks) + def distance(track): + """Get distance between the first and last point.""" + delta_x = abs(track[-1][0] - track[0][0]) + delta_y = abs(track[-1][1] - track[0][1]) + return sqrt(delta_x*delta_x + delta_y*delta_y) + for track in tracks: + length += distance(track) + return length / len(tracks) + def draw_track(self, image): """Draw track lines on image.""" cv2.polylines(image, [np.int32(track)