Initial commit

c02_pseudo_imu_rectifier.py

@@ -0,0 +1,59 @@
+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