mirror of
https://github.com/azaion/gps-denied-onboard.git
synced 2026-04-22 22:26:38 +00:00
Denys Zaitsev
59 lines
2.5 KiB
Python
59 lines
2.5 KiB
Python
import cv2
|
|
import numpy as np
|
|
import torch
|
|
import logging
|
|
from typing import Tuple
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class PseudoImuRectifier:
|
|
"""
|
|
Estimates the horizon/tilt of the UAV camera from a single monocular image
|
|
and rectifies it to a pseudo-nadir view to prevent tracking loss during sharp banks.
|
|
"""
|
|
def __init__(self, device: str = "cuda", tilt_threshold_deg: float = 15.0):
|
|
self.device = torch.device(device if torch.cuda.is_available() else "cpu")
|
|
self.tilt_threshold_deg = tilt_threshold_deg
|
|
|
|
logger.info(f"Initializing Pseudo-IMU Horizon Estimator on {self.device}")
|
|
# In a full deployment, this loads a lightweight CNN like HorizonNet or DepthAnythingV2
|
|
# self.horizon_model = load_horizon_model().to(self.device)
|
|
|
|
def estimate_attitude(self, image: np.ndarray) -> Tuple[float, float]:
|
|
"""
|
|
Estimates pitch and roll from the image's vanishing points/horizon.
|
|
Returns: (pitch_degrees, roll_degrees)
|
|
"""
|
|
# Placeholder for deep-learning based horizon estimation tensor operations.
|
|
# Returns mocked 0.0 for pitch/roll unless the model detects extreme banking.
|
|
pitch_deg = 0.0
|
|
roll_deg = 0.0
|
|
return pitch_deg, roll_deg
|
|
|
|
def compute_rectification_homography(self, pitch: float, roll: float, K: np.ndarray) -> np.ndarray:
|
|
"""Computes the homography matrix to un-warp perspective distortion."""
|
|
p = np.deg2rad(pitch)
|
|
r = np.deg2rad(roll)
|
|
|
|
# Rotation matrices for pitch (X-axis) and roll (Z-axis)
|
|
Rx = np.array([[1, 0, 0], [0, np.cos(p), -np.sin(p)], [0, np.sin(p), np.cos(p)]])
|
|
Rz = np.array([[np.cos(r), -np.sin(r), 0], [np.sin(r), np.cos(r), 0], [0, 0, 1]])
|
|
|
|
R = Rz @ Rx
|
|
|
|
# Homography: H = K * R * K^-1
|
|
K_inv = np.linalg.inv(K)
|
|
H = K @ R @ K_inv
|
|
return H
|
|
|
|
def rectify_image(self, image: np.ndarray, K: np.ndarray) -> Tuple[np.ndarray, bool]:
|
|
pitch, roll = self.estimate_attitude(image)
|
|
|
|
if abs(pitch) < self.tilt_threshold_deg and abs(roll) < self.tilt_threshold_deg:
|
|
return image, False # No rectification needed
|
|
|
|
logger.warning(f"Extreme tilt detected (Pitch: {pitch:.1f}, Roll: {roll:.1f}). Rectifying.")
|
|
H = self.compute_rectification_homography(-pitch, -roll, K)
|
|
|
|
rectified_image = cv2.warpPerspective(image, H, (image.shape[1], image.shape[0]), flags=cv2.INTER_LINEAR)
|
|
return rectified_image, True |