You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

285 lines
9.9 KiB
Python

This file contains ambiguous Unicode characters!

This file contains ambiguous Unicode characters that may be confused with others in your current locale. If your use case is intentional and legitimate, you can safely ignore this warning. Use the Escape button to highlight these characters.

import cv2
import numpy as np
from config import *
from helpers import clamp
# Camera motion compensation
# =========================
# Поддерживает три режима оценки ego-motion:
# - affine (2x3, 4 DoF через estimateAffinePartial2D)
# - homography (3x3, 8 DoF через findHomography)
# - adaptive: пытаемся homography, при плохом качестве падаем на affine
#
# Основные функции:
# estimate_global_affine_ex — старая добрая affine + outliers
# estimate_global_homography_ex — homography + outliers + meta качества
# estimate_ego_motion_adaptive — автоматический выбор
# =========================
# Пороги качества homography
HOMO_MIN_INLIERS = 50
HOMO_MIN_FILLED_CELLS = 4 # из 3x3 = 9 ячеек
HOMO_GRID_SIZE = 3
HOMO_MAX_CORNER_SHIFT = 120.0 # px — защита от вырожденных H
def _rect_mask_exclude(gray, rects_norm):
h, w = gray.shape[:2]
mask = np.ones((h, w), dtype=np.uint8) * 255
for (x1n, y1n, x2n, y2n) in rects_norm:
x1 = int(clamp(x1n * w, 0, w - 1))
y1 = int(clamp(y1n * h, 0, h - 1))
x2 = int(clamp(x2n * w, 0, w))
y2 = int(clamp(y2n * h, 0, h))
cv2.rectangle(mask, (x1, y1), (x2, y2), 0, -1)
return mask
def _lk_correspondences(prev_gray, gray):
"""Общий блок: feature detection + LK tracking. (p0, p1) или (None, None)."""
if prev_gray is None or gray is None:
return None, None
mask = _rect_mask_exclude(prev_gray, CAM_MOTION_EXCLUDE_NORM) if CAM_MOTION_EXCLUDE_NORM else None
pts = cv2.goodFeaturesToTrack(
prev_gray,
maxCorners=int(CAM_MOTION_MAX_CORNERS),
qualityLevel=float(CAM_MOTION_QL),
minDistance=int(CAM_MOTION_MIN_DIST),
mask=mask,
)
if pts is None or len(pts) < 25:
return None, None
nxt, st, err = cv2.calcOpticalFlowPyrLK(
prev_gray, gray, pts, None, winSize=(21, 21), maxLevel=3
)
if nxt is None or st is None:
return None, None
st = st.reshape(-1).astype(bool)
p0 = pts[st].reshape(-1, 2).astype(np.float32)
p1 = nxt[st].reshape(-1, 2).astype(np.float32)
if len(p0) < 25:
return None, None
return p0, p1
def _classify_outliers(p0, p1, inlier_mask):
"""
Разбивает точки текущего кадра (p1) на inliers/outliers.
Отфильтровывает outliers с слишком малым/большим потоком (LK шум).
"""
inlier_mask = inlier_mask.reshape(-1).astype(bool)
inliers_curr = p1[inlier_mask]
outliers_curr = p1[~inlier_mask]
if len(outliers_curr) == 0:
return inliers_curr, outliers_curr
flow = p1[~inlier_mask] - p0[~inlier_mask]
mag = np.linalg.norm(flow, axis=1)
if len(inliers_curr) > 0:
ego_flow = p1[inlier_mask] - p0[inlier_mask]
ego_med = float(np.linalg.norm(np.median(ego_flow, axis=0)))
else:
ego_med = 0.0
valid = (mag > max(1.5, ego_med * 1.8)) & (mag < 80.0)
outliers_curr = outliers_curr[valid]
return inliers_curr, outliers_curr
# ─────────────────────────────────────────────────────────────
# AFFINE
# ─────────────────────────────────────────────────────────────
def estimate_global_affine_ex(prev_gray, gray):
empty = (None, np.zeros((0, 2), dtype=np.float32), np.zeros((0, 2), dtype=np.float32))
p0, p1 = _lk_correspondences(prev_gray, gray)
if p0 is None:
return empty
A, inlier_mask = cv2.estimateAffinePartial2D(
p0, p1,
method=cv2.RANSAC,
ransacReprojThreshold=float(CAM_MOTION_RANSAC_THR),
)
if A is None or inlier_mask is None:
return (A, np.zeros((0, 2), dtype=np.float32), p1)
inliers_curr, outliers_curr = _classify_outliers(p0, p1, inlier_mask)
return A, inliers_curr, outliers_curr
def estimate_global_affine(prev_gray, gray):
"""Backward-compatible wrapper."""
A, _, _ = estimate_global_affine_ex(prev_gray, gray)
return A
# ─────────────────────────────────────────────────────────────
# HOMOGRAPHY
# ─────────────────────────────────────────────────────────────
def estimate_global_homography_ex(prev_gray, gray):
"""
Homography 3x3. Возвращает (H, inliers_curr, outliers_curr, meta).
meta: n_inliers, filled_cells, spatial_ok, corner_shift_ok
"""
empty_meta = {
"n_inliers": 0,
"filled_cells": 0,
"spatial_ok": False,
"corner_shift_ok": False,
}
empty = (
None,
np.zeros((0, 2), dtype=np.float32),
np.zeros((0, 2), dtype=np.float32),
empty_meta,
)
p0, p1 = _lk_correspondences(prev_gray, gray)
if p0 is None:
return empty
H, inlier_mask = cv2.findHomography(
p0, p1,
method=cv2.RANSAC,
ransacReprojThreshold=float(CAM_MOTION_RANSAC_THR),
maxIters=2000,
confidence=0.995,
)
if H is None or inlier_mask is None:
return empty
inlier_mask_flat = inlier_mask.reshape(-1).astype(bool)
n_inliers = int(inlier_mask_flat.sum())
# Распределение inliers по 3x3 сетке
h, w = gray.shape[:2]
gs = int(HOMO_GRID_SIZE)
cells = np.zeros((gs, gs), dtype=np.uint8)
if n_inliers > 0:
inl = p1[inlier_mask_flat]
gx = np.clip((inl[:, 0] / w * gs).astype(int), 0, gs - 1)
gy = np.clip((inl[:, 1] / h * gs).astype(int), 0, gs - 1)
for x, y in zip(gx, gy):
cells[y, x] = 1
filled_cells = int(cells.sum())
spatial_ok = filled_cells >= int(HOMO_MIN_FILLED_CELLS)
# Проверка вырождения: сдвиг углов кадра
corners = np.array(
[[0, 0], [w - 1, 0], [w - 1, h - 1], [0, h - 1]], dtype=np.float32
).reshape(-1, 1, 2)
try:
warped_corners = cv2.perspectiveTransform(corners, H).reshape(-1, 2)
shifts = np.linalg.norm(warped_corners - corners.reshape(-1, 2), axis=1)
corner_shift_ok = bool(np.max(shifts) < float(HOMO_MAX_CORNER_SHIFT))
except cv2.error:
corner_shift_ok = False
meta = {
"n_inliers": n_inliers,
"filled_cells": filled_cells,
"spatial_ok": spatial_ok,
"corner_shift_ok": corner_shift_ok,
}
inliers_curr, outliers_curr = _classify_outliers(p0, p1, inlier_mask_flat)
return H, inliers_curr, outliers_curr, meta
def homography_is_plausible(H, meta):
"""Проходит ли homography пороги качества."""
if H is None:
return False
if meta.get("n_inliers", 0) < int(HOMO_MIN_INLIERS):
return False
if not meta.get("spatial_ok", False):
return False
if not meta.get("corner_shift_ok", False):
return False
try:
det = float(np.linalg.det(H[:2, :2]))
except Exception:
return False
if not (0.3 < abs(det) < 3.0):
return False
return True
# ─────────────────────────────────────────────────────────────
# ADAPTIVE SELECTION
# ─────────────────────────────────────────────────────────────
def estimate_ego_motion_adaptive(prev_gray, gray):
"""
Сначала homography, при плохом качестве — affine.
Returns:
M — 3x3 (homography) / 2x3 (affine) / None
kind — "homography" | "affine" | "none"
inliers — Nx2 точки статичного мира
outliers — Nx2 точки независимо движущихся объектов
"""
H, h_in, h_out, meta = estimate_global_homography_ex(prev_gray, gray)
if homography_is_plausible(H, meta):
return H, "homography", h_in, h_out
A, a_in, a_out = estimate_global_affine_ex(prev_gray, gray)
if A is not None and affine_is_plausible(A):
return A, "affine", a_in, a_out
empty = np.zeros((0, 2), dtype=np.float32)
return None, "none", empty, empty
# ─────────────────────────────────────────────────────────────
# Utilities / backward compat
# ─────────────────────────────────────────────────────────────
def apply_affine_to_point(A, x, y):
return (
A[0, 0] * x + A[0, 1] * y + A[0, 2],
A[1, 0] * x + A[1, 1] * y + A[1, 2],
)
def affine_is_plausible(A):
if A is None:
return False
tx = float(A[0, 2])
ty = float(A[1, 2])
shift = float(np.hypot(tx, ty))
if shift > float(CAM_MOTION_MAX_SHIFT_PX):
return False
a00 = float(A[0, 0])
a10 = float(A[1, 0])
col0_norm = float(np.hypot(a00, a10))
if col0_norm < 1e-6:
return False
if not (float(CAM_MOTION_SCALE_MIN) <= col0_norm <= float(CAM_MOTION_SCALE_MAX)):
return False
ang = float(np.degrees(np.arctan2(a10, a00)))
if abs(ang) > float(CAM_MOTION_MAX_ROT_DEG):
return False
return True
def warp_is_plausible(M, kind):
if kind == "homography":
return M is not None
if kind == "affine":
return affine_is_plausible(M)
return False