<s> '''<STR_LIT>''' <EOL> import numpy as np <EOL> import cv2 <EOL> from collections import namedtuple <EOL> import video <EOL> import common <EOL> FLANN_INDEX_KDTREE = <NUM_LIT:1> <EOL> FLANN_INDEX_LSH = <NUM_LIT:6> <EOL> flann_params = dict ( algorithm = FLANN_INDEX_LSH , <EOL> table_number = <NUM_LIT:6> , <EOL> key_size = <NUM_LIT:12> , <EOL> multi_probe_level = <NUM_LIT:1> ) <EOL> MIN_MATCH_COUNT = <NUM_LIT:10> <EOL> '''<STR_LIT>''' <EOL> PlanarTarget = namedtuple ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> '''<STR_LIT>''' <EOL> TrackedTarget = namedtuple ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> class PlaneTracker : <EOL> def __init__ ( self ) : <EOL> self . detector = cv2 . ORB ( nfeatures = <NUM_LIT:1000> ) <EOL> self . matcher = cv2 . FlannBasedMatcher ( flann_params , { } ) <EOL> self . targets = [ ] <EOL> def add_target ( self , image , rect , data = None ) : <EOL> '''<STR_LIT>''' <EOL> x0 , y0 , x1 , y1 = rect <EOL> raw_points , raw_descrs = self . detect_features ( image ) <EOL> points , descs = [ ] , [ ] <EOL> for kp , desc in zip ( raw_points , raw_descrs ) : <EOL> x , y = kp . pt <EOL> if x0 <= x <= x1 and y0 <= y <= y1 : <EOL> points . append ( kp ) <EOL> descs . append ( desc ) <EOL> descs = np . uint8 ( descs ) <EOL> self . matcher . add ( [ descs ] ) <EOL> target = PlanarTarget ( image = image , rect = rect , keypoints = points , descrs = descs , data = None ) <EOL> self . targets . append ( target ) <EOL> def clear ( self ) : <EOL> '''<STR_LIT>''' <EOL> self . targets = [ ] <EOL> self . matcher . clear ( ) <EOL> def track ( self , frame ) : <EOL> '''<STR_LIT>''' <EOL> self . frame_points , self . frame_descrs = self . detect_features ( frame ) <EOL> if len ( self . frame_points ) < MIN_MATCH_COUNT : <EOL> return [ ] <EOL> matches = self . matcher . knnMatch ( self . frame_descrs , k = <NUM_LIT:2> ) <EOL> matches = [ m [ <NUM_LIT:0> ] for m in matches if len ( m ) == <NUM_LIT:2> and m [ <NUM_LIT:0> ] . distance < m [ <NUM_LIT:1> ] . distance * <NUM_LIT> ] <EOL> if len ( matches ) < MIN_MATCH_COUNT : <EOL> return [ ] <EOL> matches_by_id = [ [ ] for _ in xrange ( len ( self . targets ) ) ] <EOL> for m in matches : <EOL> matches_by_id [ m . imgIdx ] . append ( m ) <EOL> tracked = [ ] <EOL> for imgIdx , matches in enumerate ( matches_by_id ) : <EOL> if len ( matches ) < MIN_MATCH_COUNT : <EOL> continue <EOL> target = self . targets [ imgIdx ] <EOL> p0 = [ target . keypoints [ m . trainIdx ] . pt for m in matches ] <EOL> p1 = [ self . frame_points [ m . queryIdx ] . pt for m in matches ] <EOL> p0 , p1 = np . float32 ( ( p0 , p1 ) ) <EOL> H , status = cv2 . findHomography ( p0 , p1 , cv2 . RANSAC , <NUM_LIT> ) <EOL> status = status . ravel ( ) != <NUM_LIT:0> <EOL> if status . sum ( ) < MIN_MATCH_COUNT : <EOL> continue <EOL> p0 , p1 = p0 [ status ] , p1 [ status ] <EOL> x0 , y0 , x1 , y1 = target . rect <EOL> quad = np . float32 ( [ [ x0 , y0 ] , [ x1 , y0 ] , [ x1 , y1 ] , [ x0 , y1 ] ] ) <EOL> quad = cv2 . perspectiveTransform ( quad . reshape ( <NUM_LIT:1> , - <NUM_LIT:1> , <NUM_LIT:2> ) , H ) . reshape ( - <NUM_LIT:1> , <NUM_LIT:2> ) <EOL> track = TrackedTarget ( target = target , p0 = p0 , p1 = p1 , H = H , quad = quad ) <EOL> tracked . append ( track ) <EOL> tracked . sort ( key = lambda t : len ( t . p0 ) , reverse = True ) <EOL> return tracked <EOL> def detect_features ( self , frame ) : <EOL> '''<STR_LIT>''' <EOL> keypoints , descrs = self . detector . detectAndCompute ( frame , None ) <EOL> if descrs is None : <EOL> descrs = [ ] <EOL> return keypoints , descrs <EOL> class App : <EOL> def __init__ ( self , src ) : <EOL> self . cap = video . create_capture ( src ) <EOL> self . frame = None <EOL> self . paused = False <EOL> self . tracker = PlaneTracker ( ) <EOL> cv2 . namedWindow ( '<STR_LIT>' ) <EOL> self . rect_sel = common . RectSelector ( '<STR_LIT>' , self . on_rect ) <EOL> def on_rect ( self , rect ) : <EOL> self . tracker . add_target ( self . frame , rect ) <EOL> def run ( self ) : <EOL> while True : <EOL> playing = not self . paused and not self . rect_sel . dragging <EOL> if playing or self . frame is None : <EOL> ret , frame = self . cap . read ( ) <EOL> if not ret : <EOL> break <EOL> self . frame = frame . copy ( ) <EOL> vis = self . frame . copy ( ) <EOL> if playing : <EOL> tracked = self . tracker . track ( self . frame ) <EOL> for tr in tracked : <EOL> cv2 . polylines ( vis , [ np . int32 ( tr . quad ) ] , True , ( <NUM_LIT:255> , <NUM_LIT:255> , <NUM_LIT:255> ) , <NUM_LIT:2> ) <EOL> for ( x , y ) in np . int32 ( tr . p1 ) : <EOL> cv2 . circle ( vis , ( x , y ) , <NUM_LIT:2> , ( <NUM_LIT:255> , <NUM_LIT:255> , <NUM_LIT:255> ) ) <EOL> self . rect_sel . draw ( vis ) <EOL> cv2 . imshow ( '<STR_LIT>' , vis ) <EOL> ch = cv2 . waitKey ( <NUM_LIT:1> ) <EOL> if ch == ord ( '<STR_LIT:U+0020>' ) : <EOL> self . paused = not self . paused <EOL> if ch == ord ( '<STR_LIT:c>' ) : <EOL> self . tracker . clear ( ) <EOL> if ch == <NUM_LIT> : <EOL> break <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> print __doc__ <EOL> import sys <EOL> try : <EOL> video_src = sys . argv [ <NUM_LIT:1> ] <EOL> except : <EOL> video_src = <NUM_LIT:0> <EOL> App ( video_src ) . run ( ) </s>
