gps-denied-onboard/f13_coordinate_transformer.py

import math
import numpy as np
import logging
from typing import Tuple, List, Optional, Dict
from pydantic import BaseModel
from abc import ABC, abstractmethod

from f02_1_flight_lifecycle_manager import GPSPoint, CameraParameters
from f10_factor_graph_optimizer import Pose
from h01_camera_model import CameraModel
from h02_gsd_calculator import GSDCalculator

logger = logging.getLogger(__name__)

class OriginNotSetError(Exception):
    pass

class FlightConfig(BaseModel):
    camera_params: CameraParameters
    altitude: float
    
class ICoordinateTransformer(ABC):
    @abstractmethod
    def set_enu_origin(self, flight_id: str, origin_gps: GPSPoint) -> None: pass
    @abstractmethod
    def get_enu_origin(self, flight_id: str) -> GPSPoint: pass
    @abstractmethod
    def gps_to_enu(self, flight_id: str, gps: GPSPoint) -> Tuple[float, float, float]: pass
    @abstractmethod
    def enu_to_gps(self, flight_id: str, enu: Tuple[float, float, float]) -> GPSPoint: pass
    @abstractmethod
    def pixel_to_gps(self, flight_id: str, pixel: Tuple[float, float], frame_pose: Pose, camera_params: CameraParameters, altitude: float) -> GPSPoint: pass
    @abstractmethod
    def gps_to_pixel(self, flight_id: str, gps: GPSPoint, frame_pose: Pose, camera_params: CameraParameters, altitude: float) -> Tuple[float, float]: pass
    @abstractmethod
    def image_object_to_gps(self, flight_id: str, frame_id: int, object_pixel: Tuple[float, float]) -> GPSPoint: pass
    @abstractmethod
    def transform_points(self, points: List[Tuple[float, float]], transformation: np.ndarray) -> List[Tuple[float, float]]: pass
    @abstractmethod
    def calculate_meters_per_pixel(self, lat: float, zoom: int) -> float: pass
    @abstractmethod
    def calculate_haversine_distance(self, gps1: GPSPoint, gps2: GPSPoint) -> float: pass


class CoordinateTransformer(ICoordinateTransformer):
    """
    F13: Coordinate Transformer
    Provides geometric and geospatial coordinate mappings, relying on ground plane assumptions,
    camera intrinsics (H01), and the optimized Factor Graph trajectory (F10).
    """
    def __init__(self, f10_optimizer=None, f17_config=None, camera_model=None, gsd_calculator=None):
        self.f10 = f10_optimizer
        self.f17 = f17_config
        self.camera_model = camera_model or CameraModel()
        self.gsd_calculator = gsd_calculator or GSDCalculator()
        self._origins: Dict[str, GPSPoint] = {}

    # --- 13.01 ENU Coordinate Management ---

    def set_enu_origin(self, flight_id: str, origin_gps: GPSPoint) -> None:
        self._origins[flight_id] = origin_gps

    def get_enu_origin(self, flight_id: str) -> GPSPoint:
        if flight_id not in self._origins:
            raise OriginNotSetError(f"ENU Origin not set for flight {flight_id}")
        return self._origins[flight_id]

    def gps_to_enu(self, flight_id: str, gps: GPSPoint) -> Tuple[float, float, float]:
        origin = self.get_enu_origin(flight_id)
        delta_lat = gps.lat - origin.lat
        delta_lon = gps.lon - origin.lon
        east = delta_lon * math.cos(math.radians(origin.lat)) * 111319.5
        north = delta_lat * 111319.5
        return (east, north, 0.0)

    def enu_to_gps(self, flight_id: str, enu: Tuple[float, float, float]) -> GPSPoint:
        origin = self.get_enu_origin(flight_id)
        east, north, _ = enu
        delta_lat = north / 111319.5
        delta_lon = east / (math.cos(math.radians(origin.lat)) * 111319.5)
        return GPSPoint(lat=origin.lat + delta_lat, lon=origin.lon + delta_lon)

    # --- 13.02 Pixel-GPS Projection ---

    def _intersect_ray_ground_plane(self, ray_origin: np.ndarray, ray_direction: np.ndarray, ground_z: float = 0.0) -> np.ndarray:
        if abs(ray_direction[2]) < 1e-6:
            return ray_origin
        t = (ground_z - ray_origin[2]) / ray_direction[2]
        return ray_origin + t * ray_direction

    def pixel_to_gps(self, flight_id: str, pixel: Tuple[float, float], frame_pose: Pose, camera_params: CameraParameters, altitude: float) -> GPSPoint:
        ray_cam = self.camera_model.unproject(pixel, 1.0, camera_params)
        ray_enu_dir = frame_pose.orientation @ ray_cam
        
        # Origin of ray in ENU is the camera position. Using predefined altitude.
        ray_origin = np.copy(frame_pose.position)
        ray_origin[2] = altitude
        
        point_enu = self._intersect_ray_ground_plane(ray_origin, ray_enu_dir, 0.0)
        return self.enu_to_gps(flight_id, (point_enu[0], point_enu[1], point_enu[2]))

    def gps_to_pixel(self, flight_id: str, gps: GPSPoint, frame_pose: Pose, camera_params: CameraParameters, altitude: float) -> Tuple[float, float]:
        enu = self.gps_to_enu(flight_id, gps)
        point_enu = np.array(enu)
        # Transform ENU to Camera Frame
        point_cam = frame_pose.orientation.T @ (point_enu - frame_pose.position)
        return self.camera_model.project(point_cam, camera_params)

    def image_object_to_gps(self, flight_id: str, frame_id: int, object_pixel: Tuple[float, float]) -> GPSPoint:
        if not self.f10 or not self.f17:
            raise RuntimeError("Missing F10 or F17 dependencies for image_object_to_gps.")
        trajectory = self.f10.get_trajectory(flight_id)
        if frame_id not in trajectory:
            raise ValueError(f"Frame {frame_id} not found in optimized trajectory.")
            
        flight_config = self.f17.get_flight_config(flight_id)
        return self.pixel_to_gps(flight_id, object_pixel, trajectory[frame_id], flight_config.camera_params, flight_config.altitude)

    def transform_points(self, points: List[Tuple[float, float]], transformation: np.ndarray) -> List[Tuple[float, float]]:
        if not points: return []
        homog = np.hstack([np.array(points, dtype=np.float64), np.ones((len(points), 1))])
        trans = (transformation @ homog.T).T
        if transformation.shape == (3, 3):
            return [(p[0]/p[2], p[1]/p[2]) for p in trans]
        return [(p[0], p[1]) for p in trans]

    def calculate_meters_per_pixel(self, lat: float, zoom: int) -> float:
        return self.gsd_calculator.meters_per_pixel(lat, zoom)

    def calculate_haversine_distance(self, gps1: GPSPoint, gps2: GPSPoint) -> float:
        R = 6371000.0 # Earth radius in meters
        phi1 = math.radians(gps1.lat)
        phi2 = math.radians(gps2.lat)
        delta_phi = math.radians(gps2.lat - gps1.lat)
        delta_lambda = math.radians(gps2.lon - gps1.lon)
        a = math.sin(delta_phi / 2.0)**2 + math.cos(phi1) * math.cos(phi2) * math.sin(delta_lambda / 2.0)**2
        c = 2 * math.atan2(math.sqrt(a), math.sqrt(1 - a))
        return R * c