|
|
# ============================================================
|
|
|
# trackers_hybrid.py (v2 — affine motion for maneuvering targets)
|
|
|
#
|
|
|
# Ключевые отличия от v1:
|
|
|
# 1. KLT путь оценивает ПОЛНУЮ аффинку (translation + rotation
|
|
|
# + scale) через cv2.estimateAffinePartial2D, а не только
|
|
|
# медиану сдвигов. Это необходимо когда цель поворачивается
|
|
|
# в кадре — при rigid translation inlier-точки разлетаются
|
|
|
# и квалити падает.
|
|
|
# 2. Применение аффинки к центру и к размеру bbox с clamp'ом
|
|
|
# на изменение масштаба за кадр (защита от RANSAC-сбоев).
|
|
|
# 3. Более частая переинициализация features (8 вместо 15)
|
|
|
# для быстрого подстраивания под вращающуюся цель.
|
|
|
# 4. Phase correlation теперь работает на НЕСКОЛЬКИХ масштабах
|
|
|
# шаблона (0.9x, 1.0x, 1.1x) чтобы устойчиво ловить цель при
|
|
|
# приближении/удалении.
|
|
|
# ============================================================
|
|
|
|
|
|
import cv2
|
|
|
import numpy as np
|
|
|
|
|
|
from config import *
|
|
|
from helpers import clamp
|
|
|
|
|
|
|
|
|
# ─── Параметры ───────────────────────────────────────────────
|
|
|
HYBRID_CLAHE_CLIP = 2.5
|
|
|
HYBRID_CLAHE_TILE = 4
|
|
|
HYBRID_KLT_QUALITY = 0.005
|
|
|
HYBRID_KLT_QUALITY_RETRY = 0.002
|
|
|
HYBRID_KLT_MIN_DIST = 2
|
|
|
HYBRID_KLT_MAX_CORNERS = 50
|
|
|
HYBRID_KLT_MIN_INLIERS = 3
|
|
|
HYBRID_KLT_MIN_FOR_AFFINE = 6 # меньше — fallback на median translation
|
|
|
HYBRID_KLT_OUTLIER_THR = 3.5
|
|
|
HYBRID_PAD_FACTOR = 0.25
|
|
|
HYBRID_SEARCH_FACTOR = 2.4 # чуть больше search window для манёвров
|
|
|
|
|
|
HYBRID_PHASE_MIN_RESP = 0.22 # чуть ослаблен
|
|
|
HYBRID_TM_MIN_SCORE = 0.42
|
|
|
HYBRID_TM_SCALES = (0.9, 1.0, 1.1)
|
|
|
|
|
|
# Защита от RANSAC-сбоев (sanity check, размер bbox НЕ обновляется)
|
|
|
HYBRID_MAX_SCALE_PER_FRAME = 1.25 # допустимый диапазон нормы аффинки
|
|
|
HYBRID_MIN_SCALE_PER_FRAME = 0.80 # иначе считаем RANSAC сбоем
|
|
|
HYBRID_MAX_TRANSLATION_PX = 80.0 # максимальный сдвиг центра за кадр
|
|
|
|
|
|
HYBRID_TEMPLATE_UPDATE_EMA = 0.80 # чуть быстрее обновлять (было 0.85)
|
|
|
HYBRID_REINIT_FEATURES_EVERY = 8 # было 15 — для динамичной цели
|
|
|
|
|
|
|
|
|
class HybridTracker:
|
|
|
"""
|
|
|
Многоуровневый трекер для мелких манёвренных целей.
|
|
|
|
|
|
API:
|
|
|
trk = HybridTracker()
|
|
|
trk.init(frame_bgr, box)
|
|
|
new_box = trk.update(frame_bgr) # None если потеряли
|
|
|
Поля:
|
|
|
self.box — текущий bbox
|
|
|
self.quality — 0..1
|
|
|
self.last_method — "klt-affine" / "klt-translate" / "phasecorr" / "template" / "lost"
|
|
|
self.good_count — число inliers KLT на последнем кадре
|
|
|
"""
|
|
|
|
|
|
def __init__(self):
|
|
|
self._clahe = cv2.createCLAHE(
|
|
|
clipLimit=float(HYBRID_CLAHE_CLIP),
|
|
|
tileGridSize=(int(HYBRID_CLAHE_TILE), int(HYBRID_CLAHE_TILE)),
|
|
|
)
|
|
|
self.reset()
|
|
|
|
|
|
# ─── Public API ──────────────────────────────────────────
|
|
|
def reset(self):
|
|
|
self.prev_gray = None
|
|
|
self.prev_gray_eq = None
|
|
|
self.pts = None
|
|
|
self.box = None
|
|
|
self.template = None
|
|
|
self.good_count = 0
|
|
|
self.quality = 0.0
|
|
|
self.last_method = "none"
|
|
|
self._frames_since_feature_init = 0
|
|
|
self._success_streak = 0
|
|
|
|
|
|
def init(self, frame_bgr, box):
|
|
|
if frame_bgr is None or box is None:
|
|
|
return
|
|
|
gray = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2GRAY)
|
|
|
gray_eq = self._clahe.apply(gray)
|
|
|
|
|
|
self.box = np.array(box, dtype=np.float32)
|
|
|
self.prev_gray = gray
|
|
|
self.prev_gray_eq = gray_eq
|
|
|
|
|
|
self._extract_features(gray_eq, self.box)
|
|
|
self._update_template(gray_eq, self.box, reset=True)
|
|
|
self._frames_since_feature_init = 0
|
|
|
self._success_streak = 0
|
|
|
self.last_method = "init"
|
|
|
|
|
|
def update(self, frame_bgr):
|
|
|
if self.box is None or self.prev_gray is None:
|
|
|
return None
|
|
|
|
|
|
gray = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2GRAY)
|
|
|
gray_eq = self._clahe.apply(gray)
|
|
|
|
|
|
# ── Попытка 1: KLT (affine или fallback translate) ──
|
|
|
new_box_klt, q_klt, klt_method = self._try_klt(gray, gray_eq)
|
|
|
|
|
|
# ── Попытка 2: Phase correlation ────────────────────
|
|
|
new_box_pc, q_pc = (None, 0.0)
|
|
|
if new_box_klt is None or q_klt < 0.25:
|
|
|
new_box_pc, q_pc = self._try_phase_correlation(gray_eq)
|
|
|
|
|
|
# ── Попытка 3: Template matching (multi-scale) ──────
|
|
|
new_box_tm, q_tm = (None, 0.0)
|
|
|
if new_box_klt is None and new_box_pc is None:
|
|
|
new_box_tm, q_tm = self._try_template_match(gray_eq)
|
|
|
|
|
|
# ── Выбор лучшего ─────────────────────────────────────
|
|
|
candidates = []
|
|
|
if new_box_klt is not None:
|
|
|
candidates.append((klt_method, new_box_klt, q_klt))
|
|
|
if new_box_pc is not None:
|
|
|
candidates.append(("phasecorr", new_box_pc, q_pc))
|
|
|
if new_box_tm is not None:
|
|
|
candidates.append(("template", new_box_tm, q_tm))
|
|
|
|
|
|
if not candidates:
|
|
|
self.last_method = "lost"
|
|
|
self.quality = 0.0
|
|
|
self.prev_gray = gray
|
|
|
self.prev_gray_eq = gray_eq
|
|
|
return None
|
|
|
|
|
|
candidates.sort(key=lambda c: c[2], reverse=True)
|
|
|
method, new_box, q = candidates[0]
|
|
|
|
|
|
self.box = new_box
|
|
|
self.quality = float(q)
|
|
|
self.last_method = method
|
|
|
self._success_streak += 1
|
|
|
self._frames_since_feature_init += 1
|
|
|
|
|
|
# Переинициализация features и обновление template
|
|
|
if self._frames_since_feature_init >= int(HYBRID_REINIT_FEATURES_EVERY):
|
|
|
self._extract_features(gray_eq, self.box)
|
|
|
self._update_template(gray_eq, self.box, reset=False)
|
|
|
self._frames_since_feature_init = 0
|
|
|
elif method in ("phasecorr", "template") and self._success_streak % 4 == 0:
|
|
|
self._extract_features(gray_eq, self.box)
|
|
|
self._update_template(gray_eq, self.box, reset=False)
|
|
|
|
|
|
self.prev_gray = gray
|
|
|
self.prev_gray_eq = gray_eq
|
|
|
return self.box
|
|
|
|
|
|
# ─── Level 1a: KLT with AFFINE motion ────────────────────
|
|
|
def _try_klt(self, gray, gray_eq):
|
|
|
"""
|
|
|
Returns (new_box, quality, method_name).
|
|
|
method_name ∈ {"klt-affine", "klt-translate"} или (None, 0, "none").
|
|
|
"""
|
|
|
if self.pts is None or len(self.pts) == 0:
|
|
|
return None, 0.0, "none"
|
|
|
|
|
|
next_pts, st, err = cv2.calcOpticalFlowPyrLK(
|
|
|
self.prev_gray_eq, gray_eq, self.pts, None,
|
|
|
winSize=(21, 21),
|
|
|
maxLevel=3,
|
|
|
criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 30, 0.01),
|
|
|
)
|
|
|
if next_pts is None or st is None:
|
|
|
self.good_count = 0
|
|
|
return None, 0.0, "none"
|
|
|
|
|
|
st = st.reshape(-1).astype(bool)
|
|
|
good_old = self.pts[st].reshape(-1, 2)
|
|
|
good_new = next_pts[st].reshape(-1, 2)
|
|
|
self.good_count = int(len(good_new))
|
|
|
|
|
|
if self.good_count < int(HYBRID_KLT_MIN_INLIERS):
|
|
|
return None, 0.0, "none"
|
|
|
|
|
|
# ── Путь А: affine через estimateAffinePartial2D (RANSAC) ──
|
|
|
if self.good_count >= int(HYBRID_KLT_MIN_FOR_AFFINE):
|
|
|
M, inlier_mask = cv2.estimateAffinePartial2D(
|
|
|
good_old, good_new,
|
|
|
method=cv2.RANSAC,
|
|
|
ransacReprojThreshold=float(HYBRID_KLT_OUTLIER_THR),
|
|
|
maxIters=200,
|
|
|
confidence=0.99,
|
|
|
)
|
|
|
if M is not None and inlier_mask is not None:
|
|
|
inlier_cnt = int(inlier_mask.sum())
|
|
|
if inlier_cnt >= int(HYBRID_KLT_MIN_INLIERS):
|
|
|
new_box = self._apply_affine_to_box(M, self.box)
|
|
|
if new_box is not None and self._is_sane_box_update(self.box, new_box):
|
|
|
# quality: fraction of inliers * абсолютное число
|
|
|
frac = float(inlier_cnt) / max(1, self.good_count)
|
|
|
quality = float(clamp(
|
|
|
frac * (inlier_cnt / max(6.0, HYBRID_KLT_MAX_CORNERS * 0.4)),
|
|
|
0.0, 1.0,
|
|
|
))
|
|
|
# Обновляем pts только inliers для следующего кадра
|
|
|
mask_flat = inlier_mask.reshape(-1).astype(bool)
|
|
|
self.pts = good_new[mask_flat].reshape(-1, 1, 2).astype(np.float32)
|
|
|
return new_box, quality, "klt-affine"
|
|
|
|
|
|
# ── Путь Б: fallback median translation ─────────────
|
|
|
disp = good_new - good_old
|
|
|
med = np.median(disp, axis=0)
|
|
|
d = np.linalg.norm(disp - med[None, :], axis=1)
|
|
|
inliers = d < float(HYBRID_KLT_OUTLIER_THR)
|
|
|
in_cnt = int(np.count_nonzero(inliers))
|
|
|
|
|
|
if in_cnt < int(HYBRID_KLT_MIN_INLIERS):
|
|
|
return None, 0.0, "none"
|
|
|
|
|
|
disp_in = disp[inliers]
|
|
|
med = np.median(disp_in, axis=0)
|
|
|
dx, dy = float(med[0]), float(med[1])
|
|
|
|
|
|
if abs(dx) > HYBRID_MAX_TRANSLATION_PX or abs(dy) > HYBRID_MAX_TRANSLATION_PX:
|
|
|
return None, 0.0, "none"
|
|
|
|
|
|
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) / max(1, self.good_count)
|
|
|
quality = float(clamp(
|
|
|
frac * (in_cnt / max(6.0, HYBRID_KLT_MAX_CORNERS * 0.5)),
|
|
|
0.0, 0.85, # translate даёт меньший макс quality чем affine
|
|
|
))
|
|
|
|
|
|
self.pts = good_new[inliers].reshape(-1, 1, 2).astype(np.float32)
|
|
|
return new_box, quality, "klt-translate"
|
|
|
|
|
|
def _apply_affine_to_box(self, M, box):
|
|
|
"""
|
|
|
Применяет аффинку ТОЛЬКО к центру bbox (translation + rotation).
|
|
|
РАЗМЕР bbox НЕ ОБНОВЛЯЕТСЯ — остаётся фиксированным.
|
|
|
|
|
|
Причина: когда KLT-трекер масштабировал бокс своим scale, он
|
|
|
постепенно расходился по размеру с тем, что детектирует YOLO,
|
|
|
в результате ByteTrack не ассоциировал детекты с треком
|
|
|
(hit_streak падал до 0). Размер бокса должен обновляться
|
|
|
только от YOLO — это единственный надёжный источник истинного
|
|
|
масштаба цели.
|
|
|
|
|
|
scale от аффинки используется только для sanity check: если
|
|
|
он слишком далёк от 1.0, значит RANSAC сбился, возвращаем None.
|
|
|
"""
|
|
|
x1, y1, x2, y2 = [float(v) for v in box]
|
|
|
cx = 0.5 * (x1 + x2)
|
|
|
cy = 0.5 * (y1 + y2)
|
|
|
w = x2 - x1
|
|
|
h = y2 - y1
|
|
|
|
|
|
# Центр через аффинку
|
|
|
new_cx = M[0, 0] * cx + M[0, 1] * cy + M[0, 2]
|
|
|
new_cy = M[1, 0] * cx + M[1, 1] * cy + M[1, 2]
|
|
|
|
|
|
# Sanity check через норму столбца (должна быть ~1.0 для
|
|
|
# кадра к кадру; большое отклонение = RANSAC сбой)
|
|
|
scale = float(np.hypot(M[0, 0], M[1, 0]))
|
|
|
if not (float(HYBRID_MIN_SCALE_PER_FRAME) <= scale <= float(HYBRID_MAX_SCALE_PER_FRAME)):
|
|
|
return None
|
|
|
|
|
|
# Размер НЕ меняем — используем старые w, h
|
|
|
return np.array(
|
|
|
[new_cx - w * 0.5, new_cy - h * 0.5,
|
|
|
new_cx + w * 0.5, new_cy + h * 0.5],
|
|
|
dtype=np.float32,
|
|
|
)
|
|
|
|
|
|
def _is_sane_box_update(self, old_box, new_box):
|
|
|
"""Защита от RANSAC-выбросов."""
|
|
|
ox1, oy1, ox2, oy2 = [float(v) for v in old_box]
|
|
|
nx1, ny1, nx2, ny2 = [float(v) for v in new_box]
|
|
|
ocx = 0.5 * (ox1 + ox2)
|
|
|
ocy = 0.5 * (oy1 + oy2)
|
|
|
ncx = 0.5 * (nx1 + nx2)
|
|
|
ncy = 0.5 * (ny1 + ny2)
|
|
|
shift = float(np.hypot(ncx - ocx, ncy - ocy))
|
|
|
if shift > float(HYBRID_MAX_TRANSLATION_PX):
|
|
|
return False
|
|
|
return True
|
|
|
|
|
|
# ─── Level 2: Phase correlation ──────────────────────────
|
|
|
def _try_phase_correlation(self, gray_eq):
|
|
|
if self.template is None or self.template.size == 0:
|
|
|
return None, 0.0
|
|
|
|
|
|
search, (sx, sy) = self._crop_search_window(gray_eq, self.box)
|
|
|
if search is None:
|
|
|
return None, 0.0
|
|
|
h_t, w_t = self.template.shape[:2]
|
|
|
h_s, w_s = search.shape[:2]
|
|
|
if h_s < h_t or w_s < w_t:
|
|
|
return None, 0.0
|
|
|
|
|
|
cy_s = h_s // 2
|
|
|
cx_s = w_s // 2
|
|
|
y0 = max(0, cy_s - h_t // 2)
|
|
|
x0 = max(0, cx_s - w_t // 2)
|
|
|
y1 = min(h_s, y0 + h_t)
|
|
|
x1 = min(w_s, x0 + w_t)
|
|
|
patch = search[y0:y1, x0:x1]
|
|
|
if patch.shape != self.template.shape:
|
|
|
patch = cv2.copyMakeBorder(
|
|
|
patch, 0, h_t - patch.shape[0], 0, w_t - patch.shape[1],
|
|
|
cv2.BORDER_REPLICATE,
|
|
|
)
|
|
|
|
|
|
try:
|
|
|
t_f = self.template.astype(np.float32)
|
|
|
p_f = patch.astype(np.float32)
|
|
|
(shift_x, shift_y), response = cv2.phaseCorrelate(t_f, p_f)
|
|
|
except cv2.error:
|
|
|
return None, 0.0
|
|
|
|
|
|
if response < float(HYBRID_PHASE_MIN_RESP):
|
|
|
return None, 0.0
|
|
|
if abs(shift_x) > HYBRID_MAX_TRANSLATION_PX or abs(shift_y) > HYBRID_MAX_TRANSLATION_PX:
|
|
|
return None, 0.0
|
|
|
|
|
|
x1b, y1b, x2b, y2b = self.box
|
|
|
new_box = np.array(
|
|
|
[x1b + shift_x, y1b + shift_y, x2b + shift_x, y2b + shift_y],
|
|
|
dtype=np.float32,
|
|
|
)
|
|
|
quality = float(clamp(response * 0.75, 0.0, 0.75))
|
|
|
return new_box, quality
|
|
|
|
|
|
# ─── Level 3: Multi-scale template matching ──────────────
|
|
|
def _try_template_match(self, gray_eq):
|
|
|
if self.template is None or self.template.size == 0:
|
|
|
return None, 0.0
|
|
|
|
|
|
search, (sx, sy) = self._crop_search_window(gray_eq, self.box)
|
|
|
if search is None:
|
|
|
return None, 0.0
|
|
|
|
|
|
best = None # (score, new_box)
|
|
|
|
|
|
for scale in HYBRID_TM_SCALES:
|
|
|
if scale == 1.0:
|
|
|
tmpl = self.template
|
|
|
else:
|
|
|
nh = max(4, int(self.template.shape[0] * scale))
|
|
|
nw = max(4, int(self.template.shape[1] * scale))
|
|
|
tmpl = cv2.resize(self.template, (nw, nh), interpolation=cv2.INTER_LINEAR)
|
|
|
|
|
|
h_t, w_t = tmpl.shape[:2]
|
|
|
if search.shape[0] < h_t or search.shape[1] < w_t:
|
|
|
continue
|
|
|
|
|
|
try:
|
|
|
result = cv2.matchTemplate(search, tmpl, cv2.TM_CCOEFF_NORMED)
|
|
|
min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
|
|
|
except cv2.error:
|
|
|
continue
|
|
|
|
|
|
if max_val < float(HYBRID_TM_MIN_SCORE):
|
|
|
continue
|
|
|
|
|
|
mx, my = max_loc
|
|
|
nx1 = sx + mx
|
|
|
ny1 = sy + my
|
|
|
new_box = np.array(
|
|
|
[nx1, ny1, nx1 + w_t, ny1 + h_t], dtype=np.float32
|
|
|
)
|
|
|
if best is None or max_val > best[0]:
|
|
|
best = (float(max_val), new_box)
|
|
|
|
|
|
if best is None:
|
|
|
return None, 0.0
|
|
|
|
|
|
score, new_box = best
|
|
|
quality = float(clamp(score * 0.72, 0.0, 0.72))
|
|
|
return new_box, quality
|
|
|
|
|
|
# ─── Helpers ─────────────────────────────────────────────
|
|
|
def _extract_features(self, gray_eq, box):
|
|
|
h, w = gray_eq.shape[:2]
|
|
|
x1, y1, x2, y2 = box
|
|
|
bw = max(1.0, x2 - x1)
|
|
|
bh = max(1.0, y2 - y1)
|
|
|
pad_x = bw * float(HYBRID_PAD_FACTOR)
|
|
|
pad_y = bh * float(HYBRID_PAD_FACTOR)
|
|
|
|
|
|
mx1 = int(clamp(x1 - pad_x, 0, w - 1))
|
|
|
my1 = int(clamp(y1 - pad_y, 0, h - 1))
|
|
|
mx2 = int(clamp(x2 + pad_x, 1, w))
|
|
|
my2 = int(clamp(y2 + pad_y, 1, h))
|
|
|
|
|
|
mask = np.zeros_like(gray_eq)
|
|
|
mask[my1:my2, mx1:mx2] = 255
|
|
|
|
|
|
pts = cv2.goodFeaturesToTrack(
|
|
|
gray_eq,
|
|
|
maxCorners=int(HYBRID_KLT_MAX_CORNERS),
|
|
|
qualityLevel=float(HYBRID_KLT_QUALITY),
|
|
|
minDistance=int(HYBRID_KLT_MIN_DIST),
|
|
|
mask=mask,
|
|
|
)
|
|
|
if pts is None or len(pts) < 6:
|
|
|
pts = cv2.goodFeaturesToTrack(
|
|
|
gray_eq,
|
|
|
maxCorners=int(HYBRID_KLT_MAX_CORNERS),
|
|
|
qualityLevel=float(HYBRID_KLT_QUALITY_RETRY),
|
|
|
minDistance=int(HYBRID_KLT_MIN_DIST),
|
|
|
mask=mask,
|
|
|
)
|
|
|
self.pts = pts.astype(np.float32) if pts is not None else None
|
|
|
|
|
|
def _update_template(self, gray_eq, box, reset=False):
|
|
|
h, w = gray_eq.shape[:2]
|
|
|
x1 = int(clamp(box[0], 0, w - 1))
|
|
|
y1 = int(clamp(box[1], 0, h - 1))
|
|
|
x2 = int(clamp(box[2], x1 + 2, w))
|
|
|
y2 = int(clamp(box[3], y1 + 2, h))
|
|
|
crop = gray_eq[y1:y2, x1:x2]
|
|
|
if crop.size == 0:
|
|
|
return
|
|
|
|
|
|
if reset or self.template is None or self.template.shape != crop.shape:
|
|
|
self.template = crop.copy()
|
|
|
else:
|
|
|
a = float(HYBRID_TEMPLATE_UPDATE_EMA)
|
|
|
self.template = (a * self.template.astype(np.float32)
|
|
|
+ (1.0 - a) * crop.astype(np.float32)).astype(np.uint8)
|
|
|
|
|
|
def _crop_search_window(self, gray_eq, box):
|
|
|
h, w = gray_eq.shape[:2]
|
|
|
cx = 0.5 * (float(box[0]) + float(box[2]))
|
|
|
cy = 0.5 * (float(box[1]) + float(box[3]))
|
|
|
bw = max(8.0, float(box[2] - box[0]))
|
|
|
bh = max(8.0, float(box[3] - box[1]))
|
|
|
|
|
|
sw = bw * float(HYBRID_SEARCH_FACTOR)
|
|
|
sh = bh * float(HYBRID_SEARCH_FACTOR)
|
|
|
|
|
|
sx = int(clamp(cx - sw * 0.5, 0, w - 1))
|
|
|
sy = int(clamp(cy - sh * 0.5, 0, h - 1))
|
|
|
ex = int(clamp(cx + sw * 0.5, sx + 2, w))
|
|
|
ey = int(clamp(cy + sh * 0.5, sy + 2, h))
|
|
|
|
|
|
search = gray_eq[sy:ey, sx:ex]
|
|
|
if search.size == 0:
|
|
|
return None, (0, 0)
|
|
|
return search, (sx, sy)
|