Initial commit

f11_failure_recovery_coordinator.py

@@ -0,0 +1,328 @@
+import time
+import logging
+import os
+from datetime import datetime
+from typing import List, Optional, Tuple, Dict, Any
+import numpy as np
+from pydantic import BaseModel, Field
+
+from f02_1_flight_lifecycle_manager import GPSPoint
+from f04_satellite_data_manager import TileCoords
+
+logger = logging.getLogger(__name__)
+
+# --- Data Models ---
+
+class RelativePose(BaseModel):
+    transform: np.ndarray
+    inlier_count: int
+    reprojection_error: float
+
+    model_config = {"arbitrary_types_allowed": True}
+
+class Sim3Transform(BaseModel):
+    translation: np.ndarray
+    rotation: np.ndarray
+    scale: float
+
+    model_config = {"arbitrary_types_allowed": True}
+
+class AlignmentResult(BaseModel):
+    matched: bool
+    gps: GPSPoint
+    confidence: float
+    inlier_count: int
+    transform: np.ndarray
+
+    model_config = {"arbitrary_types_allowed": True}
+
+class ConfidenceAssessment(BaseModel):
+    overall_confidence: float
+    vo_confidence: float
+    litesam_confidence: float
+    inlier_count: int
+    tracking_status: str  # "good", "degraded", "lost"
+
+class SearchSession(BaseModel):
+    session_id: str
+    flight_id: str
+    frame_id: int
+    center_gps: GPSPoint
+    current_grid_size: int
+    max_grid_size: int
+    found: bool
+    exhausted: bool
+
+class SearchStatus(BaseModel):
+    current_grid_size: int
+    found: bool
+    exhausted: bool
+
+class TileCandidate(BaseModel):
+    tile_id: str
+    score: float
+    gps: GPSPoint
+
+class UserInputRequest(BaseModel):
+    request_id: str
+    flight_id: str
+    frame_id: int
+    uav_image: Any = Field(exclude=True)
+    candidate_tiles: List[TileCandidate]
+    message: str
+    created_at: datetime
+
+    model_config = {"arbitrary_types_allowed": True}
+
+class UserAnchor(BaseModel):
+    uav_pixel: Tuple[float, float]
+    satellite_gps: GPSPoint
+    confidence: float = 1.0
+
+class ChunkHandle(BaseModel):
+    chunk_id: str
+    flight_id: str
+    start_frame_id: int = 0
+    end_frame_id: Optional[int] = None
+    frames: List[int] = []
+    is_active: bool = True
+    has_anchor: bool = False
+    anchor_frame_id: Optional[int] = None
+    anchor_gps: Optional[GPSPoint] = None
+    matching_status: str = "unanchored"  # "unanchored", "matching", "anchored", "merged"
+
+class ChunkAlignmentResult(BaseModel):
+    matched: bool
+    chunk_id: str
+    chunk_center_gps: GPSPoint
+    rotation_angle: float
+    confidence: float
+    inlier_count: int
+    transform: Sim3Transform
+
+class RecoveryStatus(BaseModel):
+    success: bool
+    method: str
+    gps: Optional[GPSPoint]
+    chunk_id: Optional[str]
+    message: Optional[str]
+
+# --- Implementation ---
+
+class FailureRecoveryCoordinator:
+    """
+    Coordinates failure recovery strategies (progressive search, chunk matching, user input).
+    Pure logic component: decides what to do, delegates execution to dependencies.
+    """
+    def __init__(self, deps: Dict[str, Any]):
+        # Dependencies injected via constructor dictionary to prevent circular imports
+        self.f04 = deps.get("satellite_data_manager")
+        self.f06 = deps.get("image_rotation_manager")
+        self.f08 = deps.get("global_place_recognition")
+        self.f09 = deps.get("metric_refinement")
+        self.f10 = deps.get("factor_graph_optimizer")
+        self.f12 = deps.get("route_chunk_manager")
+
+    # --- Status Checks ---
+
+    def check_confidence(self, vo_result: RelativePose, litesam_result: Optional[AlignmentResult]) -> ConfidenceAssessment:
+        inliers = vo_result.inlier_count if vo_result else 0
+        
+        if inliers > 50:
+            status = "good"
+            vo_conf = 1.0
+        elif inliers >= 20:
+            status = "degraded"
+            vo_conf = inliers / 50.0
+        else:
+            status = "lost"
+            vo_conf = 0.0
+            
+        ls_conf = litesam_result.confidence if litesam_result else 0.0
+        overall = max(vo_conf, ls_conf)
+        
+        return ConfidenceAssessment(
+            overall_confidence=overall,
+            vo_confidence=vo_conf,
+            litesam_confidence=ls_conf,
+            inlier_count=inliers,
+            tracking_status=status
+        )
+
+    def detect_tracking_loss(self, confidence: ConfidenceAssessment) -> bool:
+        return confidence.tracking_status == "lost"
+
+    # --- Search & Recovery ---
+
+    def start_search(self, flight_id: str, frame_id: int, estimated_gps: GPSPoint) -> SearchSession:
+        logger.info(f"Starting progressive search for flight {flight_id}, frame {frame_id} at {estimated_gps}")
+        return SearchSession(
+            session_id=f"search_{flight_id}_{frame_id}",
+            flight_id=flight_id,
+            frame_id=frame_id,
+            center_gps=estimated_gps,
+            current_grid_size=1,
+            max_grid_size=25,
+            found=False,
+            exhausted=False
+        )
+
+    def expand_search_radius(self, session: SearchSession) -> List[TileCoords]:
+        grid_progression = [1, 4, 9, 16, 25]
+        try:
+            idx = grid_progression.index(session.current_grid_size)
+            if idx + 1 < len(grid_progression):
+                new_size = grid_progression[idx + 1]
+                
+                # Mocking tile expansion assuming F04 has compute_tile_coords
+                center_tc = self.f04.compute_tile_coords(session.center_gps.lat, session.center_gps.lon, zoom=18)
+                new_tiles = self.f04.expand_search_grid(center_tc, session.current_grid_size, new_size)
+                
+                session.current_grid_size = new_size
+                return new_tiles
+        except ValueError:
+            pass
+            
+        session.exhausted = True
+        return []
+
+    def try_current_grid(self, session: SearchSession, tiles: Dict[str, Tuple[np.ndarray, Any]], uav_image: np.ndarray) -> Optional[AlignmentResult]:
+        if os.environ.get("USE_MOCK_MODELS") == "1":
+            # Fake a successful satellite recovery to keep the simulation moving automatically
+            mock_res = AlignmentResult(
+                matched=True,
+                gps=session.center_gps,
+                confidence=0.9,
+                inlier_count=100,
+                transform=np.eye(3)
+            )
+            self.mark_found(session, mock_res)
+            return mock_res
+            
+        for tile_id, (tile_img, bounds) in tiles.items():
+            if self.f09:
+                result = self.f09.align_to_satellite(uav_image, tile_img, bounds)
+            if result and result.confidence > 0.7:
+                self.mark_found(session, result)
+                return result
+        return None
+
+    def mark_found(self, session: SearchSession, result: AlignmentResult) -> bool:
+        session.found = True
+        logger.info(f"Search session {session.session_id} succeeded at grid size {session.current_grid_size}.")
+        return True
+
+    def get_search_status(self, session: SearchSession) -> SearchStatus:
+        return SearchStatus(
+            current_grid_size=session.current_grid_size,
+            found=session.found,
+            exhausted=session.exhausted
+        )
+
+    def create_user_input_request(self, flight_id: str, frame_id: int, uav_image: np.ndarray, candidate_tiles: List[TileCandidate]) -> UserInputRequest:
+        return UserInputRequest(
+            request_id=f"usr_req_{flight_id}_{frame_id}",
+            flight_id=flight_id,
+            frame_id=frame_id,
+            uav_image=uav_image,
+            candidate_tiles=candidate_tiles,
+            message="Tracking lost and automatic recovery failed. Please provide a location anchor.",
+            created_at=datetime.utcnow()
+        )
+
+    def apply_user_anchor(self, flight_id: str, frame_id: int, anchor: UserAnchor) -> bool:
+        logger.info(f"Applying user anchor for frame {frame_id} at {anchor.satellite_gps}")
+        gps_array = np.array([anchor.satellite_gps.lat, anchor.satellite_gps.lon, 400.0]) # Defaulting alt
+        
+        # Delegate to Factor Graph Optimizer to add hard constraint
+        # Note: In a real integration, we'd need to find which chunk this frame belongs to
+        chunk_id = self.f12.get_chunk_for_frame(flight_id, frame_id)
+        if chunk_id:
+            self.f10.add_chunk_anchor(chunk_id, frame_id, gps_array)
+            return True
+        return False
+
+    # --- Chunk Recovery ---
+
+    def create_chunk_on_tracking_loss(self, flight_id: str, frame_id: int) -> ChunkHandle:
+        logger.warning(f"Creating proactive recovery chunk starting at frame {frame_id}")
+        chunk = self.f12.create_chunk(flight_id, frame_id)
+        return chunk
+
+    def try_chunk_semantic_matching(self, chunk_id: str) -> Optional[List[TileCandidate]]:
+        """Attempts semantic matching for a whole chunk using aggregate descriptors."""
+        logger.info(f"Attempting semantic matching for chunk {chunk_id}.")
+        if not hasattr(self.f12, 'get_chunk_images'):
+            return None
+            
+        chunk_images = self.f12.get_chunk_images(chunk_id)
+        if not chunk_images:
+            return None
+            
+        candidates = self.f08.retrieve_candidate_tiles_for_chunk(chunk_images)
+        return candidates if candidates else None
+
+    def try_chunk_litesam_matching(self, chunk_id: str, candidate_tiles: List[TileCandidate]) -> Optional[ChunkAlignmentResult]:
+        """Attempts LiteSAM matching across candidate tiles with rotation sweeps."""
+        logger.info(f"Attempting LiteSAM rotation sweeps for chunk {chunk_id}")
+        
+        if not hasattr(self.f12, 'get_chunk_images'):
+            return None
+            
+        chunk_images = self.f12.get_chunk_images(chunk_id)
+        if not chunk_images:
+            return None
+            
+        for candidate in candidate_tiles:
+            tile_img = self.f04.fetch_tile(candidate.gps.lat, candidate.gps.lon, zoom=18)
+            if tile_img is not None:
+                coords = self.f04.compute_tile_coords(candidate.gps.lat, candidate.gps.lon, zoom=18)
+                bounds = self.f04.compute_tile_bounds(coords)
+                
+                rot_result = self.f06.try_chunk_rotation_steps(chunk_images, tile_img, bounds, self.f09)
+                if rot_result and rot_result.matched:
+                    sim3 = self.f09._compute_sim3_transform(rot_result.homography, bounds) if hasattr(self.f09, '_compute_sim3_transform') else Sim3Transform(translation=np.zeros(3), rotation=np.eye(3), scale=1.0)
+                    gps = self.f09._get_chunk_center_gps(rot_result.homography, bounds, chunk_images) if hasattr(self.f09, '_get_chunk_center_gps') else candidate.gps
+                    return ChunkAlignmentResult(
+                        matched=True,
+                        chunk_id=chunk_id,
+                        chunk_center_gps=gps,
+                        rotation_angle=rot_result.precise_angle,
+                        confidence=rot_result.confidence,
+                        inlier_count=rot_result.inlier_count,
+                        transform=sim3,
+                        reprojection_error=0.0
+                    )
+        return None
+
+    def merge_chunk_to_trajectory(self, flight_id: str, chunk_id: str, alignment_result: ChunkAlignmentResult) -> bool:
+        logger.info(f"Merging chunk {chunk_id} to global trajectory.")
+        
+        main_chunk_id = self.f12.get_preceding_chunk(flight_id, chunk_id) if hasattr(self.f12, 'get_preceding_chunk') else "main"
+        if not main_chunk_id:
+            main_chunk_id = "main"
+            
+        anchor_gps = np.array([alignment_result.chunk_center_gps.lat, alignment_result.chunk_center_gps.lon, 400.0])
+        if hasattr(self.f12, 'mark_chunk_anchored'):
+            self.f12.mark_chunk_anchored(chunk_id, anchor_gps)
+        
+        if hasattr(self.f12, 'merge_chunks'):
+            return self.f12.merge_chunks(main_chunk_id, chunk_id, alignment_result.transform)
+        return False
+
+    def process_unanchored_chunks(self, flight_id: str) -> None:
+        """Background task loop structure designed to be called by a worker thread."""
+        if not hasattr(self.f12, 'get_chunks_for_matching'):
+            return
+            
+        unanchored_chunks = self.f12.get_chunks_for_matching(flight_id)
+        for chunk in unanchored_chunks:
+            if hasattr(self.f12, 'is_chunk_ready_for_matching') and self.f12.is_chunk_ready_for_matching(chunk.chunk_id):
+                if hasattr(self.f12, 'mark_chunk_matching'):
+                    self.f12.mark_chunk_matching(chunk.chunk_id)
+                
+                candidates = self.try_chunk_semantic_matching(chunk.chunk_id)
+                if candidates:
+                    alignment = self.try_chunk_litesam_matching(chunk.chunk_id, candidates)
+                    if alignment:
+                        self.merge_chunk_to_trajectory(flight_id, chunk.chunk_id, alignment)