-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathstabilizer.py
129 lines (100 loc) · 4.71 KB
/
stabilizer.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
"""
Using Kalman Filter as a point stabilizer to stabiliz a 2D point.
"""
import numpy as np
import cv2
class Stabilizer:
"""Using Kalman filter as a point stabilizer."""
def __init__(self,
state_num=4,
measure_num=2,
cov_process=0.0001,
cov_measure=0.1):
"""Initialization"""
# Currently we only support scalar and point, so check user input first.
assert state_num == 4 or state_num == 2, "Only scalar and point supported, Check state_num please."
# Store the parameters.
self.state_num = state_num
self.measure_num = measure_num
# The filter itself.
self.filter = cv2.KalmanFilter(state_num, measure_num, 0)
# Store the state.
self.state = np.zeros((state_num, 1), dtype=np.float32)
# Store the measurement result.
self.measurement = np.array((measure_num, 1), np.float32)
# Store the prediction.
self.prediction = np.zeros((state_num, 1), np.float32)
# Kalman parameters setup for scalar.
if self.measure_num == 1:
self.filter.transitionMatrix = np.array([[1, 1],
[0, 1]], np.float32)
self.filter.measurementMatrix = np.array([[1, 1]], np.float32)
self.filter.processNoiseCov = np.array([[1, 0],
[0, 1]], np.float32) * cov_process
self.filter.measurementNoiseCov = np.array(
[[1]], np.float32) * cov_measure
# Kalman parameters setup for point.
if self.measure_num == 2:
self.filter.transitionMatrix = np.array([[1, 0, 1, 0],
[0, 1, 0, 1],
[0, 0, 1, 0],
[0, 0, 0, 1]], np.float32)
self.filter.measurementMatrix = np.array([[1, 0, 0, 0],
[0, 1, 0, 0]], np.float32)
self.filter.processNoiseCov = np.array([[1, 0, 0, 0],
[0, 1, 0, 0],
[0, 0, 1, 0],
[0, 0, 0, 1]], np.float32) * cov_process
self.filter.measurementNoiseCov = np.array([[1, 0],
[0, 1]], np.float32) * cov_measure
def update(self, measurement):
"""Update the filter"""
# Make kalman prediction
self.prediction = self.filter.predict()
# Get new measurement
if self.measure_num == 1:
self.measurement = np.array([[np.float32(measurement[0])]])
else:
self.measurement = np.array([[np.float32(measurement[0])],
[np.float32(measurement[1])]])
# Correct according to mesurement
self.filter.correct(self.measurement)
# Update state value.
self.state = self.filter.statePost
def set_q_r(self, cov_process=0.1, cov_measure=0.001):
"""Set new value for processNoiseCov and measurementNoiseCov."""
if self.measure_num == 1:
self.filter.processNoiseCov = np.array([[1, 0],
[0, 1]], np.float32) * cov_process
self.filter.measurementNoiseCov = np.array(
[[1]], np.float32) * cov_measure
else:
self.filter.processNoiseCov = np.array([[1, 0, 0, 0],
[0, 1, 0, 0],
[0, 0, 1, 0],
[0, 0, 0, 1]], np.float32) * cov_process
self.filter.measurementNoiseCov = np.array([[1, 0],
[0, 1]], np.float32) * cov_measure
def main():
"""Test code"""
global mp
mp = np.array((2, 1), np.float32) # measurement
def onmouse(k, x, y, s, p):
global mp
mp = np.array([[np.float32(x)], [np.float32(y)]])
cv2.namedWindow("kalman")
cv2.setMouseCallback("kalman", onmouse)
kalman = Stabilizer(4, 2)
frame = np.zeros((480, 640, 3), np.uint8) # drawing canvas
while True:
kalman.update(mp)
point = kalman.prediction
state = kalman.filter.statePost
cv2.circle(frame, (state[0], state[1]), 2, (255, 0, 0), -1)
cv2.circle(frame, (point[0], point[1]), 2, (0, 255, 0), -1)
cv2.imshow("kalman", frame)
k = cv2.waitKey(30) & 0xFF
if k == 27:
break
if __name__ == '__main__':
main()