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"""
14import cv2
15from math import cos, sin
16import numpy as np
17
18if __name__ == "__main__":
19
20    img_height = 500
21    img_width = 500
22    kalman = cv2.KalmanFilter(2, 1, 0)
23
24    code = -1L
25
26    cv2.namedWindow("Kalman")
27
28    while True:
29        state = 0.1 * np.random.randn(2, 1)
30
31        kalman.transitionMatrix = np.array([[1., 1.], [0., 1.]])
32        kalman.measurementMatrix = 1. * np.ones((1, 2))
33        kalman.processNoiseCov = 1e-5 * np.eye(2)
34        kalman.measurementNoiseCov = 1e-1 * np.ones((1, 1))
35        kalman.errorCovPost = 1. * np.ones((2, 2))
36        kalman.statePost = 0.1 * np.random.randn(2, 1)
37
38        while True:
39            def calc_point(angle):
40                return (np.around(img_width/2 + img_width/3*cos(angle), 0).astype(int),
41                        np.around(img_height/2 - img_width/3*sin(angle), 1).astype(int))
42
43            state_angle = state[0, 0]
44            state_pt = calc_point(state_angle)
45
46            prediction = kalman.predict()
47            predict_angle = prediction[0, 0]
48            predict_pt = calc_point(predict_angle)
49
50            measurement = kalman.measurementNoiseCov * np.random.randn(1, 1)
51
52            # generate measurement
53            measurement = np.dot(kalman.measurementMatrix, state) + measurement
54
55            measurement_angle = measurement[0, 0]
56            measurement_pt = calc_point(measurement_angle)
57
58            # plot points
59            def draw_cross(center, color, d):
60                cv2.line(img,
61                         (center[0] - d, center[1] - d), (center[0] + d, center[1] + d),
62                         color, 1, cv2.LINE_AA, 0)
63                cv2.line(img,
64                         (center[0] + d, center[1] - d), (center[0] - d, center[1] + d),
65                         color, 1, cv2.LINE_AA, 0)
66
67            img = np.zeros((img_height, img_width, 3), np.uint8)
68            draw_cross(np.int32(state_pt), (255, 255, 255), 3)
69            draw_cross(np.int32(measurement_pt), (0, 0, 255), 3)
70            draw_cross(np.int32(predict_pt), (0, 255, 0), 3)
71
72            cv2.line(img, state_pt, measurement_pt, (0, 0, 255), 3, cv2.LINE_AA, 0)
73            cv2.line(img, state_pt, predict_pt, (0, 255, 255), 3, cv2.LINE_AA, 0)
74
75            kalman.correct(measurement)
76
77            process_noise = kalman.processNoiseCov * np.random.randn(2, 1)
78            state = np.dot(kalman.transitionMatrix, state) + process_noise
79
80            cv2.imshow("Kalman", img)
81
82            code = cv2.waitKey(100) % 0x100
83            if code != -1:
84                break
85
86        if code in [27, ord('q'), ord('Q')]:
87            break
88
89    cv2.destroyWindow("Kalman")
90