kalman.py 3.71 KB
Newer Older
a  
Kai Westerkamp committed
1
#!/usr/bin/python
wester committed
2 3 4 5 6 7 8 9 10 11 12 13
"""
   Tracking of rotating point.
   Rotation speed is constant.
   Both state and measurements vectors are 1D (a point angle),
   Measurement is the real point angle + gaussian noise.
   The real and the estimated points are connected with yellow line segment,
   the real and the measured points are connected with red line segment.
   (if Kalman filter works correctly,
    the yellow segment should be shorter than the red one).
   Pressing any key (except ESC) will reset the tracking with a different speed.
   Pressing ESC will stop the program.
"""
wester committed
14 15 16
import urllib2
import cv2.cv as cv
from math import cos, sin, sqrt
wester committed
17 18 19
import sys

if __name__ == "__main__":
wester committed
20
    A = [ [1, 1], [0, 1] ]
wester committed
21

wester committed
22 23 24 25 26 27 28
    img = cv.CreateImage((500, 500), 8, 3)
    kalman = cv.CreateKalman(2, 1, 0)
    state = cv.CreateMat(2, 1, cv.CV_32FC1)  # (phi, delta_phi)
    process_noise = cv.CreateMat(2, 1, cv.CV_32FC1)
    measurement = cv.CreateMat(1, 1, cv.CV_32FC1)
    rng = cv.RNG(-1)
    code = -1L
wester committed
29

wester committed
30 31
    cv.Zero(measurement)
    cv.NamedWindow("Kalman", 1)
wester committed
32 33

    while True:
wester committed
34 35 36 37 38 39 40 41 42 43 44 45
        cv.RandArr(rng, state, cv.CV_RAND_NORMAL, cv.RealScalar(0), cv.RealScalar(0.1))

        kalman.transition_matrix[0,0] = 1
        kalman.transition_matrix[0,1] = 1
        kalman.transition_matrix[1,0] = 0
        kalman.transition_matrix[1,1] = 1

        cv.SetIdentity(kalman.measurement_matrix, cv.RealScalar(1))
        cv.SetIdentity(kalman.process_noise_cov, cv.RealScalar(1e-5))
        cv.SetIdentity(kalman.measurement_noise_cov, cv.RealScalar(1e-1))
        cv.SetIdentity(kalman.error_cov_post, cv.RealScalar(1))
        cv.RandArr(rng, kalman.state_post, cv.CV_RAND_NORMAL, cv.RealScalar(0), cv.RealScalar(0.1))
wester committed
46 47 48 49


        while True:
            def calc_point(angle):
wester committed
50 51
                return (cv.Round(img.width/2 + img.width/3*cos(angle)),
                         cv.Round(img.height/2 - img.width/3*sin(angle)))
wester committed
52

wester committed
53
            state_angle = state[0,0]
wester committed
54 55
            state_pt = calc_point(state_angle)

wester committed
56
            prediction = cv.KalmanPredict(kalman)
wester committed
57 58 59
            predict_angle = prediction[0, 0]
            predict_pt = calc_point(predict_angle)

wester committed
60 61
            cv.RandArr(rng, measurement, cv.CV_RAND_NORMAL, cv.RealScalar(0),
                       cv.RealScalar(sqrt(kalman.measurement_noise_cov[0, 0])))
wester committed
62 63

            # generate measurement
wester committed
64
            cv.MatMulAdd(kalman.measurement_matrix, state, measurement, measurement)
wester committed
65 66 67 68 69 70

            measurement_angle = measurement[0, 0]
            measurement_pt = calc_point(measurement_angle)

            # plot points
            def draw_cross(center, color, d):
wester committed
71 72 73 74
                cv.Line(img, (center[0] - d, center[1] - d),
                             (center[0] + d, center[1] + d), color, 1, cv.CV_AA, 0)
                cv.Line(img, (center[0] + d, center[1] - d),
                             (center[0] - d, center[1] + d), color, 1, cv.CV_AA, 0)
wester committed
75

wester committed
76 77 78 79 80 81
            cv.Zero(img)
            draw_cross(state_pt, cv.CV_RGB(255, 255, 255), 3)
            draw_cross(measurement_pt, cv.CV_RGB(255, 0,0), 3)
            draw_cross(predict_pt, cv.CV_RGB(0, 255, 0), 3)
            cv.Line(img, state_pt, measurement_pt, cv.CV_RGB(255, 0,0), 3, cv. CV_AA, 0)
            cv.Line(img, state_pt, predict_pt, cv.CV_RGB(255, 255, 0), 3, cv. CV_AA, 0)
wester committed
82

wester committed
83
            cv.KalmanCorrect(kalman, measurement)
wester committed
84

wester committed
85 86 87
            cv.RandArr(rng, process_noise, cv.CV_RAND_NORMAL, cv.RealScalar(0),
                       cv.RealScalar(sqrt(kalman.process_noise_cov[0, 0])))
            cv.MatMulAdd(kalman.transition_matrix, state, process_noise, state)
wester committed
88

wester committed
89
            cv.ShowImage("Kalman", img)
wester committed
90

wester committed
91
            code = cv.WaitKey(100) % 0x100
wester committed
92 93 94 95 96 97
            if code != -1:
                break

        if code in [27, ord('q'), ord('Q')]:
            break

wester committed
98
    cv.DestroyWindow("Kalman")