|
|
import cv2
|
|
|
import numpy as np
|
|
|
|
|
|
from config import *
|
|
|
from helpers import clamp, box_center, box_wh, clip_box
|
|
|
|
|
|
# Kalman and KLT trackers
|
|
|
# =========================
|
|
|
# KALMAN 8D
|
|
|
# =========================
|
|
|
|
|
|
class Kalman8D:
|
|
|
def __init__(self):
|
|
|
self.x = np.zeros((8, 1), dtype=np.float32)
|
|
|
self.P = np.eye(8, dtype=np.float32) * 50.0
|
|
|
|
|
|
self.H = np.zeros((4, 8), dtype=np.float32)
|
|
|
self.H[0, 0] = 1.0
|
|
|
self.H[1, 1] = 1.0
|
|
|
self.H[2, 2] = 1.0
|
|
|
self.H[3, 3] = 1.0
|
|
|
|
|
|
self.R = np.eye(4, dtype=np.float32)
|
|
|
# Снижено с 16→9 и 25→16: новая YOLO даёт точные детекты,
|
|
|
# Kalman должен им больше доверять и быстрее подстраиваться.
|
|
|
self.R[0, 0] = 9.0
|
|
|
self.R[1, 1] = 9.0
|
|
|
self.R[2, 2] = 16.0
|
|
|
self.R[3, 3] = 16.0
|
|
|
|
|
|
self.initialized = False
|
|
|
# Процесс-шум повышен для реактивности на резкие манёвры.
|
|
|
# Было: q_pos=2.0, q_size=3.0, q_vel=4.0 — для плавного tracking'а
|
|
|
# Стало: q_pos=5.0, q_size=3.5, q_vel=12.0 — цель до 160 px/s со швырками
|
|
|
self.q_pos = 5.0
|
|
|
self.q_size = 3.5
|
|
|
self.q_vel = 12.0
|
|
|
|
|
|
def F(self, dt):
|
|
|
F = np.eye(8, dtype=np.float32)
|
|
|
F[0, 4] = dt
|
|
|
F[1, 5] = dt
|
|
|
F[2, 6] = dt
|
|
|
F[3, 7] = dt
|
|
|
return F
|
|
|
|
|
|
def Q(self, dt, q_scale=1.0):
|
|
|
q_scale = float(max(0.25, q_scale))
|
|
|
q = np.zeros((8, 8), dtype=np.float32)
|
|
|
s1 = dt * dt
|
|
|
s2 = dt
|
|
|
q[0, 0] = self.q_pos * s1 * q_scale
|
|
|
q[1, 1] = self.q_pos * s1 * q_scale
|
|
|
q[2, 2] = self.q_size * s1 * q_scale
|
|
|
q[3, 3] = self.q_size * s1 * q_scale
|
|
|
q[4, 4] = self.q_vel * s2 * q_scale
|
|
|
q[5, 5] = self.q_vel * s2 * q_scale
|
|
|
q[6, 6] = self.q_vel * s2 * q_scale
|
|
|
q[7, 7] = self.q_vel * s2 * q_scale
|
|
|
return q
|
|
|
|
|
|
def predict(self, dt, q_scale=1.0):
|
|
|
dt = float(max(1e-3, dt))
|
|
|
F = self.F(dt)
|
|
|
self.x = F @ self.x
|
|
|
self.P = F @ self.P @ F.T + self.Q(dt, q_scale=q_scale)
|
|
|
return self.x[:4].flatten()
|
|
|
|
|
|
def update(self, z):
|
|
|
z = np.array(z, dtype=np.float32).reshape(4, 1)
|
|
|
y = z - self.H @ self.x
|
|
|
S = self.H @ self.P @ self.H.T + self.R
|
|
|
K = self.P @ self.H.T @ np.linalg.inv(S)
|
|
|
self.x = self.x + K @ y
|
|
|
I = np.eye(8, dtype=np.float32)
|
|
|
self.P = (I - K @ self.H) @ self.P
|
|
|
|
|
|
def init_from_box(self, box):
|
|
|
cx, cy = box_center(box)
|
|
|
w, h = box_wh(box)
|
|
|
self.x[:] = 0
|
|
|
self.x[0, 0] = cx
|
|
|
self.x[1, 0] = cy
|
|
|
self.x[2, 0] = w
|
|
|
self.x[3, 0] = h
|
|
|
self.P = np.eye(8, dtype=np.float32) * 50.0
|
|
|
self.initialized = True
|
|
|
|
|
|
def to_box(self):
|
|
|
cx = float(self.x[0, 0])
|
|
|
cy = float(self.x[1, 0])
|
|
|
w = float(max(2.0, self.x[2, 0]))
|
|
|
h = float(max(2.0, self.x[3, 0]))
|
|
|
return np.array([cx - w * 0.5, cy - h * 0.5, cx + w * 0.5, cy + h * 0.5], dtype=np.float32)
|
|
|
|
|
|
def uncertainty(self):
|
|
|
return float(self.P[0, 0] + self.P[1, 1] + self.P[2, 2] + self.P[3, 3])
|
|
|
|
|
|
|
|
|
# =========================
|
|
|
# KLT TRACKER
|
|
|
# =========================
|
|
|
|
|
|
class KLTTracker:
|
|
|
def __init__(self):
|
|
|
self.prev_gray = None
|
|
|
self.pts = None
|
|
|
self.box = None
|
|
|
self.good_count = 0
|
|
|
self.quality = 0.0
|
|
|
|
|
|
def reset(self):
|
|
|
self.prev_gray = None
|
|
|
self.pts = None
|
|
|
self.box = None
|
|
|
self.good_count = 0
|
|
|
self.quality = 0.0
|
|
|
|
|
|
def init(self, frame_bgr, box):
|
|
|
gray = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2GRAY)
|
|
|
x1, y1, x2, y2 = map(int, box)
|
|
|
mask = np.zeros_like(gray)
|
|
|
cv2.rectangle(mask, (x1, y1), (x2, y2), 255, -1)
|
|
|
|
|
|
pts = cv2.goodFeaturesToTrack(
|
|
|
gray,
|
|
|
maxCorners=KLT_MAX_CORNERS,
|
|
|
qualityLevel=KLT_QUALITY,
|
|
|
minDistance=KLT_MIN_DIST,
|
|
|
mask=mask
|
|
|
)
|
|
|
self.prev_gray = gray
|
|
|
self.pts = pts
|
|
|
self.box = box.copy()
|
|
|
self.good_count = 0
|
|
|
self.quality = 0.0
|
|
|
|
|
|
def update(self, frame_bgr):
|
|
|
if self.prev_gray is None or self.pts is None or len(self.pts) == 0 or self.box is None:
|
|
|
self.quality = 0.0
|
|
|
self.good_count = 0
|
|
|
return None
|
|
|
|
|
|
gray = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2GRAY)
|
|
|
next_pts, st, err = cv2.calcOpticalFlowPyrLK(
|
|
|
self.prev_gray, gray, self.pts, None,
|
|
|
winSize=KLT_WIN,
|
|
|
maxLevel=KLT_MAX_LEVEL,
|
|
|
criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 30, 0.01)
|
|
|
)
|
|
|
|
|
|
if next_pts is None or st is None:
|
|
|
self.quality = 0.0
|
|
|
self.good_count = 0
|
|
|
self.prev_gray = gray
|
|
|
return None
|
|
|
|
|
|
st = st.reshape(-1)
|
|
|
good_old = self.pts[st == 1].reshape(-1, 2)
|
|
|
good_new = next_pts[st == 1].reshape(-1, 2)
|
|
|
|
|
|
self.good_count = int(len(good_new))
|
|
|
if self.good_count < 6:
|
|
|
self.quality = 0.0
|
|
|
self.prev_gray = gray
|
|
|
self.pts = good_new.reshape(-1, 1, 2) if self.good_count > 0 else None
|
|
|
return None
|
|
|
|
|
|
disp = good_new - good_old
|
|
|
med = np.median(disp, axis=0)
|
|
|
d = np.linalg.norm(disp - med[None, :], axis=1)
|
|
|
inliers = d < KLT_OUTLIER_THRESH
|
|
|
in_cnt = int(np.count_nonzero(inliers))
|
|
|
|
|
|
if in_cnt < 6:
|
|
|
self.quality = 0.0
|
|
|
self.prev_gray = gray
|
|
|
self.pts = good_new.reshape(-1, 1, 2)
|
|
|
return None
|
|
|
|
|
|
disp_in = disp[inliers]
|
|
|
med = np.median(disp_in, axis=0)
|
|
|
dx, dy = float(med[0]), float(med[1])
|
|
|
|
|
|
x1, y1, x2, y2 = self.box
|
|
|
new_box = np.array([x1 + dx, y1 + dy, x2 + dx, y2 + dy], dtype=np.float32)
|
|
|
|
|
|
frac = float(in_cnt) / float(len(inliers))
|
|
|
self.quality = clamp(frac * (self.good_count / max(1, KLT_MAX_CORNERS)), 0.0, 1.0)
|
|
|
|
|
|
self.prev_gray = gray
|
|
|
self.pts = good_new[inliers].reshape(-1, 1, 2)
|
|
|
self.box = new_box
|
|
|
return new_box
|
|
|
|
|
|
|
|
|
# =========================
|