stabilization_utils.py

import cv2
import numpy as np
from skimage.transform import EuclideanTransform
from sklearn.neighbors import KernelDensity


IDENTITY_AFFINE = np.array([[1, 0, 0], [0, 1, 0]], dtype=np.float32)


def identity_transform_matrix():
    return IDENTITY_AFFINE.copy()


def _as_points(points):
    return np.asarray(points, dtype=np.float32).reshape(-1, 2)


def _as_homogeneous_matrix(matrix):
    matrix = np.asarray(matrix, dtype=np.float64)

    if matrix.shape == (3, 3):
        return matrix

    if matrix.shape == (2, 3):
        return np.vstack([matrix, [0.0, 0.0, 1.0]])

    raise ValueError(f'expected a 2x3 or 3x3 transform matrix, got {matrix.shape}')


def _as_cv2_affine_matrix(transform):
    matrix = transform.params if hasattr(transform, 'params') else transform
    return np.asarray(matrix, dtype=np.float32)[:2, :]


def build_transformation_matrix(transform):
    """Convert [dx, dy, da] to the 2x3 matrix expected by cv2.warpAffine."""
    dx, dy, da = np.asarray(transform, dtype=np.float64).reshape(3)
    transform = EuclideanTransform(rotation=da, translation=(dx, dy))

    return _as_cv2_affine_matrix(transform)


def motion_from_matrix(matrix):
    """Convert a 2x3/3x3 transform matrix back to [dx, dy, da]."""
    transform = EuclideanTransform(matrix=_as_homogeneous_matrix(matrix))
    dx, dy = transform.translation

    return np.array([dx, dy, transform.rotation], dtype=np.float32)


def update_transformation_matrix(current_matrix, delta_matrix):
    """Compose an accumulated transform with a new frame-to-frame transform."""
    current = EuclideanTransform(matrix=_as_homogeneous_matrix(current_matrix))
    delta = EuclideanTransform(matrix=_as_homogeneous_matrix(delta_matrix))
    composed = EuclideanTransform(matrix=delta.params @ current.params)

    return _as_cv2_affine_matrix(composed)


def estimate_partial_transform(
        matched_keypoints,
        return_confidence=False,
        ransac_reproj_threshold=3.0,
        max_iters=2000,
        confidence=0.995):
    """Estimate translation + rotation from matched keypoints.

    OpenCV does the robust RANSAC fit; scikit-image is used for the reusable
    transform representation and conversion back to motion parameters.
    """
    prev_matched_kp, cur_matched_kp = matched_keypoints
    prev_matched_kp = _as_points(prev_matched_kp)
    cur_matched_kp = _as_points(cur_matched_kp)

    if prev_matched_kp.shape[0] < 3 or cur_matched_kp.shape[0] < 3:
        motion = np.zeros(3, dtype=np.float32)
        return (motion, 0.0) if return_confidence else motion

    matrix, inliers = cv2.estimateAffinePartial2D(
        prev_matched_kp,
        cur_matched_kp,
        method=cv2.RANSAC,
        ransacReprojThreshold=ransac_reproj_threshold,
        maxIters=max_iters,
        confidence=confidence,
        refineIters=10)

    if matrix is None:
        motion = np.zeros(3, dtype=np.float32)
        inlier_ratio = 0.0
    else:
        motion = motion_from_matrix(matrix)
        inlier_ratio = float(np.mean(inliers)) if inliers is not None else 1.0

    return (motion, inlier_ratio) if return_confidence else motion


def check_dy_dx_da(dy, dx, da, d_max=20.0, d_min=-20.0):
    dy, dx, da = np.clip([dy, dx, da], d_min, d_max)

    return float(dy), float(dx), float(da)


def remove_motion_outliers(prev_pts, curr_pts, min_keep=8, low_density_quantile=0.15):
    prev_pts = _as_points(prev_pts)
    curr_pts = _as_points(curr_pts)

    if prev_pts.shape[0] != curr_pts.shape[0]:
        raise ValueError('previous and current point arrays must have the same length')

    if prev_pts.shape[0] < min_keep:
        return prev_pts, curr_pts

    distances = np.linalg.norm(prev_pts - curr_pts, axis=1)
    bandwidth = max(0.5, float(np.std(distances)) * 0.25)
    distances_2d = distances.reshape(-1, 1)
    kde = KernelDensity(kernel='gaussian', bandwidth=bandwidth).fit(distances_2d)
    density = np.exp(kde.score_samples(distances_2d))

    keep = density >= np.quantile(density, low_density_quantile)
    if np.count_nonzero(keep) < min_keep:
        return prev_pts, curr_pts

    return prev_pts[keep], curr_pts[keep]


removeOutliers = remove_motion_outliers