| 37 | blockSize = 7 ) |
| 38 | |
| 39 | class App: |
| 40 | def __init__(self, video_src): |
| 41 | self.track_len = 10 |
| 42 | self.detect_interval = 5 |
| 43 | self.tracks = [] |
| 44 | self.cam = video.create_capture(video_src) |
| 45 | self.frame_idx = 0 |
| 46 | |
| 47 | def run(self): |
| 48 | while True: |
| 49 | ret, frame = self.cam.read() |
| 50 | frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) |
| 51 | vis = frame.copy() |
| 52 | |
| 53 | if len(self.tracks) > 0: |
| 54 | img0, img1 = self.prev_gray, frame_gray |
| 55 | p0 = np.float32([tr[-1] for tr in self.tracks]).reshape(-1, 1, 2) |
| 56 | p1, st, err = cv2.calcOpticalFlowPyrLK(img0, img1, p0, None, **lk_params) |
| 57 | p0r, st, err = cv2.calcOpticalFlowPyrLK(img1, img0, p1, None, **lk_params) |
| 58 | d = abs(p0-p0r).reshape(-1, 2).max(-1) |
| 59 | good = d < 1 |
| 60 | new_tracks = [] |
| 61 | for tr, (x, y), good_flag in zip(self.tracks, p1.reshape(-1, 2), good): |
| 62 | if not good_flag: |
| 63 | continue |
| 64 | tr.append((x, y)) |
| 65 | if len(tr) > self.track_len: |
| 66 | del tr[0] |
| 67 | new_tracks.append(tr) |
| 68 | cv2.circle(vis, (x, y), 2, (0, 255, 0), -1) |
| 69 | self.tracks = new_tracks |
| 70 | cv2.polylines(vis, [np.int32(tr) for tr in self.tracks], False, (0, 255, 0)) |
| 71 | draw_str(vis, (20, 20), 'track count: %d' % len(self.tracks)) |
| 72 | |
| 73 | if self.frame_idx % self.detect_interval == 0: |
| 74 | mask = np.zeros_like(frame_gray) |
| 75 | mask[:] = 255 |
| 76 | for x, y in [np.int32(tr[-1]) for tr in self.tracks]: |
| 77 | cv2.circle(mask, (x, y), 5, 0, -1) |
| 78 | p = cv2.goodFeaturesToTrack(frame_gray, mask = mask, **feature_params) |
| 79 | if p is not None: |
| 80 | for x, y in np.float32(p).reshape(-1, 2): |
| 81 | self.tracks.append([(x, y)]) |
| 82 | |
| 83 | |
| 84 | self.frame_idx += 1 |
| 85 | self.prev_gray = frame_gray |
| 86 | cv2.imshow('lk_track', vis) |
| 87 | |
| 88 | ch = cv2.waitKey(1) |
| 89 | if ch == 27: |
| 90 | break |
| 91 | |
| 92 | def main(): |
| 93 | import sys |