class KalmanFilterTracker(QWidget):
def __init__(self, title='', image=''):
super().__init__()
self.title = title
self.desktop = QDesktopWidget()
self.screen = self.desktop.availableGeometry()
self.logger = logging.getLogger(self.__class__.__name__)
# drawing setting
self._drawing = False
self._image = image
self._w, self._h, self._c = (1280, 720, 3)
self._pen_measure = QPen(Qt.black, 3, cap=Qt.RoundCap)
self._pen_measure_line = QPen(Qt.black, 1, style=Qt.SolidLine)
self._pen_predict = QPen(Qt.blue, 3, cap=Qt.RoundCap)
self._pen_predict_line = QPen(Qt.blue, 1, style=Qt.SolidLine)
self._pen_correct = QPen(Qt.red, 5, cap=Qt.RoundCap)
self._pen_correct_line = QPen(Qt.red, 1, style=Qt.SolidLine)
self._measure_points = QPolygon()
self._predict_points = QPolygon()
self._correct_points = QPolygon()
self._measure_text = f"Mouse measurement - ()"
self._predict_text = f"Kalman Filter predict - ()"
self._correct_text = f"Kalman Filter correct - ()"
# init
self.init_ui()
def _ndarray_to_qimage(self, arr):
w, h = arr.shape[:2]
return QImage(arr.data, w, h, w * 3, QImage.Format_RGB888)
def _qimage_to_qpixmap(self, qimg):
return QPixmap.fromImage(qimg)
def _reset_qpixmap(self):
self._pixmap = self._qimage_to_qpixmap(self._image)
def _reset_polygon(self):
self._measure_points = QPolygon()
self._correct_points = QPolygon()
self._predict_points = QPolygon()
def init_ui(self):
# set windows
self.setWindowTitle(self.title)
self.resize(self._w, self._h)
frame_geo = self.frameGeometry()
frame_geo.moveCenter(self.screen.center())
self.move(frame_geo.topLeft())
# set layout
self.root = QGridLayout()
self._hbox_body = QHBoxLayout()
self.setLayout(self.root)
self.show()
# set image
if self._image:
self._w, self.h = self._image.shape[:2]
else:
self._image = np.zeros((self._w, self._h, self._c),
dtype='float32')
self._image = self._ndarray_to_qimage(self._image)
self._pixmap = self._qimage_to_qpixmap(self._image)
def init_kalman_filter(self, x: float, y: float):
dim_x = 4 # Number of state variables (x, y, x_velocity, y_velocity)
dim_z = 2 # Number of of measurement inputs (x, y)
self.kf = KalmanFilter(dim_x=dim_x, dim_z=dim_z)
self.kf.x = np.array([x, y, 0., 0.
]) # init state (position, velocity) (dim_x, 1)
self.kf.F = np.array([
[1., 0., 1., 0.], # state transition matrix (dim_x, dim_x)
[0., 1., 0., 1.],
[0., 0., 1., 0.],
[0., 0., 0., 1.]
])
self.kf.H = np.array([
[1., 0., 0., 0.], # measurement function, ((dim_z, dim_x))
[0., 1., 0., 0.]
])
self.kf.P = np.eye(dim_x) * 1000. # covariance matrix (dim_x, dim_x)
self.kf.R = np.eye(
dim_z) * 10. # measurement noise covariance (dim_z, dim_z)
self.kf.Q = np.eye(dim_x) * 1e-3 # process uncertainty (dim_x, dim_x)
self._predict_text = f"Kalman Filter predict - ()"
self._correct_text = f"Kalman Filter correct - ()"
def paintEvent(self, e):
painter = QPainter(self)
painter.setRenderHint(QPainter.Antialiasing)
assert self._measure_points.count() == \
self._predict_points.count() == \
self._correct_points.count()
for i in range(self._measure_points.count()):
painter.setPen(self._pen_measure)
painter.drawPoint(self._measure_points.point(i))
painter.setPen(self._pen_predict)
painter.drawPoint(self._predict_points.point(i))
painter.setPen(self._pen_correct)
painter.drawPoint(self._correct_points.point(i))
if i:
painter.setPen(self._pen_measure_line)
painter.drawLine(self._measure_points.point(i - 1),
self._measure_points.point(i))
painter.setPen(self._pen_predict_line)
painter.drawLine(self._predict_points.point(i - 1),
self._predict_points.point(i))
painter.setPen(self._pen_correct_line)
painter.drawLine(self._correct_points.point(i - 1),
self._correct_points.point(i))
painter.setPen(self._pen_measure_line)
painter.drawText(50, 60, self._measure_text)
painter.setPen(self._pen_predict_line)
painter.drawText(50, 80, self._predict_text)
painter.setPen(self._pen_correct_line)
painter.drawText(50, 100, self._correct_text)
def mousePressEvent(self, e):
if e.button() == Qt.LeftButton:
self._drawing = True
self.logger.info(f"========== Reset QPolygon ==========")
x, y = float(e.x()), float(e.y())
self._reset_polygon()
self.init_kalman_filter(x, y)
self._measure_text = f"Mouse measurement - ({e.x()}, {e.y()})"
def mouseMoveEvent(self, e):
if e.buttons() and Qt.LeftButton and self._drawing:
self.logger.debug(f"Record {e.pos()}")
# measurement
self._measure_points << e.pos()
self._measure_text = f"Mouse measurement - ({e.x()}, {e.y()})"
# predict
self.kf.predict()
predict_pts = self.kf.x[:2]
predict_x, predict_y = predict_pts
predict_x, predict_y = int(predict_x), int(predict_y)
self._predict_points << QPoint(predict_x, predict_y)
self._predict_text = f"Kalman Filter predict - ({predict_pts[0]:.4f}, {predict_pts[1]:4f})"
# correct
self.kf.update(
np.array([e.x(), e.y()], dtype='float32').reshape((2, 1)))
correct_pts = self.kf.x[:2]
correct_x, correct_y = correct_pts
correct_x, correct_y = int(correct_x), int(correct_y)
self._correct_points << QPoint(correct_x, correct_y)
self._correct_text = f"Kalman Filter correct - ({correct_pts[0]:4f}, {correct_pts[1]:4f})"
# draw
self.logger.info(
f"mouse ({e.x()}, {e.y()}); KF predict {predict_pts}; KF correct {correct_pts}"
)
self.update()
def mouseReleaseEvent(self, e):
if e.button() == Qt.LeftButton:
self._drawing = False
