diff --git a/samples/python2/common.py b/samples/python2/common.py index 6b62f4304b..883aa9ae23 100644 --- a/samples/python2/common.py +++ b/samples/python2/common.py @@ -204,3 +204,9 @@ def getsize(img): def mdot(*args): return reduce(np.dot, args) + +def draw_keypoints(vis, keypoints, color = (0, 255, 255)): + for kp in keypoints: + x, y = kp.pt + cv2.circle(vis, (int(x), int(y)), 2, color) + diff --git a/samples/python2/feature_homography.py b/samples/python2/feature_homography.py index 182b5698f5..eec36d74e9 100644 --- a/samples/python2/feature_homography.py +++ b/samples/python2/feature_homography.py @@ -19,122 +19,23 @@ import numpy as np import cv2 import video import common -from collections import namedtuple -from common import getsize, Bunch - - -FLANN_INDEX_KDTREE = 1 -FLANN_INDEX_LSH = 6 -flann_params= dict(algorithm = FLANN_INDEX_LSH, - table_number = 6, # 12 - key_size = 12, # 20 - multi_probe_level = 1) #2 - -MIN_MATCH_COUNT = 10 - - -ar_verts = np.float32([[0, 0, 0], [0, 1, 0], [1, 1, 0], [1, 0, 0], - [0, 0, 1], [0, 1, 1], [1, 1, 1], [1, 0, 1], - [0, 0.5, 2], [1, 0.5, 2]]) -ar_edges = [(0, 1), (1, 2), (2, 3), (3, 0), - (4, 5), (5, 6), (6, 7), (7, 4), - (0, 4), (1, 5), (2, 6), (3, 7), - (4, 8), (5, 8), (6, 9), (7, 9), (8, 9)] +from common import getsize, draw_keypoints +from plane_tracker import PlaneTracker - -def draw_keypoints(vis, keypoints, color = (0, 255, 255)): - for kp in keypoints: - x, y = kp.pt - cv2.circle(vis, (int(x), int(y)), 2, color) - class App: def __init__(self, src): self.cap = video.create_capture(src) self.frame = None self.paused = False - self.ref_frames = [] - - self.detector = cv2.ORB( nfeatures = 1000 ) - self.matcher = cv2.FlannBasedMatcher(flann_params, {}) # bug : need to pass empty dict (#1329) + self.tracker = PlaneTracker() cv2.namedWindow('plane') self.rect_sel = common.RectSelector('plane', self.on_rect) - - - def match_frames(self): - if len(self.frame_desc) < MIN_MATCH_COUNT or len(self.frame_desc) < MIN_MATCH_COUNT: - return - - matches = self.matcher.knnMatch(self.frame_desc, k = 2) - matches = [m[0] for m in matches if len(m) == 2 and m[0].distance < m[1].distance * 0.75] - if len(matches) < MIN_MATCH_COUNT: - return - img_ids = [m.imgIdx for m in matches] - match_counts = np.bincount(img_ids, minlength=len(self.ref_frames)) - bast_id = match_counts.argmax() - if match_counts[bast_id] < MIN_MATCH_COUNT: - return - ref_frame = self.ref_frames[bast_id] - matches = [m for m in matches if m.imgIdx == bast_id] - p0 = [ref_frame.points[m.trainIdx].pt for m in matches] - p1 = [self.frame_points[m.queryIdx].pt for m in matches] - p0, p1 = np.float32((p0, p1)) - if len(p0) < MIN_MATCH_COUNT: - return - - H, status = cv2.findHomography(p0, p1, cv2.RANSAC, 4.0) - status = status.ravel() != 0 - if status.sum() < MIN_MATCH_COUNT: - return - p0, p1 = p0[status], p1[status] - return ref_frame, p0, p1, H - - - def on_frame(self, vis): - match = self.match_frames() - if match is None: - return - - w, h = getsize(self.frame) - ref_frame, p0, p1, H = match - vis[:h,w:] = ref_frame.frame - draw_keypoints(vis[:,w:], ref_frame.points) - x0, y0, x1, y1 = ref_frame.rect - cv2.rectangle(vis, (x0+w, y0), (x1+w, y1), (0, 255, 0), 2) - corners0 = np.float32([[x0, y0], [x1, y0], [x1, y1], [x0, y1]]) - img_corners = cv2.perspectiveTransform(corners0.reshape(1, -1, 2), H) - cv2.polylines(vis, [np.int32(img_corners)], True, (255, 255, 255), 2) - - for (x0, y0), (x1, y1) in zip(np.int32(p0), np.int32(p1)): - cv2.line(vis, (x0+w, y0), (x1, y1), (0, 255, 0)) - - ''' - corners3d = np.hstack([corners0, np.zeros((4, 1), np.float32)]) - fx = 0.9 - K = np.float64([[fx*w, 0, 0.5*(w-1)], - [0, fx*w, 0.5*(h-1)], - [0.0,0.0, 1.0]]) - dist_coef = np.zeros(4) - ret, rvec, tvec = cv2.solvePnP(corners3d, img_corners, K, dist_coef) - verts = ar_verts * [(x1-x0), (y1-y0), -(x1-x0)*0.3] + (x0, y0, 0) - verts = cv2.projectPoints(verts, rvec, tvec, K, dist_coef)[0].reshape(-1, 2) - for i, j in ar_edges: - (x0, y0), (x1, y1) = verts[i], verts[j] - cv2.line(vis, (int(x0), int(y0)), (int(x1), int(y1)), (255, 255, 0), 2) - ''' + def on_rect(self, rect): - x0, y0, x1, y1 = rect - points, descs = [], [] - for kp, desc in zip(self.frame_points, self.frame_desc): - x, y = kp.pt - if x0 <= x <= x1 and y0 <= y <= y1: - points.append(kp) - descs.append(desc) - descs = np.uint8(descs) - frame_data = Bunch(frame = self.frame, rect=rect, points = points, descs=descs) - self.ref_frames.append(frame_data) - self.matcher.add([descs]) + self.tracker.clear() + self.tracker.add_target(self.frame, rect) def run(self): while True: @@ -143,19 +44,27 @@ class App: ret, frame = self.cap.read() if not ret: break - self.frame = np.fliplr(frame).copy() - self.frame_points, self.frame_desc = self.detector.detectAndCompute(self.frame, None) - if self.frame_desc is None: # detectAndCompute returns descs=None if not keypoints found - self.frame_desc = [] + self.frame = np.frame.copy() w, h = getsize(self.frame) vis = np.zeros((h, w*2, 3), np.uint8) vis[:h,:w] = self.frame - draw_keypoints(vis, self.frame_points) + if len(self.tracker.targets) > 0: + target = self.tracker.targets[0] + vis[:,w:] = target.image + draw_keypoints(vis[:,w:], target.keypoints) + x0, y0, x1, y1 = target.rect + cv2.rectangle(vis, (x0+w, y0), (x1+w, y1), (0, 255, 0), 2) if playing: - self.on_frame(vis) - + tracked = self.tracker.track(self.frame) + if len(tracked) > 0: + tracked = tracked[0] + cv2.polylines(vis, [np.int32(tracked.quad)], True, (255, 255, 255), 2) + for (x0, y0), (x1, y1) in zip(np.int32(tracked.p0), np.int32(tracked.p1)): + cv2.line(vis, (x0+w, y0), (x1, y1), (0, 255, 0)) + draw_keypoints(vis, self.tracker.frame_points) + self.rect_sel.draw(vis) cv2.imshow('plane', vis) ch = cv2.waitKey(1) @@ -164,6 +73,7 @@ class App: if ch == 27: break + if __name__ == '__main__': print __doc__ diff --git a/samples/python2/plane_ar.py b/samples/python2/plane_ar.py new file mode 100644 index 0000000000..27667270fc --- /dev/null +++ b/samples/python2/plane_ar.py @@ -0,0 +1,81 @@ +import numpy as np +import cv2 +import video +import common +from plane_tracker import PlaneTracker + + +ar_verts = np.float32([[0, 0, 0], [0, 1, 0], [1, 1, 0], [1, 0, 0], + [0, 0, 1], [0, 1, 1], [1, 1, 1], [1, 0, 1], + [0, 0.5, 2], [1, 0.5, 2]]) +ar_edges = [(0, 1), (1, 2), (2, 3), (3, 0), + (4, 5), (5, 6), (6, 7), (7, 4), + (0, 4), (1, 5), (2, 6), (3, 7), + (4, 8), (5, 8), (6, 9), (7, 9), (8, 9)] + +class App: + def __init__(self, src): + self.cap = video.create_capture(src) + self.frame = None + self.paused = False + self.tracker = PlaneTracker() + + cv2.namedWindow('plane') + cv2.createTrackbar('focal', 'plane', 25, 50, common.nothing) + self.rect_sel = common.RectSelector('plane', self.on_rect) + + def on_rect(self, rect): + self.tracker.add_target(self.frame, rect) + + def run(self): + while True: + playing = not self.paused and not self.rect_sel.dragging + if playing or self.frame is None: + ret, frame = self.cap.read() + if not ret: + break + self.frame = frame.copy() + + vis = self.frame.copy() + if playing: + tracked = self.tracker.track(self.frame) + for tr in tracked: + cv2.polylines(vis, [np.int32(tr.quad)], True, (255, 255, 255), 2) + for (x, y) in np.int32(tr.p1): + cv2.circle(vis, (x, y), 2, (255, 255, 255)) + self.draw_overlay(vis, tr) + + self.rect_sel.draw(vis) + cv2.imshow('plane', vis) + ch = cv2.waitKey(1) + if ch == ord(' '): + self.paused = not self.paused + if ch == ord('c'): + self.tracker.clear() + if ch == 27: + break + + def draw_overlay(self, vis, tracked): + x0, y0, x1, y1 = tracked.target.rect + quad_3d = np.float32([[x0, y0, 0], [x1, y0, 0], [x1, y1, 0], [x0, y1, 0]]) + fx = 0.5 + cv2.getTrackbarPos('focal', 'plane') / 50.0 + h, w = vis.shape[:2] + K = np.float64([[fx*w, 0, 0.5*(w-1)], + [0, fx*w, 0.5*(h-1)], + [0.0,0.0, 1.0]]) + dist_coef = np.zeros(4) + ret, rvec, tvec = cv2.solvePnP(quad_3d, tracked.quad, K, dist_coef) + verts = ar_verts * [(x1-x0), (y1-y0), -(x1-x0)*0.3] + (x0, y0, 0) + verts = cv2.projectPoints(verts, rvec, tvec, K, dist_coef)[0].reshape(-1, 2) + for i, j in ar_edges: + (x0, y0), (x1, y1) = verts[i], verts[j] + cv2.line(vis, (int(x0), int(y0)), (int(x1), int(y1)), (255, 255, 0), 2) + + +if __name__ == '__main__': + print __doc__ + + import sys + try: video_src = sys.argv[1] + except: video_src = 0 + App(video_src).run() diff --git a/samples/python2/plane_tracker.py b/samples/python2/plane_tracker.py new file mode 100644 index 0000000000..336654e9bf --- /dev/null +++ b/samples/python2/plane_tracker.py @@ -0,0 +1,136 @@ +import numpy as np +import cv2 +from collections import namedtuple +import video +import common + + +FLANN_INDEX_KDTREE = 1 +FLANN_INDEX_LSH = 6 +flann_params= dict(algorithm = FLANN_INDEX_LSH, + table_number = 6, # 12 + key_size = 12, # 20 + multi_probe_level = 1) #2 + +MIN_MATCH_COUNT = 10 + + +PlanarTarget = namedtuple('PlaneTarget', 'image, rect, keypoints, descrs, data') +TrackedTarget = namedtuple('TrackedTarget', 'target, p0, p1, H, quad') + +class PlaneTracker: + def __init__(self): + self.detector = cv2.ORB( nfeatures = 1000 ) + self.matcher = cv2.FlannBasedMatcher(flann_params, {}) # bug : need to pass empty dict (#1329) + self.targets = [] + + def add_target(self, image, rect, data=None): + '''Add a new tracking target.''' + x0, y0, x1, y1 = rect + raw_points, raw_descrs = self.detect_features(image) + points, descs = [], [] + for kp, desc in zip(raw_points, raw_descrs): + x, y = kp.pt + if x0 <= x <= x1 and y0 <= y <= y1: + points.append(kp) + descs.append(desc) + descs = np.uint8(descs) + self.matcher.add([descs]) + target = PlanarTarget(image = image, rect=rect, keypoints = points, descrs=descs, data=None) + self.targets.append(target) + + def clear(self): + '''Remove all targets''' + self.targets = [] + self.matcher.clear() + + def track(self, frame): + '''Returns a list of detected TrackedTarget objects''' + self.frame_points, self.frame_descrs = self.detect_features(frame) + if len(self.frame_points) < MIN_MATCH_COUNT: + return [] + matches = self.matcher.knnMatch(self.frame_descrs, k = 2) + matches = [m[0] for m in matches if len(m) == 2 and m[0].distance < m[1].distance * 0.75] + if len(matches) < MIN_MATCH_COUNT: + return [] + matches_by_id = [[] for _ in xrange(len(self.targets))] + for m in matches: + matches_by_id[m.imgIdx].append(m) + tracked = [] + for imgIdx, matches in enumerate(matches_by_id): + if len(matches) < MIN_MATCH_COUNT: + continue + target = self.targets[imgIdx] + p0 = [target.keypoints[m.trainIdx].pt for m in matches] + p1 = [self.frame_points[m.queryIdx].pt for m in matches] + p0, p1 = np.float32((p0, p1)) + H, status = cv2.findHomography(p0, p1, cv2.RANSAC, 3.0) + status = status.ravel() != 0 + if status.sum() < MIN_MATCH_COUNT: + continue + p0, p1 = p0[status], p1[status] + + x0, y0, x1, y1 = target.rect + quad = np.float32([[x0, y0], [x1, y0], [x1, y1], [x0, y1]]) + quad = cv2.perspectiveTransform(quad.reshape(1, -1, 2), H).reshape(-1, 2) + + track = TrackedTarget(target=target, p0=p0, p1=p1, H=H, quad=quad) + tracked.append(track) + tracked.sort(key = lambda t: len(t.p0), reverse=True) + return tracked + + def detect_features(self, frame): + '''detect_features(self, frame) -> keypoints, descrs''' + keypoints, descrs = self.detector.detectAndCompute(frame, None) + if descrs is None: # detectAndCompute returns descs=None if not keypoints found + descrs = [] + return keypoints, descrs + + +class App: + def __init__(self, src): + self.cap = video.create_capture(src) + self.frame = None + self.paused = False + self.tracker = PlaneTracker() + + cv2.namedWindow('plane') + self.rect_sel = common.RectSelector('plane', self.on_rect) + + def on_rect(self, rect): + self.tracker.add_target(self.frame, rect) + + def run(self): + while True: + playing = not self.paused and not self.rect_sel.dragging + if playing or self.frame is None: + ret, frame = self.cap.read() + if not ret: + break + self.frame = frame.copy() + + vis = self.frame.copy() + if playing: + tracked = self.tracker.track(self.frame) + for tr in tracked: + cv2.polylines(vis, [np.int32(tr.quad)], True, (255, 255, 255), 2) + for (x, y) in np.int32(tr.p1): + cv2.circle(vis, (x, y), 2, (255, 255, 255)) + + self.rect_sel.draw(vis) + cv2.imshow('plane', vis) + ch = cv2.waitKey(1) + if ch == ord(' '): + self.paused = not self.paused + if ch == ord('c'): + self.tracker.clear() + if ch == 27: + break + +if __name__ == '__main__': + print __doc__ + + import sys + try: video_src = sys.argv[1] + except: video_src = 0 + App(video_src).run()