56 lines
1.8 KiB
Python
56 lines
1.8 KiB
Python
import cv2
|
||
import numpy as np
|
||
from config import KF_PROCESS_NOISE, KF_MEASUREMENT_NOISE
|
||
|
||
class KalmanFilter:
|
||
"""卡尔曼滤波器：平滑并预测装甲板位置"""
|
||
def __init__(self):
|
||
self.kf = cv2.KalmanFilter(4, 2) # 4状态(x,y,dx,dy)，2测量(x,y)
|
||
self._init_params()
|
||
self.initialized = False
|
||
self.last_measurement = None
|
||
self.last_prediction = None
|
||
|
||
def _init_params(self):
|
||
"""初始化卡尔曼参数"""
|
||
self.kf.transitionMatrix = np.array([
|
||
[1, 0, 1, 0],
|
||
[0, 1, 0, 1],
|
||
[0, 0, 1, 0],
|
||
[0, 0, 0, 1]
|
||
], np.float32)
|
||
|
||
self.kf.measurementMatrix = np.array([
|
||
[1, 0, 0, 0],
|
||
[0, 1, 0, 0]
|
||
], np.float32)
|
||
|
||
self.kf.processNoiseCov = np.eye(4, dtype=np.float32) * KF_PROCESS_NOISE
|
||
self.kf.measurementNoiseCov = np.eye(2, dtype=np.float32) * KF_MEASUREMENT_NOISE
|
||
self.kf.errorCovPost = np.eye(4, dtype=np.float32)
|
||
|
||
def update(self, measurement):
|
||
"""用测量值更新滤波器"""
|
||
measurement = np.array([[np.float32(measurement[0])], [np.float32(measurement[1])]])
|
||
|
||
if not self.initialized:
|
||
self.kf.statePost = np.array([
|
||
[measurement[0, 0]],
|
||
[measurement[1, 0]],
|
||
[0.0],
|
||
[0.0]
|
||
], np.float32)
|
||
self.initialized = True
|
||
else:
|
||
self.kf.correct(measurement)
|
||
|
||
self.last_measurement = (measurement[0, 0], measurement[1, 0])
|
||
|
||
def predict(self):
|
||
"""预测下一帧位置"""
|
||
if not self.initialized:
|
||
return None
|
||
|
||
prediction = self.kf.predict()
|
||
self.last_prediction = (prediction[0, 0], prediction[1, 0])
|
||
return self.last_prediction |