vcs.maemo.org Git - opencv/blob - samples/swig_python/kalman.py

   1 #!/usr/bin/python
   2 """
   3    Tracking of rotating point.
   4    Rotation speed is constant.
   5    Both state and measurements vectors are 1D (a point angle),
   6    Measurement is the real point angle + gaussian noise.
   7    The real and the estimated points are connected with yellow line segment,
   8    the real and the measured points are connected with red line segment.
   9    (if Kalman filter works correctly,
  10     the yellow segment should be shorter than the red one).
  11    Pressing any key (except ESC) will reset the tracking with a different speed.
  12    Pressing ESC will stop the program.
  13 """
  14 from opencv.cv import *
  15 from opencv.highgui import *
  16 from math import cos, sin, sqrt
  17
  18 if __name__ == "__main__":
  19     A = [ [1, 1], [0, 1] ]
  20
  21     img = cvCreateImage( cvSize(500,500), 8, 3 )
  22     kalman = cvCreateKalman( 2, 1, 0 )
  23     state = cvCreateMat( 2, 1, CV_32FC1 )  # (phi, delta_phi)
  24     process_noise = cvCreateMat( 2, 1, CV_32FC1 )
  25     measurement = cvCreateMat( 1, 1, CV_32FC1 )
  26     rng = cvRNG(-1)
  27     code = -1L
  28
  29     cvZero( measurement )
  30     cvNamedWindow( "Kalman", 1 )
  31
  32     while True:
  33         cvRandArr( rng, state, CV_RAND_NORMAL, cvRealScalar(0), cvRealScalar(0.1) )
  34
  35         kalman.transition_matrix[:] = A
  36         cvSetIdentity( kalman.measurement_matrix, cvRealScalar(1) )
  37         cvSetIdentity( kalman.process_noise_cov, cvRealScalar(1e-5) )
  38         cvSetIdentity( kalman.measurement_noise_cov, cvRealScalar(1e-1) )
  39         cvSetIdentity( kalman.error_cov_post, cvRealScalar(1))
  40         cvRandArr( rng, kalman.state_post, CV_RAND_NORMAL, cvRealScalar(0), cvRealScalar(0.1) )
  41
  42         while True:
  43             def calc_point(angle):
  44                 return cvPoint( cvRound(img.width/2 + img.width/3*cos(angle)),
  45                          cvRound(img.height/2 - img.width/3*sin(angle)))
  46
  47             state_angle = state[0]
  48             state_pt = calc_point(state_angle)
  49
  50             prediction = cvKalmanPredict( kalman )
  51             predict_angle = prediction[0,0]
  52             predict_pt = calc_point(predict_angle)
  53
  54             cvRandArr( rng, measurement, CV_RAND_NORMAL, cvRealScalar(0),
  55                        cvRealScalar(sqrt(kalman.measurement_noise_cov[0,0])) )
  56
  57             # generate measurement
  58             cvMatMulAdd( kalman.measurement_matrix, state, measurement, measurement )
  59
  60             measurement_angle = measurement[0,0]
  61             measurement_pt = calc_point(measurement_angle)
  62
  63             # plot points
  64             def draw_cross( center, color, d ):
  65                 cvLine( img, cvPoint( center.x - d, center.y - d ),
  66                              cvPoint( center.x + d, center.y + d ), color, 1, CV_AA, 0)
  67                 cvLine( img, cvPoint( center.x + d, center.y - d ),
  68                              cvPoint( center.x - d, center.y + d ), color, 1, CV_AA, 0 )
  69
  70             cvZero( img )
  71             draw_cross( state_pt, CV_RGB(255,255,255), 3 )
  72             draw_cross( measurement_pt, CV_RGB(255,0,0), 3 )
  73             draw_cross( predict_pt, CV_RGB(0,255,0), 3 )
  74             cvLine( img, state_pt, measurement_pt, CV_RGB(255,0,0), 3, CV_AA, 0 )
  75             cvLine( img, state_pt, predict_pt, CV_RGB(255,255,0), 3, CV_AA, 0 )
  76
  77             cvKalmanCorrect( kalman, measurement )
  78
  79             cvRandArr( rng, process_noise, CV_RAND_NORMAL, cvRealScalar(0),
  80                        cvRealScalar(sqrt(kalman.process_noise_cov[0,0])))
  81             cvMatMulAdd( kalman.transition_matrix, state, process_noise, state )
  82
  83             cvShowImage( "Kalman", img )
  84
  85             code = str(cvWaitKey( 100 ))
  86             if( code != '-1'):
  87                 break
  88
  89         if( code == '\x1b' or code == 'q' or code == 'Q' ):
  90             break
  91
  92     cvDestroyWindow("Kalman")