import json import math import os import time import contextlib import io import warnings from dataclasses import asdict, dataclass from glob import glob from typing import Dict, List, Optional, Sequence, Tuple import cv2 import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np import torch from scipy.optimize import least_squares from scipy.sparse import lil_matrix from scipy.spatial.transform import Rotation from horizonstream.data import HorizonStreamDataLoader, HorizonStreamSequenceInfo from horizonstream.eval.io import frame_stems, read_opencv_camera_yml, read_pred_w2c_txt from horizonstream.eval.metrics import ate_rmse, transform_points from horizonstream.loop.pypose_optimizer import optimize_pose_graph_pypose from horizonstream.streaming.keyframe_selector import KeyframeSelector @dataclass class LoopConfig: verbose: bool = True pose_graph_backend: str = "pypose" pose_graph_model: str = "se3" pose_graph_update_mode: str = "all" pose_graph_solver_verbose: bool = True keyframe_stride: int = 1 min_keyframe_gap: int = 10 dbow_temporal_exclusion: int = 10 loop_chunk_size: int = 20 retrieval_top_k: int = 5 dbow_num_repeat: int = 3 retrieval_score_thresh_dbow: float = 0.034 retrieval_score_thresh_salad: float = 0.85 nms_radius: int = 25 max_candidates_per_method: int = 1000 max_verified_loops_per_method: int = 1000 dbow_vocab_path: Optional[str] = None salad_ckpt_path: str = "weights/dino_salad.ckpt" salad_dino_weights_path: str = "weights/dinov2_vitb14_pretrain.pth" salad_backbone: str = "dinov2_vitb14" salad_image_size: Tuple[int, int] = (336, 336) salad_batch_size: int = 32 orb_features: int = 4096 patch_match_thresh: float = 0.40 max_pair_matches: int = 256 orb_ratio_test: float = 0.80 rigid_ransac_iters: int = 1 rigid_inlier_thresh: float = 0.1 rigid_min_inliers: int = 24 loop_edge_min_separation: int = 30 sim3_irls_delta: float = 0.1 sim3_irls_max_iters: int = 5 sim3_irls_tol: float = 1e-9 pose_graph_rot_weight: float = 1.0 pose_graph_trans_weight: float = 2.0 pose_graph_scale_weight: float = 1.0 pose_graph_loop_weight: float = 0.1 pose_graph_max_nfev: int = 200 pose_graph_max_iterations: int = 30 pose_graph_lambda_init: float = 1e-6 def _log(msg: str, enabled: bool = True) -> None: if not enabled: return stamp = time.strftime("%H:%M:%S") print(f"[loop_runtime {stamp}] {msg}", flush=True) def _log_every( label: str, index: int, total: int, enabled: bool = True, step: int = 10, ) -> None: if not enabled: return if total <= 0: return if index == 1 or index == total or index % max(1, step) == 0: _log(f"{label}: {index}/{total}", enabled=enabled) @dataclass class SequenceArtifacts: seq_dir: str pose_variant: str pose_path: str intri_path: str image_paths: List[str] depth_paths: List[str] frame_ids: List[int] w2c: np.ndarray c2w: np.ndarray intrinsics: np.ndarray image_hw: Tuple[int, int] @dataclass class LoopCandidate: src_pos: int dst_pos: int src_frame: int dst_frame: int score: float method: str @dataclass class LoopEdge: src_pos: int dst_pos: int src_frame: int dst_frame: int score: float inliers: int method: str transform_ji: np.ndarray @dataclass class PoseGraphSummary: method: str num_candidates: int num_verified_loops: int mean_loop_score: float mean_loop_inliers: float trajectory_path: str loop_plot_path: str @dataclass class GroundTruthTrajectory: seq_name: str scene_root: str camera: Optional[str] c2w: np.ndarray valid_mask: np.ndarray def _read_intri_txt(path: str) -> Tuple[List[int], np.ndarray]: frames = [] intri = [] with open(path, "r") as f: for line in f: line = line.strip() if not line or line.startswith("#"): continue vals = [float(x) for x in line.split()] if len(vals) != 5: continue frame = int(vals[0]) K = np.eye(3, dtype=np.float64) K[0, 0] = vals[1] K[1, 1] = vals[2] K[0, 2] = vals[3] K[1, 2] = vals[4] frames.append(frame) intri.append(K) return frames, np.asarray(intri, dtype=np.float64) def _sorted_frame_files(path_pattern: str) -> List[str]: def frame_key(path: str) -> Tuple[int, str]: stem = os.path.splitext(os.path.basename(path))[0] digits = "".join(ch for ch in stem if ch.isdigit()) return (int(digits) if digits else -1, path) return sorted(glob(path_pattern), key=frame_key) def _resolve_pose_and_intri_paths( pose_dir: str, pose_variant: str = "main", ) -> Tuple[str, str, str]: variant = str(pose_variant or "main").strip().lower() if variant in {"", "main", "default"}: variant = "main" pose_path = os.path.join(pose_dir, "abs_pose.txt") intri_path = os.path.join(pose_dir, "intri.txt") else: pose_path = os.path.join(pose_dir, f"{variant}_abs_pose.txt") intri_path = os.path.join(pose_dir, f"{variant}_intri.txt") if not os.path.exists(pose_path): raise FileNotFoundError( f"Pose file not found for pose_variant='{variant}': {pose_path}" ) if not os.path.exists(intri_path): fallback_intri = os.path.join(pose_dir, "intri.txt") if os.path.exists(fallback_intri): intri_path = fallback_intri else: raise FileNotFoundError( f"Intrinsic file not found for pose_variant='{variant}': {intri_path}" ) return variant, pose_path, intri_path def load_sequence_artifacts( seq_dir: str, image_dir: Optional[str] = None, depth_dir: Optional[str] = None, pose_variant: str = "main", ) -> SequenceArtifacts: pose_dir = os.path.join(seq_dir, "poses") image_dir = image_dir or os.path.join(seq_dir, "images", "rgb") depth_dir = depth_dir or os.path.join(seq_dir, "depth", "dpt") pose_variant, pose_path, intri_path = _resolve_pose_and_intri_paths( pose_dir, pose_variant=pose_variant, ) frame_ids, w2c_list = read_pred_w2c_txt(pose_path) intri_frames, intri = _read_intri_txt(intri_path) if frame_ids != intri_frames: raise ValueError("Frame ids in abs_pose.txt and intri.txt do not match") image_paths = _sorted_frame_files(os.path.join(image_dir, "frame_*.*")) depth_paths = _sorted_frame_files(os.path.join(depth_dir, "frame_*.npy")) if len(image_paths) != len(frame_ids): raise ValueError( f"Expected {len(frame_ids)} RGB frames in {image_dir}, found {len(image_paths)}" ) if len(depth_paths) != len(frame_ids): raise ValueError( f"Expected {len(frame_ids)} depth frames in {depth_dir}, found {len(depth_paths)}" ) sample = cv2.imread(image_paths[0], cv2.IMREAD_COLOR) if sample is None: raise FileNotFoundError(image_paths[0]) h, w = sample.shape[:2] w2c = np.asarray(w2c_list, dtype=np.float64) c2w = np.linalg.inv(w2c) return SequenceArtifacts( seq_dir=seq_dir, pose_variant=pose_variant, pose_path=pose_path, intri_path=intri_path, image_paths=image_paths, depth_paths=depth_paths, frame_ids=frame_ids, w2c=w2c, c2w=c2w, intrinsics=intri, image_hw=(h, w), ) class SequenceCache: def __init__(self, artifacts: SequenceArtifacts): self.artifacts = artifacts self._rgb: Dict[int, np.ndarray] = {} self._gray: Dict[int, np.ndarray] = {} self._depth: Dict[int, np.ndarray] = {} def rgb(self, frame_idx: int) -> np.ndarray: if frame_idx not in self._rgb: image = cv2.imread(self.artifacts.image_paths[frame_idx], cv2.IMREAD_COLOR) if image is None: raise FileNotFoundError(self.artifacts.image_paths[frame_idx]) self._rgb[frame_idx] = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) return self._rgb[frame_idx] def gray(self, frame_idx: int) -> np.ndarray: if frame_idx not in self._gray: self._gray[frame_idx] = cv2.cvtColor(self.rgb(frame_idx), cv2.COLOR_RGB2GRAY) return self._gray[frame_idx] def depth(self, frame_idx: int) -> np.ndarray: if frame_idx not in self._depth: self._depth[frame_idx] = np.load(self.artifacts.depth_paths[frame_idx]).astype( np.float32 ) return self._depth[frame_idx] def select_keyframes(num_frames: int, stride: int) -> Tuple[np.ndarray, np.ndarray]: selector = KeyframeSelector( min_interval=stride, max_interval=stride, force_first=True, mode="fixed", ) is_keyframe, keyframe_indices = selector.select_keyframes( num_frames, batch_size=1, device=torch.device("cpu") ) return is_keyframe[0].numpy(), keyframe_indices[0].numpy() def pose_matrix_to_vec(pose: np.ndarray) -> np.ndarray: rotvec = Rotation.from_matrix(pose[:3, :3]).as_rotvec() return np.concatenate([rotvec, pose[:3, 3]], axis=0) def pose_vec_to_matrix(vec: np.ndarray) -> np.ndarray: pose = np.eye(4, dtype=np.float64) pose[:3, :3] = Rotation.from_rotvec(vec[:3]).as_matrix() pose[:3, 3] = vec[3:6] return pose def project_to_rotation(mat: np.ndarray) -> np.ndarray: u, _, vt = np.linalg.svd(np.asarray(mat, dtype=np.float64)) rot = u @ vt if np.linalg.det(rot) < 0: vt[-1] *= -1 rot = u @ vt return rot.astype(np.float64, copy=False) def similarity_components(mat: np.ndarray) -> Tuple[float, np.ndarray, np.ndarray]: linear = np.asarray(mat[:3, :3], dtype=np.float64) trans = np.asarray(mat[:3, 3], dtype=np.float64) col_norms = np.linalg.norm(linear, axis=0) scale = float(np.mean(col_norms)) if not np.isfinite(scale) or scale < 1e-12: scale = 1.0 rot = project_to_rotation(linear / scale) return scale, rot, trans def similarity_matrix(scale: float, rot: np.ndarray, trans: np.ndarray) -> np.ndarray: out = np.eye(4, dtype=np.float64) out[:3, :3] = float(scale) * np.asarray(rot, dtype=np.float64) out[:3, 3] = np.asarray(trans, dtype=np.float64) return out def normalize_pose_matrix(mat: np.ndarray) -> np.ndarray: _, rot, trans = similarity_components(mat) out = np.eye(4, dtype=np.float64) out[:3, :3] = rot out[:3, 3] = trans return out def relative_measurement_from_c2w(src_c2w: np.ndarray, dst_c2w: np.ndarray) -> np.ndarray: return np.linalg.inv(dst_c2w) @ src_c2w def camera_centers_from_w2c(w2c: np.ndarray) -> np.ndarray: rot = w2c[:, :3, :3] t = w2c[:, :3, 3] return -(np.transpose(rot, (0, 2, 1)) @ t[..., None])[..., 0] def _top_down_xy(c2w: np.ndarray) -> np.ndarray: xyz = c2w[:, :3, 3] return xyz[:, [0, 2]] def _normalize_rows(x: np.ndarray, eps: float = 1e-8) -> np.ndarray: norms = np.linalg.norm(x, axis=1, keepdims=True) return x / np.clip(norms, eps, None) def _copy_loop_cfg(cfg: LoopConfig) -> LoopConfig: return LoopConfig(**asdict(cfg)) def _imagenet_normalize(images: torch.Tensor) -> torch.Tensor: mean = torch.tensor([0.485, 0.456, 0.406], device=images.device).view(1, 3, 1, 1) std = torch.tensor([0.229, 0.224, 0.225], device=images.device).view(1, 3, 1, 1) return (images - mean) / std def _resize_rgb_batch(batch: np.ndarray, image_size: Tuple[int, int]) -> np.ndarray: height, width = int(image_size[0]), int(image_size[1]) return np.stack( [ cv2.resize(image, (width, height), interpolation=cv2.INTER_LINEAR) for image in batch ], axis=0, ) def _batch_to_imagenet_tensor(batch: np.ndarray, device: torch.device) -> torch.Tensor: images = torch.from_numpy(batch).float().permute(0, 3, 1, 2).to(device) / 255.0 return _imagenet_normalize(images) def salad_log_otp_solver( log_a: torch.Tensor, log_b: torch.Tensor, scores: torch.Tensor, num_iters: int = 20, reg: float = 1.0, ) -> torch.Tensor: scores = scores / reg u = torch.zeros_like(log_a) v = torch.zeros_like(log_b) for _ in range(num_iters): u = log_a - torch.logsumexp(scores + v.unsqueeze(1), dim=2) v = log_b - torch.logsumexp(scores + u.unsqueeze(2), dim=1) return scores + u.unsqueeze(2) + v.unsqueeze(1) def salad_matching_probs( scores: torch.Tensor, dustbin_score: torch.Tensor, num_iters: int = 3, reg: float = 1.0, ) -> torch.Tensor: batch_size, m, n = scores.size() scores_aug = torch.empty(batch_size, m + 1, n, dtype=scores.dtype, device=scores.device) scores_aug[:, :m, :n] = scores scores_aug[:, m, :] = dustbin_score norm = -torch.tensor(math.log(n + m), device=scores.device) log_a = norm.expand(m + 1).contiguous() log_b = norm.expand(n).contiguous() log_a[-1] = log_a[-1] + math.log(n - m) log_a = log_a.expand(batch_size, -1) log_b = log_b.expand(batch_size, -1) return salad_log_otp_solver(log_a, log_b, scores_aug, num_iters=num_iters, reg=reg) - norm class SaladAggregator(torch.nn.Module): def __init__( self, num_channels: int = 768, num_clusters: int = 64, cluster_dim: int = 128, token_dim: int = 256, dropout: float = 0.3, ) -> None: super().__init__() dropout_layer: torch.nn.Module dropout_layer = torch.nn.Dropout(dropout) if dropout > 0 else torch.nn.Identity() self.num_channels = int(num_channels) self.num_clusters = int(num_clusters) self.cluster_dim = int(cluster_dim) self.token_dim = int(token_dim) self.token_features = torch.nn.Sequential( torch.nn.Linear(self.num_channels, 512), torch.nn.ReLU(), torch.nn.Linear(512, self.token_dim), ) self.cluster_features = torch.nn.Sequential( torch.nn.Conv2d(self.num_channels, 512, 1), dropout_layer, torch.nn.ReLU(), torch.nn.Conv2d(512, self.cluster_dim, 1), ) self.score = torch.nn.Sequential( torch.nn.Conv2d(self.num_channels, 512, 1), dropout_layer, torch.nn.ReLU(), torch.nn.Conv2d(512, self.num_clusters, 1), ) self.dust_bin = torch.nn.Parameter(torch.tensor(1.0)) def forward(self, features_and_token: Tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor: features, token = features_and_token cluster_features = self.cluster_features(features).flatten(2) scores = self.score(features).flatten(2) token = self.token_features(token) probs = torch.exp(salad_matching_probs(scores, self.dust_bin, 3)) probs = probs[:, :-1, :] probs = probs.unsqueeze(1).repeat(1, self.cluster_dim, 1, 1) cluster_features = cluster_features.unsqueeze(2).repeat(1, 1, self.num_clusters, 1) descriptor = torch.cat( [ torch.nn.functional.normalize(token, p=2, dim=-1), torch.nn.functional.normalize( (cluster_features * probs).sum(dim=-1), p=2, dim=1 ).flatten(1), ], dim=-1, ) return torch.nn.functional.normalize(descriptor, p=2, dim=-1) class DinoV2BackboneForSalad(torch.nn.Module): CHANNELS = { "dinov2_vits14": 384, "dinov2_vitb14": 768, "dinov2_vitl14": 1024, "dinov2_vitg14": 1536, } def __init__(self, model_name: str, weights_path: Optional[str], num_trainable_blocks: int = 4): super().__init__() if model_name not in self.CHANNELS: raise ValueError(f"Unsupported SALAD backbone: {model_name}") self.model_name = model_name self.num_channels = int(self.CHANNELS[model_name]) self.num_trainable_blocks = int(num_trainable_blocks) with warnings.catch_warnings(), contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()): warnings.filterwarnings("ignore", message=".*xFormers is not available.*") self.model = torch.hub.load("facebookresearch/dinov2", model_name, pretrained=False) if weights_path and os.path.exists(weights_path): state = torch.load(weights_path, map_location="cpu") self.model.load_state_dict(state, strict=True) def forward(self, images: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: batch, _, height, width = images.shape tokens = self.model.prepare_tokens_with_masks(images) for block in self.model.blocks: tokens = block(tokens) tokens = self.model.norm(tokens) cls_token = tokens[:, 0] patch_tokens = tokens[:, 1:] feature_map = patch_tokens.reshape( batch, height // 14, width // 14, self.num_channels ).permute(0, 3, 1, 2) return feature_map, cls_token, patch_tokens class SaladVPRModel(torch.nn.Module): def __init__(self, backbone_name: str, dino_weights_path: Optional[str]): super().__init__() self.backbone = DinoV2BackboneForSalad(backbone_name, dino_weights_path) self.aggregator = SaladAggregator( num_channels=self.backbone.num_channels, num_clusters=64, cluster_dim=128, token_dim=256, ) def forward_features(self, images: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: feature_map, cls_token, patch_tokens = self.backbone(images) descriptor = self.aggregator((feature_map, cls_token)) return descriptor, torch.nn.functional.normalize(patch_tokens, dim=-1) def forward(self, images: torch.Tensor) -> torch.Tensor: descriptor, _ = self.forward_features(images) return descriptor def _load_salad_state_dict(model: torch.nn.Module, ckpt_path: str) -> List[str]: if not os.path.exists(ckpt_path): raise FileNotFoundError( f"SALAD checkpoint not found: {ckpt_path}. " "Download it from https://github.com/serizba/salad/releases/download/v1.0.0/dino_salad.ckpt" ) payload = torch.load(ckpt_path, map_location="cpu") state = payload.get("state_dict", payload) if isinstance(payload, dict) else payload cleaned = {} for key, value in state.items(): out_key = str(key) for prefix in ("model.", "module."): if out_key.startswith(prefix): out_key = out_key[len(prefix) :] cleaned[out_key] = value missing, unexpected = model.load_state_dict(cleaned, strict=False) missing = list(missing) missing_aggregator = [key for key in missing if key.startswith("aggregator.")] if missing_aggregator: raise RuntimeError( "SALAD checkpoint did not provide the expected aggregator weights: " + ", ".join(missing_aggregator[:5]) ) if unexpected: _log( "SALAD: ignored unexpected checkpoint keys: " + ", ".join(str(key) for key in unexpected[:5]), enabled=True, ) return missing _POPCOUNT_TABLE = np.unpackbits( np.arange(256, dtype=np.uint8)[:, None], axis=1 ).sum(axis=1).astype(np.uint8) class FlatORBVocabulary: def __init__(self, path: str): self.path = path self.centers, self.idf = self._load(path) self.num_words = int(self.centers.shape[0]) @staticmethod def _load(path: str) -> Tuple[np.ndarray, Optional[np.ndarray]]: if not os.path.exists(path): raise FileNotFoundError(path) ext = os.path.splitext(path)[1].lower() if ext == ".npz": data = np.load(path) if "centers" not in data: raise ValueError(f"Expected 'centers' in vocabulary file: {path}") centers = np.asarray(data["centers"], dtype=np.float32) idf = np.asarray(data["idf"], dtype=np.float32) if "idf" in data else None elif ext == ".npy": centers = np.asarray(np.load(path), dtype=np.float32) idf = None else: try: centers = np.asarray(np.loadtxt(path), dtype=np.float32) except Exception as exc: raise ValueError( f"Unsupported ORB vocabulary text format in {path}. " "This loader expects a flat centers matrix (.npz/.npy/.txt), not a hierarchical DBoW2 ORBvoc.txt." ) from exc idf = None if centers.ndim != 2 or centers.shape[1] != 32: raise ValueError( f"Expected ORB vocabulary centers with shape [K, 32], got {tuple(centers.shape)} from {path}" ) if idf is not None: if idf.ndim != 1 or idf.shape[0] != centers.shape[0]: raise ValueError( f"Expected idf shape [{centers.shape[0]}], got {tuple(idf.shape)} from {path}" ) return centers.astype(np.float32, copy=False), idf def quantize(self, desc: np.ndarray) -> np.ndarray: diff = desc.astype(np.float32, copy=False)[:, None, :] - self.centers[None, :, :] dist2 = np.sum(diff * diff, axis=-1) return dist2.argmin(axis=1).astype(np.int64, copy=False) def histogram(self, desc: Optional[np.ndarray]) -> np.ndarray: hist = np.zeros(self.num_words, dtype=np.float32) if desc is None or len(desc) == 0: return hist words = self.quantize(desc) np.add.at(hist, words, 1.0) if self.idf is not None: hist *= self.idf return _normalize_rows(hist[None] + 1e-12)[0] def score_matrix(self, hists: np.ndarray) -> np.ndarray: hists = _normalize_rows(hists.astype(np.float32)) return hists @ hists.T class DBoWTextVocabulary: def __init__(self, path: str): self.path = path self.nodes: List[Dict[str, object]] = [ { "id": 0, "parent": -1, "children": [], "descriptor": None, "weight": 0.0, "word_id": -1, "is_leaf": False, } ] self.word_node_ids: List[int] = [] self.k = 0 self.L = 0 self.scoring = 0 self.weighting = 0 self._load(path) self.num_words = len(self.word_node_ids) def _load(self, path: str) -> None: if not os.path.exists(path): raise FileNotFoundError(path) with open(path, "r") as f: header = f.readline().strip().split() if len(header) != 4: raise ValueError(f"Invalid ORBvoc header in {path}: {' '.join(header)}") self.k, self.L, self.scoring, self.weighting = [int(x) for x in header] for line in f: line = line.strip() if not line: continue parts = line.split() if len(parts) != 35: raise ValueError( f"Expected 35 fields per ORBvoc node, got {len(parts)} in {path}" ) parent = int(parts[0]) is_leaf = int(parts[1]) > 0 descriptor = np.asarray([int(x) for x in parts[2:34]], dtype=np.uint8) weight = float(parts[34]) node_id = len(self.nodes) node = { "id": node_id, "parent": parent, "children": [], "descriptor": descriptor, "weight": weight, "word_id": -1, "is_leaf": is_leaf, } self.nodes.append(node) self.nodes[parent]["children"].append(node_id) if is_leaf: node["word_id"] = len(self.word_node_ids) self.word_node_ids.append(node_id) if self.weighting != 0 or self.scoring != 0: raise ValueError( f"Unsupported ORBvoc scoring/weighting ({self.scoring}, {self.weighting}) in {path}. " "This loader currently supports the standard TF_IDF + L1_NORM ORBvoc only." ) def _hamming_distance(self, desc: np.ndarray, child_descs: np.ndarray) -> np.ndarray: xor = np.bitwise_xor(desc[:, None, :], child_descs[None, :, :]) return _POPCOUNT_TABLE[xor].sum(axis=2) def quantize(self, desc: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: current = np.zeros(len(desc), dtype=np.int64) while True: next_nodes = current.copy() active = False for parent_id in np.unique(current): idx = np.flatnonzero(current == parent_id) if idx.size == 0: continue children = self.nodes[parent_id]["children"] if not children: continue active = True child_descs = np.stack( [self.nodes[child_id]["descriptor"] for child_id in children], axis=0 ) dist = self._hamming_distance(desc[idx], child_descs) best = dist.argmin(axis=1) next_nodes[idx] = np.asarray(children, dtype=np.int64)[best] current = next_nodes if not active: break if all(bool(self.nodes[node_id]["is_leaf"]) for node_id in np.unique(current)): break word_ids = np.asarray( [int(self.nodes[node_id]["word_id"]) for node_id in current], dtype=np.int64 ) weights = np.asarray( [float(self.nodes[node_id]["weight"]) for node_id in current], dtype=np.float32 ) return word_ids, weights def sparse_histogram(self, desc: Optional[np.ndarray]) -> Dict[int, float]: hist: Dict[int, float] = {} if desc is None or len(desc) == 0: return hist word_ids, weights = self.quantize(desc) for word_id, weight in zip(word_ids.tolist(), weights.tolist()): hist[int(word_id)] = hist.get(int(word_id), 0.0) + float(weight) total = float(sum(hist.values())) if total > 0: hist = {word_id: value / total for word_id, value in hist.items()} return hist def score_matrix(self, hists: Sequence[Dict[int, float]]) -> np.ndarray: n = len(hists) sim = np.eye(n, dtype=np.float32) for i in range(n): hi = hists[i] for j in range(i): hj = hists[j] if len(hi) > len(hj): hi, hj = hj, hi score = 0.0 for word_id, value in hi.items(): other = hj.get(word_id) if other is not None: score += min(value, other) sim[i, j] = score sim[j, i] = score return sim def _load_orb_vocabulary(path: str): if os.path.basename(path) == "ORBvoc.txt": return DBoWTextVocabulary(path) return FlatORBVocabulary(path) def _resolve_gt_sequence_info( seq_dir: str, data_cfg: Optional[dict], seq_name_hint: Optional[str] = None, ) -> Optional[HorizonStreamSequenceInfo]: if not data_cfg: return None seq_names = [] if seq_name_hint: seq_names.append(str(seq_name_hint).replace("\\", "/").strip("/")) seq_names.append(str(os.path.basename(os.path.normpath(seq_dir))).replace("\\", "/").strip("/")) preferred_camera = data_cfg.get("camera", None) if preferred_camera is not None: preferred_camera = str(preferred_camera).strip() loader = HorizonStreamDataLoader(dict(data_cfg)) prefixed_matches: List[HorizonStreamSequenceInfo] = [] suffixed_matches: List[HorizonStreamSequenceInfo] = [] for seq_info in loader.iter_sequence_infos(): name_norm = str(seq_info.name).replace("\\", "/").strip("/") for seq_name_norm in seq_names: if name_norm == seq_name_norm: return seq_info parts = [part for part in name_norm.split("/") if part] for seq_name_norm in seq_names: if parts and parts[0] == seq_name_norm: prefixed_matches.append(seq_info) if parts and parts[-1] == seq_name_norm: suffixed_matches.append(seq_info) candidates = prefixed_matches or suffixed_matches if not candidates: return None if preferred_camera: for seq_info in candidates: if str(seq_info.camera or "") == preferred_camera: return seq_info for seq_info in candidates: if str(seq_info.camera or "") == "02": return seq_info candidates = sorted(candidates, key=lambda item: str(item.name)) if candidates: return candidates[0] return None def load_ground_truth_trajectory( seq_dir: str, data_cfg: Optional[dict], verbose: bool = True, seq_name_hint: Optional[str] = None, ) -> Optional[GroundTruthTrajectory]: seq_info = _resolve_gt_sequence_info(seq_dir, data_cfg, seq_name_hint=seq_name_hint) if seq_info is None: _log("GT: no matching sequence info found from data config", verbose) return None if seq_info.camera is not None: cam_dir = os.path.join(seq_info.scene_root, "cameras", seq_info.camera) extri_path = os.path.join(cam_dir, "extri.yml") intri_path = os.path.join(cam_dir, "intri.yml") else: extri_path = os.path.join(seq_info.scene_root, "extri.yml") intri_path = os.path.join(seq_info.scene_root, "intri.yml") if not os.path.exists(extri_path): _log(f"GT: missing pose file {extri_path}", verbose) return None extri, _, _ = read_opencv_camera_yml(extri_path, intri_path) stems = frame_stems(seq_info.image_paths) gt_c2w = [] valid_mask = [] for stem in stems: if stem in extri: gt_c2w.append(np.linalg.inv(extri[stem])) valid_mask.append(True) else: gt_c2w.append(np.eye(4, dtype=np.float64)) valid_mask.append(False) valid_mask = np.asarray(valid_mask, dtype=bool) if not np.any(valid_mask): _log( f"GT: pose file {extri_path} does not contain any matching frame ids", verbose, ) return None _log( f"GT: loaded {int(valid_mask.sum())}/{len(valid_mask)} poses from {seq_info.scene_root}" + (f" camera={seq_info.camera}" if seq_info.camera else ""), verbose, ) return GroundTruthTrajectory( seq_name=seq_info.name, scene_root=seq_info.scene_root, camera=seq_info.camera, c2w=np.asarray(gt_c2w, dtype=np.float64), valid_mask=valid_mask, ) def evaluate_trajectory_against_gt( pred_c2w: np.ndarray, frame_ids: Sequence[int], gt: Optional[GroundTruthTrajectory], ) -> Optional[Dict[str, object]]: if gt is None: return None pred_xyz = [] gt_xyz = [] used_frame_ids = [] for pred_idx, frame_id in enumerate(frame_ids): if frame_id < 0 or frame_id >= len(gt.c2w): continue if not gt.valid_mask[frame_id]: continue pred_xyz.append(pred_c2w[pred_idx, :3, 3]) gt_xyz.append(gt.c2w[frame_id, :3, 3]) used_frame_ids.append(int(frame_id)) if len(pred_xyz) < 3: return None pred_xyz = np.asarray(pred_xyz, dtype=np.float64) gt_xyz = np.asarray(gt_xyz, dtype=np.float64) metrics = ate_rmse(pred_xyz, gt_xyz, align_scale=True) scale = float(metrics["sim3_scale"]) rot = np.asarray(metrics["sim3_rotation"], dtype=np.float64) trans = np.asarray(metrics["sim3_translation"], dtype=np.float64) pred_xyz_aligned = transform_points(pred_xyz, scale, rot, trans) metrics["frame_ids"] = used_frame_ids metrics["pred_xyz_aligned"] = pred_xyz_aligned.tolist() metrics["gt_xyz"] = gt_xyz.tolist() return metrics def compact_gt_metrics(metrics: Optional[Dict[str, object]]) -> Optional[Dict[str, object]]: if metrics is None: return None out = dict(metrics) out.pop("pred_xyz_aligned", None) out.pop("gt_xyz", None) out.pop("frame_ids", None) return out class ORBBoWRetriever: def __init__(self, cfg: LoopConfig): self.cfg = cfg if not cfg.dbow_vocab_path: raise ValueError("LoopConfig.dbow_vocab_path must be set for ORB-BoW retrieval") self.orb = cv2.ORB_create( nfeatures=cfg.orb_features, scaleFactor=1.2, nlevels=8, edgeThreshold=31, patchSize=31, fastThreshold=20, ) self.keypoints: Dict[int, List[cv2.KeyPoint]] = {} self.descriptors: Dict[int, Optional[np.ndarray]] = {} self.histograms = None self.vocab = _load_orb_vocabulary(cfg.dbow_vocab_path) _log( f"ORB-BoW: loaded fixed vocabulary {cfg.dbow_vocab_path} " f"with {self.vocab.num_words} words", self.cfg.verbose, ) def fit(self, cache: SequenceCache, keyframe_ids: Sequence[int]) -> np.ndarray: _log( f"ORB-BoW: extracting ORB descriptors for {len(keyframe_ids)} keyframes", self.cfg.verbose, ) total = len(keyframe_ids) for idx, frame_idx in enumerate(keyframe_ids, start=1): gray = cache.gray(frame_idx) kps, desc = self.orb.detectAndCompute(gray, None) self.keypoints[frame_idx] = kps or [] self.descriptors[frame_idx] = desc _log_every("ORB keyframes", idx, total, self.cfg.verbose, step=20) if not any(desc is not None and len(desc) for desc in self.descriptors.values()): if isinstance(self.vocab, DBoWTextVocabulary): self.histograms = [{} for _ in keyframe_ids] else: self.histograms = np.zeros( (len(keyframe_ids), int(self.vocab.num_words)), dtype=np.float32 ) _log("ORB-BoW: no descriptors found, returning empty histograms", self.cfg.verbose) return self.histograms if isinstance(self.vocab, DBoWTextVocabulary): hists = [] for frame_idx in keyframe_ids: desc = self.descriptors[frame_idx] hists.append(self.vocab.sparse_histogram(desc)) self.histograms = hists mean_nnz = float(np.mean([len(hist) for hist in hists])) if hists else 0.0 _log( f"ORB-BoW: built sparse descriptor bank for {len(hists)} keyframes " f"(mean nnz={mean_nnz:.1f})", self.cfg.verbose, ) else: hists = [] for frame_idx in keyframe_ids: desc = self.descriptors[frame_idx] hists.append(self.vocab.histogram(desc)) self.histograms = np.asarray(hists, dtype=np.float32) _log( f"ORB-BoW: built descriptor bank with shape {tuple(self.histograms.shape)}", self.cfg.verbose, ) return self.histograms def score_matrix(self) -> np.ndarray: if self.histograms is None: raise RuntimeError("ORB-BoW histograms are not initialized") return self.vocab.score_matrix(self.histograms) def collect_3d_correspondences( self, cache: SequenceCache, artifacts: SequenceArtifacts, src_frame: int, dst_frame: int, ) -> Tuple[np.ndarray, np.ndarray]: desc_src = self.descriptors.get(src_frame) desc_dst = self.descriptors.get(dst_frame) if desc_src is None or desc_dst is None or len(desc_src) < 8 or len(desc_dst) < 8: return ( np.empty((0, 3), dtype=np.float64), np.empty((0, 3), dtype=np.float64), ) matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=False) knn_a = matcher.knnMatch(desc_src, desc_dst, k=2) knn_b = matcher.knnMatch(desc_dst, desc_src, k=2) def ratio_filter(knn_matches): passed = {} for pair in knn_matches: if len(pair) < 2: continue m, n = pair if m.distance < self.cfg.orb_ratio_test * n.distance: passed[m.queryIdx] = (m.trainIdx, float(m.distance)) return passed forward = ratio_filter(knn_a) backward = ratio_filter(knn_b) matches = [] for q_idx, (t_idx, dist) in forward.items(): rev = backward.get(t_idx) if rev is not None and rev[0] == q_idx: matches.append((q_idx, t_idx, dist)) if len(matches) < self.cfg.rigid_min_inliers: return ( np.empty((0, 3), dtype=np.float64), np.empty((0, 3), dtype=np.float64), ) kps_src = self.keypoints[src_frame] kps_dst = self.keypoints[dst_frame] uv_src = np.asarray([kps_src[q].pt for q, _, _ in matches], dtype=np.float32) uv_dst = np.asarray([kps_dst[t].pt for _, t, _ in matches], dtype=np.float32) return depth_correspondences_to_3d( uv_src, uv_dst, cache.depth(src_frame), cache.depth(dst_frame), artifacts.intrinsics[src_frame], artifacts.intrinsics[dst_frame], ) class SaladRetriever: def __init__(self, cfg: LoopConfig, device: str): self.cfg = cfg self.device = torch.device(device) if not os.path.exists(cfg.salad_ckpt_path): raise FileNotFoundError( f"SALAD checkpoint not found: {cfg.salad_ckpt_path}. " "Download it from https://github.com/serizba/salad/releases/download/v1.0.0/dino_salad.ckpt" ) self.model = SaladVPRModel( backbone_name=cfg.salad_backbone, dino_weights_path=cfg.salad_dino_weights_path, ) missing = _load_salad_state_dict(self.model, cfg.salad_ckpt_path) if any(key.startswith("backbone.") for key in missing) and not os.path.exists( cfg.salad_dino_weights_path ): raise FileNotFoundError( f"DINOv2 backbone weights not found: {cfg.salad_dino_weights_path}. " "Download it from https://dl.fbaipublicfiles.com/dinov2/dinov2_vitb14/dinov2_vitb14_pretrain.pth" ) self.model = self.model.to(self.device).eval() self.global_desc: Optional[np.ndarray] = None self.local_desc: Dict[int, np.ndarray] = {} self.patch_hw: Optional[Tuple[int, int]] = None self.patch_size = 14 def fit(self, cache: SequenceCache, keyframe_ids: Sequence[int]) -> np.ndarray: batch_size = max(1, int(self.cfg.salad_batch_size)) image_size = tuple(int(x) for x in self.cfg.salad_image_size) globals_out = [] use_amp = self.device.type == "cuda" with torch.no_grad(): for batch_idx, start in enumerate(range(0, len(keyframe_ids), batch_size), start=1): batch_ids = list(keyframe_ids[start : start + batch_size]) batch = np.stack([cache.rgb(i) for i in batch_ids], axis=0) resized = _resize_rgb_batch(batch, image_size) salad_images = _batch_to_imagenet_tensor(resized, self.device) local_images = _batch_to_imagenet_tensor(batch, self.device) with torch.autocast(device_type=self.device.type, dtype=torch.float16, enabled=use_amp): desc = self.model(salad_images) _, patch_tokens = self.model.forward_features(local_images) globals_out.append(desc.float().cpu().numpy().astype(np.float32)) patch_tokens_np = patch_tokens.float().cpu().numpy().astype(np.float32) for frame_idx, patch_desc in zip(batch_ids, patch_tokens_np): self.local_desc[frame_idx] = patch_desc if self.patch_hw is None: h, w = batch.shape[1:3] self.patch_hw = (h // self.patch_size, w // self.patch_size) self.global_desc = np.concatenate(globals_out, axis=0) self.global_desc = _normalize_rows(self.global_desc.astype(np.float32)) _log( f"SALAD: built descriptor bank with shape {tuple(self.global_desc.shape)}", self.cfg.verbose, ) return self.global_desc def collect_3d_correspondences( self, cache: SequenceCache, artifacts: SequenceArtifacts, src_frame: int, dst_frame: int, ) -> Tuple[np.ndarray, np.ndarray]: desc_src = self.local_desc.get(src_frame) desc_dst = self.local_desc.get(dst_frame) if desc_src is None or desc_dst is None: return ( np.empty((0, 3), dtype=np.float64), np.empty((0, 3), dtype=np.float64), ) sim = desc_src @ desc_dst.T dst_best = sim.argmax(axis=1) dst_score = sim[np.arange(sim.shape[0]), dst_best] src_best = sim.argmax(axis=0) matches = [] for src_idx, dst_idx in enumerate(dst_best.tolist()): if src_best[dst_idx] != src_idx: continue score = float(dst_score[src_idx]) if score < self.cfg.patch_match_thresh: continue matches.append((src_idx, dst_idx, score)) if len(matches) < self.cfg.rigid_min_inliers: return ( np.empty((0, 3), dtype=np.float64), np.empty((0, 3), dtype=np.float64), ) matches.sort(key=lambda item: item[2], reverse=True) matches = matches[: self.cfg.max_pair_matches] uv_src = np.asarray( [patch_index_to_uv(m[0], self.patch_hw, self.patch_size) for m in matches], dtype=np.float32, ) uv_dst = np.asarray( [patch_index_to_uv(m[1], self.patch_hw, self.patch_size) for m in matches], dtype=np.float32, ) return depth_correspondences_to_3d( uv_src, uv_dst, cache.depth(src_frame), cache.depth(dst_frame), artifacts.intrinsics[src_frame], artifacts.intrinsics[dst_frame], ) def patch_index_to_uv(index: int, patch_hw: Tuple[int, int], patch_size: int) -> Tuple[float, float]: h_patch, w_patch = patch_hw y = index // w_patch x = index % w_patch cx = (x + 0.5) * patch_size cy = (y + 0.5) * patch_size return float(cx), float(cy) def sample_depth_at_uv(depth: np.ndarray, uv: np.ndarray) -> np.ndarray: h, w = depth.shape[:2] x = np.clip(np.round(uv[:, 0]).astype(np.int32), 0, w - 1) y = np.clip(np.round(uv[:, 1]).astype(np.int32), 0, h - 1) return depth[y, x] def unproject_points(uv: np.ndarray, depth: np.ndarray, intrinsics: np.ndarray) -> np.ndarray: z = depth.astype(np.float64) x = (uv[:, 0] - intrinsics[0, 2]) / intrinsics[0, 0] * z y = (uv[:, 1] - intrinsics[1, 2]) / intrinsics[1, 1] * z return np.stack([x, y, z], axis=1) def depth_correspondences_to_3d( uv_src: np.ndarray, uv_dst: np.ndarray, depth_src: np.ndarray, depth_dst: np.ndarray, intr_src: np.ndarray, intr_dst: np.ndarray, ) -> Tuple[np.ndarray, np.ndarray]: z_src = sample_depth_at_uv(depth_src, uv_src) z_dst = sample_depth_at_uv(depth_dst, uv_dst) valid = np.isfinite(z_src) & np.isfinite(z_dst) & (z_src > 1e-4) & (z_dst > 1e-4) if not np.any(valid): return np.empty((0, 3), dtype=np.float64), np.empty((0, 3), dtype=np.float64) pts_src = unproject_points(uv_src[valid], z_src[valid], intr_src) pts_dst = unproject_points(uv_dst[valid], z_dst[valid], intr_dst) return pts_src, pts_dst def estimate_similarity_umeyama( src_pts: np.ndarray, dst_pts: np.ndarray, allow_scale: bool = True, weights: Optional[np.ndarray] = None, ) -> Tuple[float, np.ndarray, np.ndarray]: src = np.asarray(src_pts, dtype=np.float64) dst = np.asarray(dst_pts, dtype=np.float64) if src.shape != dst.shape or src.ndim != 2 or src.shape[1] != 3: raise ValueError("Expected src_pts and dst_pts with shape [N, 3]") if len(src) < 3: raise ValueError("At least 3 points are required") if weights is None: weights = np.ones(len(src), dtype=np.float64) else: weights = np.asarray(weights, dtype=np.float64).reshape(-1) if weights.shape[0] != src.shape[0]: raise ValueError("weights must have shape [N]") weights = np.clip(weights, 1e-12, None) weights /= weights.sum() src_mean = np.sum(src * weights[:, None], axis=0) dst_mean = np.sum(dst * weights[:, None], axis=0) src_centered = src - src_mean dst_centered = dst - dst_mean cov = (weights[:, None] * dst_centered).T @ src_centered u, singular_vals, vt = np.linalg.svd(cov) s_mat = np.eye(3, dtype=np.float64) if np.linalg.det(u) * np.linalg.det(vt) < 0: s_mat[-1, -1] = -1.0 rot = u @ s_mat @ vt if allow_scale: src_var = float(np.sum(weights * np.sum(src_centered * src_centered, axis=1))) if src_var < 1e-12: raise ValueError("Degenerate source variance for Sim3 estimation") scale = float(np.sum(singular_vals * np.diag(s_mat)) / src_var) else: scale = 1.0 trans = dst_mean - scale * (rot @ src_mean) return scale, rot.astype(np.float64, copy=False), trans.astype(np.float64, copy=False) def apply_similarity_to_points( pts: np.ndarray, scale: float, rot: np.ndarray, trans: np.ndarray, ) -> np.ndarray: return scale * (np.asarray(rot, dtype=np.float64) @ np.asarray(pts, dtype=np.float64).T).T + np.asarray( trans, dtype=np.float64 ) def robust_refine_similarity_transform( src_pts: np.ndarray, dst_pts: np.ndarray, cfg: LoopConfig, ) -> np.ndarray: allow_scale = cfg.pose_graph_model == "sim3" weights = np.ones(len(src_pts), dtype=np.float64) scale, rot, trans = estimate_similarity_umeyama( src_pts, dst_pts, allow_scale=allow_scale, weights=weights ) prev_cost = float("inf") delta = float(cfg.sim3_irls_delta) for _ in range(max(1, int(cfg.sim3_irls_max_iters))): pred = apply_similarity_to_points(src_pts, scale, rot, trans) resid = np.linalg.norm(pred - dst_pts, axis=1) huber = np.ones_like(resid) large = resid > delta huber[large] = delta / np.clip(resid[large], 1e-12, None) huber = np.clip(huber, 1e-6, None) scale_new, rot_new, trans_new = estimate_similarity_umeyama( src_pts, dst_pts, allow_scale=allow_scale, weights=huber ) current_cost = float( np.sum( np.where( resid <= delta, 0.5 * resid * resid, delta * (resid - 0.5 * delta), ) ) ) param_change = abs(scale_new - scale) + np.linalg.norm(trans_new - trans) rot_angle = Rotation.from_matrix(rot_new @ rot.T).magnitude() scale, rot, trans = scale_new, rot_new, trans_new if ( param_change < float(cfg.sim3_irls_tol) and rot_angle < np.deg2rad(0.1) ) or abs(prev_cost - current_cost) < float(cfg.sim3_irls_tol) * max(prev_cost, 1.0): break prev_cost = current_cost return similarity_matrix(scale, rot, trans) def estimate_similarity_transform_ransac( src_pts: np.ndarray, dst_pts: np.ndarray, cfg: LoopConfig, ) -> Optional[Tuple[np.ndarray, int]]: if len(src_pts) < cfg.rigid_min_inliers: return None allow_scale = cfg.pose_graph_model == "sim3" rng = np.random.default_rng(0) best_inliers: Optional[np.ndarray] = None best_count = 0 sample_size = 3 for _ in range(cfg.rigid_ransac_iters): sample = rng.choice(len(src_pts), size=sample_size, replace=False) try: scale, rot, trans = estimate_similarity_umeyama( src_pts[sample], dst_pts[sample], allow_scale=allow_scale ) except ValueError: continue pred = apply_similarity_to_points(src_pts, scale, rot, trans) err = np.linalg.norm(pred - dst_pts, axis=1) inliers = err < cfg.rigid_inlier_thresh count = int(inliers.sum()) if count > best_count: best_inliers = inliers best_count = count if best_inliers is None or best_count < cfg.rigid_min_inliers: return None refined = robust_refine_similarity_transform(src_pts[best_inliers], dst_pts[best_inliers], cfg) return refined, best_count def retrieve_dbow_candidates( sim: np.ndarray, keyframe_ids: Sequence[int], cfg: LoopConfig, ) -> List[LoopCandidate]: candidates: List[LoopCandidate] = [] found: List[Tuple[int, int, float]] = [] accepted_pairs: List[Tuple[int, int]] = [] repeat = max(1, int(cfg.dbow_num_repeat)) temporal_exclusion = max(1, int(cfg.dbow_temporal_exclusion)) nms = max(0, int(cfg.nms_radius)) for src_pos in range(len(keyframe_ids)): max_dst = src_pos - temporal_exclusion if max_dst < 0: continue prev_scores = sim[src_pos, : max_dst + 1] dst_pos = int(np.argmax(prev_scores)) score = float(prev_scores[dst_pos]) if score < cfg.retrieval_score_thresh_dbow: continue if accepted_pairs and nms > 0: dists_sq = [ (src_pos - prev_src) ** 2 + (dst_pos - prev_dst) ** 2 for prev_src, prev_dst in accepted_pairs ] if min(dists_sq, default=float("inf")) < nms * nms: continue found.append((src_pos, dst_pos, score)) if len(found) < repeat: continue latest = found[-repeat:] first_src = latest[0][0] if (1 + src_pos - first_src) != repeat: continue accepted_src, accepted_dst, accepted_score = latest[repeat // 2] accepted_pairs.append((accepted_src, accepted_dst)) candidates.append( LoopCandidate( src_pos=int(accepted_src), dst_pos=int(max(accepted_dst, 1)), src_frame=int(keyframe_ids[accepted_src]), dst_frame=int(keyframe_ids[max(accepted_dst, 1)]), score=float(accepted_score), method="dbow", ) ) found.clear() candidates.sort(key=lambda item: item.score, reverse=True) return candidates[: cfg.max_candidates_per_method] def non_max_suppress_loop_candidates( candidates: Sequence[LoopCandidate], nms_threshold: int, limit: int, ) -> List[LoopCandidate]: if not candidates or nms_threshold <= 0: return list(candidates)[:limit] sorted_candidates = sorted(candidates, key=lambda item: item.score, reverse=True) max_pos = max(max(item.src_pos, item.dst_pos) for item in sorted_candidates) filtered: List[LoopCandidate] = [] suppressed = set() for cand in sorted_candidates: low = min(cand.src_pos, cand.dst_pos) high = max(cand.src_pos, cand.dst_pos) if low in suppressed or high in suppressed: continue filtered.append(cand) start1 = max(0, low - nms_threshold) end1 = min(low + nms_threshold + 1, high) suppressed.update(range(start1, end1)) start2 = max(low + 1, high - nms_threshold) end2 = min(high + nms_threshold + 1, max_pos + 1) suppressed.update(range(start2, end2)) if len(filtered) >= limit: break return filtered def retrieve_vpr_candidates( descriptors: np.ndarray, keyframe_ids: Sequence[int], method: str, score_thresh: float, cfg: LoopConfig, ) -> List[LoopCandidate]: desc = _normalize_rows(descriptors.astype(np.float32)) sim = desc @ desc.T top_k = max(1, int(cfg.retrieval_top_k)) min_gap = max(1, int(cfg.min_keyframe_gap)) candidates_by_key: Dict[Tuple[int, int], LoopCandidate] = {} for src_pos in range(len(keyframe_ids)): order = np.argsort(sim[src_pos])[::-1] keep = 0 for dst_pos in order.tolist(): if dst_pos == src_pos: continue if keep >= top_k: break score = float(sim[src_pos, dst_pos]) if score <= score_thresh: break keep += 1 if abs(int(keyframe_ids[src_pos]) - int(keyframe_ids[dst_pos])) <= min_gap: continue high = max(src_pos, dst_pos) low = min(src_pos, dst_pos) key = (high, low) prev = candidates_by_key.get(key) if prev is not None and prev.score >= score: continue candidates_by_key[key] = LoopCandidate( src_pos=int(high), dst_pos=int(low), src_frame=int(keyframe_ids[high]), dst_frame=int(keyframe_ids[low]), score=score, method=method, ) return non_max_suppress_loop_candidates( list(candidates_by_key.values()), nms_threshold=max(0, int(cfg.nms_radius)), limit=max(1, int(cfg.max_candidates_per_method)), ) def centered_window_positions(center: int, size: int, total: int) -> np.ndarray: size = max(1, int(size)) total = max(0, int(total)) if total <= 0: return np.empty((0,), dtype=np.int64) if size >= total: return np.arange(total, dtype=np.int64) half = size // 2 start = center - half end = start + size if start < 0: start = 0 end = size if end > total: end = total start = total - size return np.arange(start, end, dtype=np.int64) def candidate_chunk_frame_pairs( candidate: LoopCandidate, keyframe_ids: Sequence[int], chunk_size: int, ) -> List[Tuple[int, int]]: src_pos = int(candidate.src_pos) dst_pos = int(candidate.dst_pos) src_window = centered_window_positions(src_pos, chunk_size, len(keyframe_ids)) dst_window = centered_window_positions(dst_pos, chunk_size, len(keyframe_ids)) span = min(len(src_window), len(dst_window)) if span <= 0: return [] return [ (int(keyframe_ids[src_window[idx]]), int(keyframe_ids[dst_window[idx]])) for idx in range(span) ] def verify_candidates( candidates: Sequence[LoopCandidate], verifier, cache: SequenceCache, artifacts: SequenceArtifacts, keyframe_ids: Sequence[int], cfg: LoopConfig, ) -> List[LoopEdge]: method = candidates[0].method if candidates else "unknown" _log( f"{method}: verifying {len(candidates)} loop candidates with geometry", cfg.verbose, ) verified: List[LoopEdge] = [] total = len(candidates) for idx, cand in enumerate(candidates, start=1): if abs(int(cand.src_pos) - int(cand.dst_pos)) < int(cfg.loop_edge_min_separation): _log_every( f"{method} verify", idx, total, cfg.verbose, step=5, ) continue frame_pairs = candidate_chunk_frame_pairs(cand, keyframe_ids, cfg.loop_chunk_size) pts_src_all = [] pts_dst_all = [] for src_frame, dst_frame in frame_pairs: pts_src, pts_dst = verifier.collect_3d_correspondences( cache, artifacts, src_frame, dst_frame ) if len(pts_src) == 0: continue pts_src_all.append(pts_src) pts_dst_all.append(pts_dst) if not pts_src_all: _log_every( f"{method} verify", idx, total, cfg.verbose, step=5, ) continue pts_src = np.concatenate(pts_src_all, axis=0) pts_dst = np.concatenate(pts_dst_all, axis=0) result = estimate_similarity_transform_ransac(pts_src, pts_dst, cfg) if result is None: _log_every( f"{method} verify", idx, total, cfg.verbose, step=5, ) continue transform_ji, inliers = result verified.append( LoopEdge( src_pos=cand.src_pos, dst_pos=cand.dst_pos, src_frame=cand.src_frame, dst_frame=cand.dst_frame, score=cand.score, inliers=inliers, method=cand.method, transform_ji=transform_ji, ) ) _log( f"{method}: accepted loop frame {cand.src_frame} -> {cand.dst_frame} " f"(score={cand.score:.3f}, inliers={inliers}, chunk={max(1, cfg.loop_chunk_size)})", cfg.verbose, ) _log_every( f"{method} verify", idx, total, cfg.verbose, step=5, ) verified.sort(key=lambda edge: (edge.inliers, edge.score), reverse=True) verified = verified[: cfg.max_verified_loops_per_method] _log( f"{method}: kept {len(verified)} verified loops after ranking", cfg.verbose, ) return verified def build_keyframe_odometry_edges(keyframe_c2w: np.ndarray) -> List[LoopEdge]: edges = [] for idx in range(1, len(keyframe_c2w)): meas = relative_measurement_from_c2w( keyframe_c2w[idx - 1], keyframe_c2w[idx] ) edges.append( LoopEdge( src_pos=idx - 1, dst_pos=idx, src_frame=idx - 1, dst_frame=idx, score=1.0, inliers=0, method="odometry", transform_ji=meas, ) ) return edges def optimize_keyframe_pose_graph( keyframe_c2w_init: np.ndarray, odom_edges: Sequence[LoopEdge], loop_edges: Sequence[LoopEdge], cfg: LoopConfig, ) -> np.ndarray: if cfg.pose_graph_backend == "pypose": all_edges = list(odom_edges) + list(loop_edges) edge_weights = [] for edge in all_edges: if edge.method == "odometry": edge_weights.append(1.0) else: edge_weights.append( cfg.pose_graph_loop_weight * math.sqrt(max(edge.inliers, 1) / max(cfg.rigid_min_inliers, 1)) ) return optimize_pose_graph_pypose( keyframe_c2w_init=keyframe_c2w_init, edge_src=[edge.src_pos for edge in all_edges], edge_dst=[edge.dst_pos for edge in all_edges], edge_measurements=[edge.transform_ji for edge in all_edges], edge_weights=edge_weights, model=cfg.pose_graph_model, update_mode=cfg.pose_graph_update_mode, trans_weight=cfg.pose_graph_trans_weight, rot_weight=cfg.pose_graph_rot_weight, scale_weight=cfg.pose_graph_scale_weight, max_iterations=cfg.pose_graph_max_iterations, lambda_init=cfg.pose_graph_lambda_init, verbose=cfg.pose_graph_solver_verbose, log_fn=lambda msg: _log(msg, cfg.verbose), ) if cfg.pose_graph_backend != "scipy": raise ValueError(f"Unsupported pose_graph_backend: {cfg.pose_graph_backend}") if cfg.pose_graph_model != "se3": raise ValueError("SciPy pose graph backend only supports pose_graph_model='se3'") if len(keyframe_c2w_init) <= 1: return keyframe_c2w_init.copy() num_nodes = len(keyframe_c2w_init) num_vars = (num_nodes - 1) * 6 num_edges = len(odom_edges) + len(loop_edges) num_residuals = num_edges * 6 _log( f"Pose graph: optimizing {num_nodes} keyframes with " f"{len(odom_edges)} odom edges and {len(loop_edges)} loop edges " f"({num_vars} vars, {num_residuals} residuals)", cfg.verbose, ) x0 = np.concatenate( [pose_matrix_to_vec(pose) for pose in keyframe_c2w_init[1:]], axis=0 ) def unpack_poses(x: np.ndarray) -> List[np.ndarray]: poses = [keyframe_c2w_init[0]] for idx in range(0, len(x), 6): poses.append(pose_vec_to_matrix(x[idx : idx + 6])) return poses def edge_residual(edge: LoopEdge, poses: Sequence[np.ndarray], loop_scale: float) -> np.ndarray: pred = np.linalg.inv(poses[edge.dst_pos]) @ poses[edge.src_pos] rot_err = Rotation.from_matrix( edge.transform_ji[:3, :3].T @ pred[:3, :3] ).as_rotvec() trans_err = pred[:3, 3] - edge.transform_ji[:3, 3] if edge.method == "odometry": w = 1.0 else: w = loop_scale * math.sqrt(max(edge.inliers, 1) / max(cfg.rigid_min_inliers, 1)) return np.concatenate( [ rot_err * cfg.pose_graph_rot_weight * w, trans_err * cfg.pose_graph_trans_weight * w, ], axis=0, ) def residual_fn(x: np.ndarray) -> np.ndarray: poses = unpack_poses(x) residuals = [] for edge in odom_edges: residuals.append(edge_residual(edge, poses, loop_scale=1.0)) for edge in loop_edges: residuals.append(edge_residual(edge, poses, cfg.pose_graph_loop_weight)) return np.concatenate(residuals, axis=0) def build_jac_sparsity() -> lil_matrix: sparsity = lil_matrix((num_residuals, num_vars), dtype=np.int8) row = 0 for edge in list(odom_edges) + list(loop_edges): for node_idx in (edge.src_pos, edge.dst_pos): if node_idx <= 0: continue col0 = (node_idx - 1) * 6 sparsity[row : row + 6, col0 : col0 + 6] = 1 row += 6 return sparsity jac_sparsity = build_jac_sparsity() _log( "Pose graph: built sparse Jacobian pattern for finite differences", cfg.verbose, ) result = least_squares( residual_fn, x0, loss="huber", f_scale=1.0, max_nfev=cfg.pose_graph_max_nfev, jac_sparsity=jac_sparsity, x_scale="jac", verbose=2 if cfg.pose_graph_solver_verbose else 0, ) poses = unpack_poses(result.x) _log( f"Pose graph: done status={result.status} cost={float(result.cost):.4f} " f"nfev={int(result.nfev)}", cfg.verbose, ) return np.asarray(poses, dtype=np.float64) def propagate_keyframe_corrections( c2w_base: np.ndarray, keyframe_ids: Sequence[int], frame_ref_keyframe_pos: np.ndarray, keyframe_c2w_opt: np.ndarray, ) -> np.ndarray: corrected = c2w_base.copy() keyframe_base = c2w_base[np.asarray(keyframe_ids)] correction = np.asarray( [ keyframe_c2w_opt[idx] @ np.linalg.inv(keyframe_base[idx]) for idx in range(len(keyframe_ids)) ], dtype=np.float64, ) for frame_idx in range(len(c2w_base)): ref_pos = int(frame_ref_keyframe_pos[frame_idx]) corrected[frame_idx] = normalize_pose_matrix(correction[ref_pos] @ c2w_base[frame_idx]) return corrected def frame_to_keyframe_position( keyframe_mask: np.ndarray, keyframe_indices: np.ndarray ) -> np.ndarray: key_positions = np.cumsum(keyframe_mask.astype(np.int64)) - 1 out = np.zeros_like(keyframe_indices, dtype=np.int64) for idx in range(len(keyframe_indices)): if keyframe_mask[idx]: out[idx] = key_positions[idx] else: ref_frame = int(keyframe_indices[idx]) out[idx] = key_positions[ref_frame] return out def save_w2c_txt(path: str, w2c: np.ndarray, frame_ids: Sequence[int]) -> None: os.makedirs(os.path.dirname(path), exist_ok=True) with open(path, "w") as f: f.write("# w2c\n") for idx, frame_id in enumerate(frame_ids): row = [frame_id] + w2c[idx, :3, :3].reshape(-1).tolist() + w2c[idx, :3, 3].tolist() f.write(" ".join(map(str, row)) + "\n") def save_loop_edges_json(path: str, edges: Sequence[LoopEdge]) -> None: os.makedirs(os.path.dirname(path), exist_ok=True) payload = [] for edge in edges: item = asdict(edge) item["transform_ji"] = edge.transform_ji.tolist() payload.append(item) with open(path, "w") as f: json.dump(payload, f, indent=2) def _xy_from_xyz(xyz: np.ndarray) -> np.ndarray: xyz = np.asarray(xyz, dtype=np.float64) return xyz[:, [0, 2]] def plot_trajectory_panels( output_path: str, baseline_c2w: np.ndarray, method_results: Dict[str, np.ndarray], loop_edges: Dict[str, Sequence[LoopEdge]], frame_ids: Sequence[int], gt_metrics: Optional[Dict[str, Dict[str, object]]] = None, selected_loop_weights: Optional[Dict[str, float]] = None, ) -> None: os.makedirs(os.path.dirname(output_path), exist_ok=True) extra_names = [name for name in method_results.keys() if name != "no_loop"] names = ["no_loop", *extra_names] titles = { "no_loop": "No Loop", "dbow": "ORB-BoW", "salad": "SALAD", } fig, axes = plt.subplots(1, len(names), figsize=(6 * len(names), 6)) if len(names) == 1: axes = [axes] base_metrics = (gt_metrics or {}).get("no_loop") if base_metrics is not None: base_xy = _xy_from_xyz(np.asarray(base_metrics["pred_xyz_aligned"], dtype=np.float64)) else: base_xy = _top_down_xy(baseline_c2w) for ax, name in zip(axes, names): method_metric = (gt_metrics or {}).get(name) if method_metric is not None: xy = _xy_from_xyz(np.asarray(method_metric["pred_xyz_aligned"], dtype=np.float64)) gt_xy = _xy_from_xyz(np.asarray(method_metric["gt_xyz"], dtype=np.float64)) frame_id_to_xy = { int(frame_id): xy[idx] for idx, frame_id in enumerate(method_metric["frame_ids"]) } ax.plot(gt_xy[:, 0], gt_xy[:, 1], color="#2ca02c", linewidth=2.0, label="gt") if name != "no_loop": ax.plot( base_xy[:, 0], base_xy[:, 1], color="#808080", linewidth=1.6, linestyle="--", label="no_loop", ) else: traj = method_results.get(name, baseline_c2w) xy = _top_down_xy(traj) gt_xy = None frame_id_to_xy = { int(frame_id): xy[idx] for idx, frame_id in enumerate(frame_ids) } if name != "no_loop": ax.plot( base_xy[:, 0], base_xy[:, 1], color="#808080", linewidth=1.6, linestyle="--", label="no_loop", ) ax.plot(xy[:, 0], xy[:, 1], color="#1f77b4", linewidth=2.0, label=name) if name != "no_loop": for edge in loop_edges.get(name, []): a = frame_id_to_xy.get(int(edge.src_frame)) b = frame_id_to_xy.get(int(edge.dst_frame)) if a is None or b is None: continue ax.plot( [a[0], b[0]], [a[1], b[1]], color="#d62728", alpha=0.35, linewidth=1.0, ) ax.set_aspect("equal") title = titles.get(name, name) if method_metric is not None: title += f"\nATE {float(method_metric['ate_rmse']):.2f} m" if name != "no_loop" and selected_loop_weights is not None and name in selected_loop_weights: title += f", w={float(selected_loop_weights[name]):.2f}" ax.set_title(title) ax.set_xlabel("x") ax.set_ylabel("z") ax.grid(True, alpha=0.3) ax.legend(loc="best", fontsize=8) plt.tight_layout() plt.savefig(output_path, dpi=220, bbox_inches="tight") plt.close(fig) def plot_overlay( output_path: str, baseline_c2w: np.ndarray, method_results: Dict[str, np.ndarray], gt_metrics: Optional[Dict[str, Dict[str, object]]] = None, ) -> None: os.makedirs(os.path.dirname(output_path), exist_ok=True) plt.figure(figsize=(7, 7)) if gt_metrics and gt_metrics.get("no_loop") is not None: gt_xy = _xy_from_xyz( np.asarray(gt_metrics["no_loop"]["gt_xyz"], dtype=np.float64) ) plt.plot(gt_xy[:, 0], gt_xy[:, 1], label="gt", linewidth=2.0, color="#2ca02c") color_map = { "no_loop": "#808080", "dbow": "#1f77b4", "salad": "#9467bd", } for name in ["no_loop", *[key for key in method_results.keys() if key != "no_loop"]]: metric = gt_metrics.get(name) if metric is None: continue xy = _xy_from_xyz(np.asarray(metric["pred_xyz_aligned"], dtype=np.float64)) plt.plot( xy[:, 0], xy[:, 1], label=name, linewidth=2.0, color=color_map.get(name, None), ) else: entries = [(baseline_c2w, "no_loop", "#808080")] if "dbow" in method_results: entries.append((method_results["dbow"], "orb_bow", "#1f77b4")) if "salad" in method_results: entries.append((method_results["salad"], "salad", "#9467bd")) for c2w, label, color in entries: xy = _top_down_xy(c2w) plt.plot(xy[:, 0], xy[:, 1], label=label, linewidth=2.0, color=color) plt.gca().set_aspect("equal") plt.xlabel("x") plt.ylabel("z") plt.title("Offline Loop Comparison") plt.grid(True, alpha=0.3) plt.legend() plt.tight_layout() plt.savefig(output_path, dpi=220, bbox_inches="tight") plt.close() def summarize_loop_edges( method: str, num_candidates: int, edges: Sequence[LoopEdge], trajectory_path: str, loop_plot_path: str, ) -> PoseGraphSummary: return PoseGraphSummary( method=method, num_candidates=num_candidates, num_verified_loops=len(edges), mean_loop_score=float(np.mean([edge.score for edge in edges])) if edges else 0.0, mean_loop_inliers=float(np.mean([edge.inliers for edge in edges])) if edges else 0.0, trajectory_path=trajectory_path, loop_plot_path=loop_plot_path, )