optical_flow_tracker.py

'''
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


class Tracker:
    """Lucas-Kanade sparse optical flow tracker"""

    def __init__(self):
        self.track_len = 5
        self.tracks = []
        self.lk_params = dict(winSize=(15, 15),
                              maxLevel=2,
                              criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03))
        self.feature_params = dict(maxCorners=500,
                                   qualityLevel=0.3,
                                   minDistance=7,
                                   blockSize=7)

    def update_tracks(self, img_old, img_new):
        """Update tracks."""
        # Get old points, using the latest one.
        points_old = np.float32([track[-1]
                                 for track in self.tracks]).reshape(-1, 1, 2)

        # Get new points from old points.
        points_new, _st, _err = cv2.calcOpticalFlowPyrLK(
            img_old, img_new, points_old, None, **self.lk_params)

        # Get inferred old points from new points.
        points_old_inferred, _st, _err = cv2.calcOpticalFlowPyrLK(
            img_new, img_old, points_new, None, **self.lk_params)

        # Compare between old points and inferred old points
        error_term = abs(
            points_old - points_old_inferred).reshape(-1, 2).max(-1)
        point_valid = error_term < 1

        new_tracks = []
        for track, (x, y), good_flag in zip(self.tracks, points_new.reshape(-1, 2), point_valid):
            # Track is good?
            if not good_flag:
                continue

            # New point is good, add to track.
            track.append((x, y))

            # Need to drop first old point?
            if len(track) > self.track_len:
                del track[0]

            # Track updated, add to track groups.
            new_tracks.append(track)

        # New track groups got, do update.
        self.tracks = new_tracks

    def get_new_tracks(self, frame, roi):
        """Get new tracks every detect_interval frames."""
        # Using mask to determine where to look for feature points.
        mask = np.zeros_like(frame)
        mask[roi[0]:roi[1], roi[2]:roi[3]] = 255

        # Get good feature points.
        feature_points = cv2.goodFeaturesToTrack(
            frame, mask=mask, **self.feature_params)

        if feature_points is not None:
            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)
                              for track in self.tracks], False, (0, 255, 0))


def main():
    """Test code"""
    import sys
    try:
        video_src = sys.argv[1]
    except:
        video_src = 0

    tracker = Tracker()

    cam = cv2.VideoCapture(video_src)
    detect_interval = 5
    frame_idx = 0

    prev_gray = cam.read()
    while True:
        _ret, frame = cam.read()
        frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        # Update tracks.
        if len(tracker.tracks) > 0:
            tracker.update_tracks(prev_gray, frame_gray)

        # Get new tracks every detect_interval frames.
        target_box = [100, 400, 100, 400]
        if frame_idx % detect_interval == 0:
            tracker.get_new_tracks(frame_gray, target_box)

        # Draw tracks
        tracker.draw_track(frame)

        frame_idx += 1
        prev_gray = frame_gray
        cv2.imshow('lk_track', frame)
        ch = cv2.waitKey(1)
        if ch == 27:
            break


if __name__ == '__main__':
    main()