88 lines
3.5 KiB
Python
88 lines
3.5 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Dict, Optional
|
|
|
|
import numpy as np
|
|
|
|
|
|
@dataclass
|
|
class LocalizationFrameState:
|
|
odom_to_map: np.ndarray = field(default_factory=lambda: np.eye(4, dtype=np.float64))
|
|
odom_to_ref: np.ndarray = field(default_factory=lambda: np.eye(4, dtype=np.float64))
|
|
last_alignment: np.ndarray = field(default_factory=lambda: np.eye(4, dtype=np.float64))
|
|
confidence: float = 0.0
|
|
ref_valid: bool = False
|
|
|
|
def reset(self) -> None:
|
|
self.odom_to_map = np.eye(4, dtype=np.float64)
|
|
self.odom_to_ref = np.eye(4, dtype=np.float64)
|
|
self.last_alignment = np.eye(4, dtype=np.float64)
|
|
self.confidence = 0.0
|
|
self.ref_valid = False
|
|
|
|
def reset_reference(self) -> None:
|
|
self.odom_to_ref = np.eye(4, dtype=np.float64)
|
|
self.last_alignment = np.eye(4, dtype=np.float64)
|
|
self.confidence = 0.0
|
|
self.ref_valid = False
|
|
|
|
def set_reference_alignment(self, odom_to_ref: np.ndarray, confidence: float) -> None:
|
|
tf = np.asarray(odom_to_ref, dtype=np.float64)
|
|
if tf.shape != (4, 4):
|
|
return
|
|
self.odom_to_ref = np.array(tf, dtype=np.float64, copy=True)
|
|
self.last_alignment = np.array(tf, dtype=np.float64, copy=True)
|
|
self.confidence = float(np.clip(confidence, 0.0, 1.0))
|
|
self.ref_valid = True
|
|
|
|
def update_confidence(self, confidence: float) -> None:
|
|
self.confidence = float(np.clip(confidence, 0.0, 1.0))
|
|
|
|
def invalidate_reference(self) -> None:
|
|
self.ref_valid = False
|
|
self.confidence = 0.0
|
|
|
|
def apply_map_correction(self, correction: np.ndarray) -> None:
|
|
corr = np.asarray(correction, dtype=np.float64)
|
|
if corr.shape != (4, 4):
|
|
return
|
|
self.odom_to_map = corr @ np.asarray(self.odom_to_map, dtype=np.float64)
|
|
|
|
def apply_loop_correction(self, correction: np.ndarray, update_reference: bool = True) -> None:
|
|
corr = np.asarray(correction, dtype=np.float64)
|
|
if corr.shape != (4, 4):
|
|
return
|
|
self.odom_to_map = corr @ np.asarray(self.odom_to_map, dtype=np.float64)
|
|
if update_reference and self.ref_valid:
|
|
try:
|
|
inv_corr = np.linalg.inv(corr)
|
|
self.odom_to_ref = np.asarray(self.odom_to_ref, dtype=np.float64) @ inv_corr
|
|
self.last_alignment = np.array(self.odom_to_ref, dtype=np.float64, copy=True)
|
|
except Exception:
|
|
pass
|
|
|
|
def map_pose(self, odom_pose: Optional[np.ndarray]) -> Optional[np.ndarray]:
|
|
if odom_pose is None:
|
|
return None
|
|
pose = np.asarray(odom_pose, dtype=np.float64)
|
|
if pose.shape != (4, 4):
|
|
return None
|
|
return np.asarray(self.odom_to_map, dtype=np.float64) @ pose
|
|
|
|
def ref_pose(self, odom_pose: Optional[np.ndarray]) -> Optional[np.ndarray]:
|
|
if odom_pose is None or not self.ref_valid:
|
|
return None
|
|
pose = np.asarray(odom_pose, dtype=np.float64)
|
|
if pose.shape != (4, 4):
|
|
return None
|
|
return np.asarray(self.odom_to_ref, dtype=np.float64) @ pose
|
|
|
|
def snapshot(self) -> Dict[str, Any]:
|
|
return {
|
|
"ref_valid": bool(self.ref_valid),
|
|
"confidence": float(self.confidence),
|
|
"odom_to_map_t": np.asarray(self.odom_to_map[:3, 3], dtype=np.float64).tolist(),
|
|
"odom_to_ref_t": np.asarray(self.odom_to_ref[:3, 3], dtype=np.float64).tolist(),
|
|
}
|