Update to 2.0.0 tree from current Fremantle build
[opencv] / samples / swig_python / motempl.py
1 #!/usr/bin/python
2 from opencv.cv import *
3 from opencv.highgui import *
4 import sys
5 import time
6 from math import cos,sin
7
8 CLOCKS_PER_SEC = 1.0
9 MHI_DURATION = 1
10 MAX_TIME_DELTA = 0.5
11 MIN_TIME_DELTA = 0.05
12 N = 4
13 buf = range(10) 
14 last = 0
15 mhi = None # MHI
16 orient = None # orientation
17 mask = None # valid orientation mask
18 segmask = None # motion segmentation map
19 storage = None # temporary storage
20
21 def update_mhi( img, dst, diff_threshold ):
22     global last
23     global mhi
24     global storage
25     global mask
26     global orient
27     global segmask
28     timestamp = time.clock()/CLOCKS_PER_SEC # get current time in seconds
29     size = cvSize(img.width,img.height) # get current frame size
30     idx1 = last
31     if not mhi or mhi.width != size.width or mhi.height != size.height: 
32         for i in range( N ):
33             buf[i] = cvCreateImage( size, IPL_DEPTH_8U, 1 )
34             cvZero( buf[i] )
35         mhi = cvCreateImage( size, IPL_DEPTH_32F, 1 )
36         cvZero( mhi ) # clear MHI at the beginning
37         orient = cvCreateImage( size, IPL_DEPTH_32F, 1 )
38         segmask = cvCreateImage( size, IPL_DEPTH_32F, 1 )
39         mask = cvCreateImage( size, IPL_DEPTH_8U, 1 )
40     
41     cvCvtColor( img, buf[last], CV_BGR2GRAY ) # convert frame to grayscale
42     idx2 = (last + 1) % N # index of (last - (N-1))th frame
43     last = idx2
44     silh = buf[idx2]
45     cvAbsDiff( buf[idx1], buf[idx2], silh ) # get difference between frames
46     cvThreshold( silh, silh, diff_threshold, 1, CV_THRESH_BINARY ) # and threshold it
47     cvUpdateMotionHistory( silh, mhi, timestamp, MHI_DURATION ) # update MHI
48     cvCvtScale( mhi, mask, 255./MHI_DURATION,
49                 (MHI_DURATION - timestamp)*255./MHI_DURATION )
50     cvZero( dst )
51     cvMerge( mask, None, None, None, dst )
52     cvCalcMotionGradient( mhi, mask, orient, MAX_TIME_DELTA, MIN_TIME_DELTA, 3 )
53     if( not storage ):
54         storage = cvCreateMemStorage(0)
55     else:
56         cvClearMemStorage(storage)
57     seq = cvSegmentMotion( mhi, segmask, storage, timestamp, MAX_TIME_DELTA )
58     for i in range(-1, seq.total):
59         if( i < 0 ):  # case of the whole image
60             comp_rect = cvRect( 0, 0, size.width, size.height )
61             color = CV_RGB(255,255,255)
62             magnitude = 100.
63         else:  # i-th motion component
64             comp_rect = seq[i].rect 
65             if( comp_rect.width + comp_rect.height < 100 ): # reject very small components
66                 continue
67             color = CV_RGB(255,0,0)
68             magnitude = 30.
69         silh_roi = cvGetSubRect(silh, comp_rect)
70         mhi_roi = cvGetSubRect( mhi, comp_rect )
71         orient_roi = cvGetSubRect( orient, comp_rect )
72         mask_roi = cvGetSubRect( mask, comp_rect )
73         angle = cvCalcGlobalOrientation( orient_roi, mask_roi, mhi_roi, timestamp, MHI_DURATION)
74         angle = 360.0 - angle  # adjust for images with top-left origin
75         count = cvNorm( silh_roi, None, CV_L1, None ) # calculate number of points within silhouette ROI
76         if( count < comp_rect.width * comp_rect.height * 0.05 ):
77             continue
78         center = cvPoint( (comp_rect.x + comp_rect.width/2),
79                           (comp_rect.y + comp_rect.height/2) )
80         cvCircle( dst, center, cvRound(magnitude*1.2), color, 3, CV_AA, 0 )
81         cvLine( dst, center, cvPoint( cvRound( center.x + magnitude*cos(angle*CV_PI/180)),
82                 cvRound( center.y - magnitude*sin(angle*CV_PI/180))), color, 3, CV_AA, 0 )
83
84 if __name__ == "__main__":
85     motion = 0
86     capture = 0
87
88     if len(sys.argv)==1:
89         capture = cvCreateCameraCapture( 0 )
90     elif len(sys.argv)==2 and sys.argv[1].isdigit():
91         capture = cvCreateCameraCapture( int(sys.argv[1]) )
92     elif len(sys.argv)==2:
93         capture = cvCreateFileCapture( sys.argv[1] ) 
94
95     if not capture:
96         print "Could not initialize capturing..."
97         sys.exit(-1)
98         
99     cvNamedWindow( "Motion", 1 )
100     while True:
101         image = cvQueryFrame( capture )
102         if( image ):
103             if( not motion ):
104                     motion = cvCreateImage( cvSize(image.width,image.height), 8, 3 )
105                     cvZero( motion )
106                     #motion.origin = image.origin
107             update_mhi( image, motion, 30 )
108             cvShowImage( "Motion", motion )
109             if( cvWaitKey(10) != -1 ):
110                 break
111         else:
112             break
113     cvDestroyWindow( "Motion" )