Skip to main content

kornia_slam/system/
mod.rs

1//! SLAM runtime: orchestrates tracking, mapping, and state transitions.
2//!
3//! The runtime flow is kept in one file so it can be read from top to bottom
4//! in the same order frames move through the system.
5
6mod config;
7mod inertial;
8mod state;
9
10pub use config::SlamConfig;
11pub use state::{TrackingResult, TrackingStatus};
12
13use inertial::{AppliedInitialization, InertialState, viba0_accel_bias_prior};
14use state::{SystemMode, SystemState};
15
16use crate::tracking::{KeyframePolicy, TrackingLossRecoveryPolicy};
17
18use crate::tracking::local_map::{LocalMapSelectionConfig, select_local_map_points};
19use crate::tracking::motion::{InertialPrediction, predict_pose};
20use crate::tracking::tracker::{FrameInput, Tracker};
21
22use std::sync::{Arc, Mutex};
23
24use crate::Frame;
25use crate::initialization::bootstrap::{
26    self, BootstrapDecision, MIN_KEYPOINTS_FOR_BOOTSTRAP, MIN_STEREO_POINTS, MIN_VALID_POINTS,
27    TooFewStereoPoints, evaluate_bootstrap,
28};
29use crate::initialization::two_view::TwoViewInitConfig;
30use crate::loop_closure::place_recognition::Vocabulary;
31use crate::loop_closure::{LoopCloser, LoopClosingContext, LoopClosureEvent};
32use crate::mapping::keyframe_mapping::{self, KeyframeGrowthResult, KeyframeInsertion};
33use crate::mapping::map::{Keyframe, Map, MapInsertion, MapMutationError, MapPoint};
34use crate::mapping::{KeyframeJob, LocalMapping};
35use crate::pose_conversion::apply_reference_pose_correction;
36use crate::sensor_rig::{ImuCalibration, SensorRig};
37use kornia_3d::camera::PinholeCamera;
38use kornia_3d::pose::Pose3d;
39use kornia_image::Image;
40use kornia_sensors::imu::ImuMeasurement;
41
42/// Top-level ORB-SLAM system: orchestrates tracking, mapping, and state transitions.
43pub struct SlamSystem {
44    // Camera model
45    rig: SensorRig,
46    // Primary pose estimator
47    tracker: Tracker,
48    // Boostrap pose estimator
49    two_view_init_config: TwoViewInitConfig,
50    // Keyframe insertion policy
51    keyframe_policy: KeyframePolicy,
52    // Recently-lost grace period policy
53    tracking_loss_recovery: TrackingLossRecoveryPolicy,
54    // mThDepth (metres): back-project close stereo points at each keyframe when set
55    stereo_close_depth: Option<f64>,
56    // Emit per-frame diagnostic logs (skip/reject reasons, growth counters)
57    debug: bool,
58    // Buffered debug messages produced during the most recent process_frame call;
59    // drained by the caller (TUI panel or stderr).
60    debug_messages: Vec<String>,
61    // Map object
62    map: Arc<Mutex<Map>>,
63    // Serializes compound map publication and short local-BA snapshot/merge phases.
64    map_publication_gate: Option<Arc<Mutex<()>>>,
65    // IMU estimates, sample buffer and the inertial-initialization schedule
66    inertial: InertialState,
67    local_mapping: LocalMapping,
68    // Place recognition for every keyframe, and loop closing when configured.
69    loop_closer: LoopCloser,
70    loop_closure_events: Vec<LoopClosureEvent>,
71    // System state
72    state: SystemState,
73}
74
75impl SlamSystem {
76    /// Creates a new system with identity pose.
77    pub fn new(camera: PinholeCamera, config: SlamConfig) -> Self {
78        Self::with_rig(SensorRig::new(camera), config)
79    }
80
81    /// Builds a system from an explicit sensor rig, so IMU noise parameters and
82    /// extrinsics can be supplied for the actual sensor.
83    pub fn with_rig(rig: SensorRig, config: SlamConfig) -> Self {
84        let camera = rig.camera.clone();
85        let map = Arc::new(Mutex::new(Map::new()));
86        let local_mapping =
87            LocalMapping::new(config.local_mapping, Arc::clone(&map), camera.clone());
88        let map_publication_gate = local_mapping.publication_gate();
89        let loop_closer = LoopCloser::new(config.pgo);
90        Self {
91            rig,
92            tracker: Tracker::new(config.map_projection),
93            two_view_init_config: config.two_view_init,
94            keyframe_policy: config.keyframe_policy,
95            tracking_loss_recovery: config.tracking_loss_recovery,
96            stereo_close_depth: config.stereo_close_depth_m,
97            debug: config.debug,
98            debug_messages: Vec::new(),
99            map,
100            map_publication_gate,
101            local_mapping,
102            state: SystemState::new(),
103            inertial: InertialState::new(),
104            loop_closer,
105            loop_closure_events: Vec::new(),
106        }
107    }
108
109    /// Enables appearance-based loop detection with a bag-of-words vocabulary.
110    /// Without it, keyframes are not indexed and no loop candidates are emitted.
111    pub fn set_vocabulary(&mut self, vocabulary: Vocabulary) {
112        self.loop_closer.set_vocabulary(vocabulary);
113    }
114
115    pub fn drain_loop_closure_events(&mut self) -> Vec<LoopClosureEvent> {
116        std::mem::take(&mut self.loop_closure_events)
117    }
118
119    /// Enables the inertial path by providing the camera-to-body extrinsic
120    /// `T_BC` (`X_body = T_BC * X_cam`). Without it, IMU samples are ignored.
121    pub fn set_imu_extrinsics(&mut self, t_bc: Pose3d) {
122        self.rig.imu = Some(ImuCalibration::new(t_bc));
123    }
124
125    /// The system's fixed sensor calibration.
126    pub fn rig(&self) -> &SensorRig {
127        &self.rig
128    }
129
130    /// Processes one frame (pre-extracted features) and returns the tracking result.
131    pub fn process_frame(
132        &mut self,
133        mut frame: Frame,
134        previous_image: Option<&Image<u8, 1>>,
135        current_image: &Image<u8, 1>,
136        timestamp_sec: f64,
137        imu_samples: Vec<ImuMeasurement>,
138    ) -> TrackingResult {
139        // Local-BA snapshots, merges, and their correction messages can only
140        // cross this boundary between complete tracking frames.
141        let publication_gate = self.map_publication_gate.clone();
142        let _publication = publication_gate
143            .as_ref()
144            .map(|gate| gate.lock().unwrap_or_else(|poisoned| poisoned.into_inner()));
145        self.apply_local_mapping_results();
146        // Fill the per-frame undistortion cache once; tracking, BA gathering,
147        // growth, and fuse all read from it.
148        frame.ensure_undistorted(&self.rig.camera);
149        self.inertial.buffer_samples(imu_samples);
150
151        match self.state.mode {
152            SystemMode::Bootstrap => self.bootstrap_step(frame, timestamp_sec),
153            SystemMode::ImuInit => {
154                self.inertial_init_step(frame, previous_image, current_image, timestamp_sec)
155            }
156            SystemMode::Tracking => {
157                self.tracking_step(frame, previous_image, current_image, timestamp_sec)
158            }
159        }
160    }
161
162    /// Runs `f` against the live map points, holding the map lock only for the
163    /// duration of the call. Avoids cloning the whole point list (descriptors
164    /// included) for read-only consumers such as viz logging and summaries.
165    pub fn with_map_points<R>(&self, f: impl FnOnce(&[MapPoint]) -> R) -> R {
166        f(self.map.lock().unwrap().map_points())
167    }
168
169    /// Returns the index of the current reference keyframe, if tracking has one.
170    pub fn current_keyframe_idx(&self) -> Option<usize> {
171        self.state.current_keyframe_idx.and_then(|ki| {
172            self.map
173                .lock()
174                .unwrap()
175                .get_keyframe(ki)
176                .map(|kf| kf.frame.idx)
177        })
178    }
179
180    /// Returns the number of active (non-culled) map points.
181    pub fn num_active_map_points(&self) -> usize {
182        self.map.lock().unwrap().num_active_map_points()
183    }
184
185    /// Drain any debug messages accumulated since the last call.
186    pub fn drain_debug_messages(&mut self) -> Vec<String> {
187        std::mem::take(&mut self.debug_messages)
188    }
189
190    /// Toggle whether the pipeline buffers per-frame debug messages.
191    pub fn set_debug(&mut self, on: bool) {
192        self.debug = on;
193        if !on {
194            self.debug_messages.clear();
195        }
196    }
197
198    fn apply_local_mapping_results(&mut self) {
199        let Some(reference_idx) = self.state.current_keyframe_idx else {
200            // Still drain results so a completed worker cannot build a result backlog.
201            let _ = self.local_mapping.drain_results();
202            return;
203        };
204
205        for result in self.local_mapping.drain_results() {
206            let Some(correction) = result
207                .keyframe_corrections
208                .iter()
209                .find(|correction| correction.kf_idx == reference_idx)
210            else {
211                continue;
212            };
213
214            self.state.pose_world_to_cam = apply_reference_pose_correction(
215                self.state.pose_world_to_cam,
216                correction.pose_before,
217                correction.pose_after,
218            );
219            self.state.velocity_world = correction.velocity_world;
220            self.inertial.bias = correction.imu_bias;
221        }
222    }
223
224    fn dbg(&mut self, msg: String) {
225        if self.debug {
226            self.debug_messages.push(msg);
227        }
228    }
229
230    fn bootstrap_step(&mut self, curr_frame: Frame, timestamp_sec: f64) -> TrackingResult {
231        // Stereo frames carry metric per-keypoint depth, so we can build a
232        // metric map from a single keyframe (ORB-SLAM3's StereoInitialization)
233        // instead of waiting for two-view parallax.
234        if curr_frame.is_stereo() {
235            return self.bootstrap_stereo(curr_frame, timestamp_sec);
236        }
237        self.bootstrap_mono(curr_frame, timestamp_sec)
238    }
239
240    /// Single-frame metric initialization from stereo depth.
241    fn bootstrap_stereo(&mut self, mut curr_frame: Frame, timestamp_sec: f64) -> TrackingResult {
242        // Build the new map in the current odometry frame (identity at start,
243        // or the recovery pose after a tracking loss).
244        curr_frame.pose_world_to_cam = self.state.pose_world_to_cam;
245        let curr_idx = curr_frame.idx;
246
247        let insertion =
248            match bootstrap::stereo_initial_map(Keyframe::from_frame(curr_frame), &self.rig.camera)
249            {
250                Ok(insertion) => insertion,
251                Err(TooFewStereoPoints { found }) => {
252                    self.dbg(format!(
253                        "[bootstrap_stereo] frame={curr_idx} skip: only {found} stereo points \
254                     (need >= {MIN_STEREO_POINTS})"
255                    ));
256                    return self.frame_result(TrackingStatus::Skipped);
257                }
258            };
259
260        // The keyframe and its seeds are published together; tracker state is
261        // adopted below only once that succeeded.
262        let published = self.map.lock().unwrap().apply_insertion(insertion);
263        let added = match published {
264            Ok(result) => result.landmark_ids.len(),
265            Err(error) => {
266                self.dbg(format!(
267                    "[bootstrap_stereo] frame={curr_idx} publication rejected: {error}"
268                ));
269                return self.frame_result(TrackingStatus::Skipped);
270            }
271        };
272
273        self.dbg(format!(
274            "[bootstrap_stereo] frame={curr_idx} metric map created with {added} points",
275        ));
276
277        self.state.velocity = None;
278        // The map is already metric (stereo baseline), but gravity, velocities,
279        // and the gyro bias still need the inertial init before IMU prediction
280        // can run; the solve there keeps scale fixed at 1.
281        self.finish_bootstrap(curr_idx, timestamp_sec);
282        self.prune_imu_before(timestamp_sec);
283
284        self.frame_result(TrackingStatus::KeyframeAccepted)
285    }
286
287    fn bootstrap_mono(&mut self, mut curr_frame: Frame, timestamp_sec: f64) -> TrackingResult {
288        // Stamp frames with current odometry pose so bootstrap builds
289        // the new map in the existing coordinate frame.
290        curr_frame.pose_world_to_cam = self.state.pose_world_to_cam;
291
292        let decision = evaluate_bootstrap(
293            self.state.bootstrap_frame.as_ref(),
294            &curr_frame,
295            &self.rig.camera,
296            &self.two_view_init_config,
297        );
298        let two_view_estimate = match decision {
299            // A stale reference is dropped with the frame: neither is viable.
300            BootstrapDecision::Unusable { keypoints } => {
301                self.dbg(format!(
302                    "[bootstrap] frame={} skip: too few keypoints ({}, need > {})",
303                    curr_frame.idx, keypoints, MIN_KEYPOINTS_FOR_BOOTSTRAP,
304                ));
305                self.state.bootstrap_frame = None;
306                return self.frame_result(TrackingStatus::Skipped);
307            }
308            BootstrapDecision::StoreAsReference => {
309                self.dbg(format!(
310                    "[bootstrap] frame={} stored as reference (awaiting second frame)",
311                    curr_frame.idx,
312                ));
313                self.state.bootstrap_frame = Some(curr_frame);
314                self.inertial.bootstrap_timestamp_sec = Some(timestamp_sec);
315                // Samples before the reference frame can never enter an edge.
316                self.prune_imu_before(timestamp_sec);
317                return self.frame_result(TrackingStatus::Skipped);
318            }
319            // The reference is kept: only the second frame was unsuitable.
320            BootstrapDecision::Rejected {
321                reference_idx,
322                reason,
323            } => {
324                self.dbg(format!(
325                    "[bootstrap] frame={} (ref={}) reject: {:?}",
326                    curr_frame.idx, reference_idx, reason,
327                ));
328                return self.frame_result(TrackingStatus::Skipped);
329            }
330            BootstrapDecision::Initialized(estimate) => *estimate,
331        };
332        let prev_bootstrap_frame = self
333            .state
334            .bootstrap_frame
335            .take()
336            .expect("an accepted bootstrap consumes the reference it matched against");
337
338        self.dbg(format!(
339            "[bootstrap] frame={} accept: model={} triangulated={} inliers={}",
340            curr_frame.idx,
341            two_view_estimate.model_kind,
342            two_view_estimate.points3d.len(),
343            two_view_estimate.estimate.inliers,
344        ));
345
346        let estimated_pose = two_view_estimate.estimate.pose;
347        let prev_pose_world_to_cam = curr_frame.pose_world_to_cam;
348        self.state.pose_world_to_cam = estimated_pose;
349        curr_frame.pose_world_to_cam = estimated_pose;
350
351        // Promote to Keyframes
352        let prev_idx = prev_bootstrap_frame.idx;
353        let reference_kf = Keyframe::from_frame(prev_bootstrap_frame);
354        let current_kf = Keyframe::from_frame(curr_frame);
355        let curr_idx = current_kf.frame.idx;
356
357        // A rejected publication writes nothing, so there is no map to
358        // evaluate; treat it like a failed health gate.
359        let insertion =
360            bootstrap::two_view_initial_map(reference_kf, current_kf, &two_view_estimate);
361        if let Err(error) = self.publish_initial_map(insertion, [prev_idx, curr_idx]) {
362            self.dbg(format!(
363                "[bootstrap] frame={curr_idx} publication rejected: {error}"
364            ));
365            self.state.reset();
366            return self.frame_result(TrackingStatus::Skipped);
367        }
368
369        // The gate runs after initial BA, which can collapse a badly
370        // conditioned pair.
371        let health = crate::initialization::initial_map_health(&self.map.lock().unwrap());
372        if health.is_degenerate() {
373            self.dbg(format!(
374                "[init_gate] reject: valid_in_both={} median_depth={:.3} (need >= {} and > 0)",
375                health.valid_in_both, health.median_depth_older_kf, MIN_VALID_POINTS,
376            ));
377            self.map.lock().unwrap().clear_active();
378            self.state.reset();
379            return self.frame_result(TrackingStatus::Skipped);
380        }
381
382        // Initial BA in publish_initial_map may have refined KF1's pose; sync state
383        // and recompute velocity from the post-BA pose.
384        if let Some(kf) = self.map.lock().unwrap().get_keyframe(curr_idx) {
385            self.state.pose_world_to_cam = kf.frame.pose_world_to_cam;
386        }
387
388        if let Some(prev_ts) = self.inertial.bootstrap_timestamp_sec {
389            if let Some(edge) = self.inertial.keyframe_edge(
390                self.rig.imu_noise(),
391                prev_idx,
392                curr_idx,
393                prev_ts,
394                timestamp_sec,
395            ) {
396                // Through the validated path: endpoints must exist, the edge
397                // must be new, and the interval finite and ordered.
398                let published = self.map.lock().unwrap().apply_insertion(MapInsertion {
399                    imu_factors: vec![edge],
400                    ..Default::default()
401                });
402                if let Err(error) = published {
403                    self.dbg(format!(
404                        "[bootstrap] frame={curr_idx} imu edge rejected: {error}"
405                    ));
406                }
407            }
408            self.prune_imu_before(timestamp_sec);
409        }
410
411        self.state.velocity = Some(Pose3d::between(
412            &prev_pose_world_to_cam,
413            &self.state.pose_world_to_cam,
414        ));
415
416        self.finish_bootstrap(curr_idx, timestamp_sec);
417
418        self.frame_result(TrackingStatus::KeyframeAccepted)
419    }
420
421    /// Publishes the two-view map, refines it with initial BA, and offers both
422    /// keyframes to place recognition so a later revisit of the start can
423    /// match them.
424    fn publish_initial_map(
425        &mut self,
426        insertion: MapInsertion,
427        keyframes: [usize; 2],
428    ) -> Result<(), MapMutationError> {
429        self.map.lock().unwrap().apply_insertion(insertion)?;
430        let refined = crate::mapping::bundle_adjustment::run_initial_ba(
431            &mut self.map.lock().unwrap(),
432            &self.rig.camera,
433        );
434        if let Err(error) = refined {
435            self.dbg(format!("[init_ba] {error}; keeping the unrefined map"));
436        }
437        for kf_idx in keyframes {
438            self.register_place_recognition(kf_idx);
439        }
440        Ok(())
441    }
442
443    /// Leaves bootstrap with `kf_idx` as the reference keyframe. Inertial
444    /// initialization follows when the camera-to-body extrinsic is known, since
445    /// it relates IMU deltas to camera poses; otherwise tracking is visual-only.
446    fn finish_bootstrap(&mut self, kf_idx: usize, timestamp_sec: f64) {
447        self.state.current_keyframe_idx = Some(kf_idx);
448        self.state.last_keyframe_idx = Some(kf_idx);
449        self.state.mode = if self.rig.camera_to_body().is_some() {
450            self.inertial.schedule.start(kf_idx, timestamp_sec);
451            SystemMode::ImuInit
452        } else {
453            SystemMode::Tracking
454        };
455        self.inertial.last_keyframe_timestamp_sec = Some(timestamp_sec);
456    }
457
458    fn frame_result(&self, status: TrackingStatus) -> TrackingResult {
459        TrackingResult {
460            pose_world_to_cam: self.state.pose_world_to_cam,
461            status,
462        }
463    }
464
465    /// Drops buffered IMU samples strictly older than `t` (typically the last
466    /// keyframe timestamp: the next edge and all per-frame windows start there).
467    fn prune_imu_before(&mut self, t: f64) {
468        self.inertial.prune_before(t);
469    }
470
471    /// Tracks the frame, then attempts VIBA0 on accepted keyframes once the
472    /// window is ready and a retry is due.
473    fn inertial_init_step(
474        &mut self,
475        frame: Frame,
476        previous_image: Option<&Image<u8, 1>>,
477        current_image: &Image<u8, 1>,
478        timestamp_sec: f64,
479    ) -> TrackingResult {
480        let result = self.tracking_step(frame, previous_image, current_image, timestamp_sec);
481
482        if result.status == TrackingStatus::KeyframeAccepted
483            && let Some(start_idx) = self.inertial.schedule.start_kf_idx()
484        {
485            // Snapshot the fields needed after releasing the map lock.
486            let kfs: Vec<usize> = self
487                .map
488                .lock()
489                .unwrap()
490                .keyframes()
491                .iter()
492                .filter(|kf| kf.frame.idx >= start_idx)
493                .map(|kf| kf.frame.idx)
494                .collect();
495            let imu_time: f64 = self
496                .map
497                .lock()
498                .unwrap()
499                .imu_factors()
500                .iter()
501                .filter(|f| f.curr_kf_idx >= start_idx)
502                .map(|f| f.preintegrated.dt)
503                .sum();
504            let gate_msg = format_imu_init_gate(
505                start_idx,
506                kfs.first().copied(),
507                kfs.last().copied(),
508                kfs.len(),
509                self.inertial.initializer.config.min_keyframes,
510                imu_time,
511                self.inertial.initializer.config.min_time_sec,
512            );
513            self.dbg(gate_msg);
514        }
515
516        let due_for_retry = self.inertial.schedule.retry_due(timestamp_sec);
517        let imu_init_ready = self.inertial.initializer.ready(
518            &self.map.lock().unwrap(),
519            self.inertial.schedule.start_kf_idx(),
520        );
521
522        if result.status == TrackingStatus::KeyframeAccepted && due_for_retry && imu_init_ready {
523            let Some(start_idx) = self.inertial.schedule.start_kf_idx() else {
524                return result;
525            };
526            self.inertial.schedule.record_attempt(timestamp_sec);
527            let is_mono = !self
528                .map
529                .lock()
530                .unwrap()
531                .keyframes()
532                .iter()
533                .find(|kf| kf.frame.idx >= start_idx)
534                .map(|kf| kf.frame.is_stereo())
535                .unwrap_or(false);
536            let prior_a0 = viba0_accel_bias_prior(is_mono);
537            // Drop the solve's map lock before applying its result with a new lock.
538            let init_result = self.inertial.initializer.try_initialize(
539                &self.map.lock().unwrap(),
540                self.rig.camera_to_body(),
541                self.inertial.bias,
542                start_idx,
543                1e2,
544                prior_a0,
545                false,
546            );
547            match init_result {
548                Some(init) => {
549                    let applied = self.inertial.apply_initialization(
550                        &mut self.map.lock().unwrap(),
551                        &mut self.state,
552                        init,
553                        start_idx,
554                    );
555                    match applied {
556                        Ok(applied) => self.resume_tracking_after_viba0(timestamp_sec, applied),
557                        // Staying in ImuInit lets the retry throttle re-solve later.
558                        Err(error) => {
559                            self.dbg(format!("[imu_init] VIBA0 alignment refused: {error}"));
560                        }
561                    }
562                }
563                None => {
564                    self.dbg("[imu_init] VIBA0 rejected: solve failed or invalid scale".into());
565                }
566            }
567        }
568
569        result
570    }
571
572    fn resume_tracking_after_viba0(&mut self, timestamp_sec: f64, applied: AppliedInitialization) {
573        // Mirrors ORB-SLAM3: IMU is marked initialized (and tracking resumes)
574        // immediately after VIBA0 succeeds — VIBA1/VIBA2 refine
575        // bg/ba/scale/gravity further in the background (see
576        // try_insert_keyframe), they don't gate resuming tracking.
577        self.state.mode = SystemMode::Tracking;
578        self.state.imu_init_timestamp_sec = Some(timestamp_sec);
579        if !self.local_mapping.submit(KeyframeJob {
580            imu_initialized: true,
581            imu_t_bc: self.rig.camera_to_body(),
582            gravity_world: self.inertial.gravity_world,
583        }) {
584            self.dbg("[local_mapping] worker is unavailable".into());
585        }
586        self.apply_local_mapping_results();
587        let AppliedInitialization {
588            scale,
589            gravity_world: gravity,
590            gyro_bias: bg,
591        } = applied;
592        self.dbg(format!(
593            "[imu_init] VIBA0 accepted: scale={scale:.4} gravity=({:.3},{:.3},{:.3}) \
594             gyro_bias=({:.4},{:.4},{:.4})",
595            gravity.x, gravity.y, gravity.z, bg.x, bg.y, bg.z
596        ));
597    }
598
599    fn tracking_step(
600        &mut self,
601        frame: Frame,
602        previous_image: Option<&Image<u8, 1>>,
603        current_image: &Image<u8, 1>,
604        timestamp_sec: f64,
605    ) -> TrackingResult {
606        let image_size = frame.image_size;
607        let pose_before = self.state.pose_world_to_cam;
608        let prev_timestamp = self.state.last_frame_timestamp_sec;
609
610        // Local BA updates keyframe state asynchronously, so refresh cached IMU state.
611        if let Some(kf_idx) = self.state.current_keyframe_idx
612            && let Some(kf) = self.map.lock().unwrap().get_keyframe(kf_idx)
613        {
614            self.state.velocity_world = kf.velocity_world;
615            self.inertial.bias = kf.imu_bias;
616        }
617
618        // Preintegration is prepared here: the system owns bias and the sample
619        // buffer. The motion model chooses between it and the visual model.
620        let preintegrated = (self.state.imu_initialized && prev_timestamp > 0.0).then(|| {
621            self.inertial
622                .preintegrate_window(self.rig.imu_noise(), prev_timestamp, timestamp_sec)
623                .0
624        });
625        let (candidate_pose, predicted_velocity) = predict_pose(
626            pose_before,
627            self.state.velocity,
628            self.state.velocity_world,
629            preintegrated
630                .as_ref()
631                .map(|preintegrated| InertialPrediction {
632                    preintegrated,
633                    camera_to_body: self.rig.camera_to_body(),
634                    gravity_world: self.inertial.gravity_world,
635                }),
636        );
637        if let Some(velocity_world) = predicted_velocity {
638            self.state.velocity_world = velocity_world; // propagate for next frame
639        }
640
641        let currently_lost_for = self
642            .state
643            .lost_since_sec
644            .map_or(0.0, |t0| timestamp_sec - t0);
645        let result = {
646            let map = self.map.lock().unwrap();
647            self.tracker.estimate(
648                FrameInput {
649                    frame: &frame,
650                    previous_image,
651                    current_image,
652                    candidate_pose,
653                    pose_before,
654                    current_keyframe_idx: self.state.current_keyframe_idx,
655                    lost_for_sec: currently_lost_for,
656                },
657                &map,
658                &self.rig.camera,
659            )
660        };
661
662        let (mut status, matches, tracked_inliers, reject_reason) = match result {
663            Ok(estimate) => {
664                self.state.pose_world_to_cam = estimate.pose;
665
666                // When IMU-initialized, velocity_world was already updated by IMU
667                // preintegration (predict_pose_imu → pred_vel) before PnP ran; don't
668                // overwrite it with a visual finite-difference, since at 30 fps the
669                // inter-frame displacement is noise-dominated during low-translation
670                // segments, which would collapse velocity to zero and permanently
671                // freeze the IMU pose prediction.
672                if !self.state.imu_initialized {
673                    self.state.velocity = Some(Pose3d::between(&pose_before, &estimate.pose));
674                }
675
676                (
677                    TrackingStatus::Tracked,
678                    estimate.matches,
679                    estimate.inliers,
680                    None,
681                )
682            }
683            Err(reason) => {
684                // Carry the predicted pose forward instead of freezing at
685                // pose_before. state.velocity_world was already advanced by
686                // predict_pose_imu above regardless of visual outcome, so
687                // anchoring the next frame's prediction on a stale position
688                // would desync position/rotation from velocity: every
689                // subsequent frame's candidate pose would drift further from
690                // reality, making the projection search miss again and
691                // compounding a single bad frame into a full tracking loss.
692                self.state.pose_world_to_cam = candidate_pose;
693                (TrackingStatus::Skipped, Vec::new(), 0, Some(reason))
694            }
695        };
696        if self.debug {
697            let msg = match reject_reason {
698                Some(reason) => format!("[track] frame={} reject: {:?}", frame.idx, reason),
699                None => format!(
700                    "[track] frame={} ok: matches={} inliers={}",
701                    frame.idx,
702                    matches.len(),
703                    tracked_inliers,
704                ),
705            };
706            self.debug_messages.push(msg);
707        }
708
709        if status == TrackingStatus::Tracked {
710            // Visibility bookkeeping over the local map only (mirrors
711            // ORB-SLAM3, which counts mnVisible on local-map points): full-map
712            // scans here would grow with trajectory length.
713            // Release this non-reentrant lock before keyframe insertion locks the map.
714            {
715                let mut map_guard = self.map.lock().unwrap();
716                let current_kf = self
717                    .state
718                    .current_keyframe_idx
719                    .and_then(|ki| map_guard.get_keyframe(ki));
720                let local_indices = select_local_map_points(
721                    &map_guard,
722                    &matches,
723                    current_kf,
724                    LocalMapSelectionConfig::default(),
725                );
726                let visible = crate::tracking::local_map::landmarks_in_frustum(
727                    &map_guard,
728                    &local_indices,
729                    &self.rig.camera,
730                    &candidate_pose,
731                    image_size,
732                );
733                map_guard.update_observation_counts(&visible, &matches);
734            }
735
736            if self.try_insert_keyframe(&frame, timestamp_sec, tracked_inliers, &matches) {
737                status = TrackingStatus::KeyframeAccepted;
738            }
739        }
740
741        if status == TrackingStatus::Skipped {
742            let policy = &self.tracking_loss_recovery;
743
744            let lost_since = *self.state.lost_since_sec.get_or_insert(timestamp_sec);
745            let recently_lost_for = timestamp_sec - lost_since;
746
747            let imu_confident = self.state.imu_initialized
748                && self
749                    .state
750                    .imu_init_timestamp_sec
751                    .is_some_and(|t0| timestamp_sec - t0 >= policy.min_imu_confidence_sec);
752            let grace_period_sec = policy.grace_period_sec(imu_confident);
753            let map_established =
754                self.map.lock().unwrap().keyframes().len() > policy.min_keyframes_for_grace;
755
756            if !map_established || recently_lost_for >= grace_period_sec {
757                self.dbg(format!(
758                    "[lost] frame={} giving up after {:.2}s (map_established={}): resetting",
759                    frame.idx, recently_lost_for, map_established,
760                ));
761                self.tracker.reset_tracks();
762                self.state.reset();
763                return self.bootstrap_step(frame, timestamp_sec);
764            }
765        } else {
766            self.state.lost_since_sec = None;
767        }
768        self.state.last_frame_timestamp_sec = timestamp_sec;
769        // Samples older than the last keyframe can't enter any future window
770        // (the next edge and all per-frame predictions start at or after it).
771        if let Some(kf_ts) = self.inertial.last_keyframe_timestamp_sec {
772            self.prune_imu_before(kf_ts.min(timestamp_sec));
773        }
774        self.frame_result(status)
775    }
776
777    fn try_insert_keyframe(
778        &mut self,
779        frame: &Frame,
780        timestamp_sec: f64,
781        tracked_inliers: usize,
782        matches: &[(usize, usize)],
783    ) -> bool {
784        let n_ref_map_points = if let Some(ki) = self.state.current_keyframe_idx {
785            let map = self.map.lock().unwrap();
786            map.get_keyframe(ki)
787                .map(|kf| kf.num_associated_points())
788                .unwrap_or(0)
789        } else {
790            0
791        };
792
793        if !self.keyframe_policy.should_insert(
794            frame.idx,
795            self.state.last_keyframe_idx,
796            tracked_inliers,
797            n_ref_map_points,
798        ) {
799            return false;
800        }
801
802        // Guard: reference KF must exist before we can triangulate.
803        if let Some(ki) = self.state.current_keyframe_idx {
804            let map = self.map.lock().unwrap();
805            if map.get_keyframe(ki).is_none() {
806                return false;
807            }
808        } else {
809            return false;
810        }
811
812        let mut curr_kf = Keyframe::from_frame(Frame {
813            pose_world_to_cam: self.state.pose_world_to_cam,
814            ..frame.clone()
815        });
816        // Seed the new keyframe with the IMU-propagated velocity and current bias
817        // estimate so that VI-BA starts from a reasonable linearisation point rather
818        // than zero, which would produce huge residuals on the newest IMU edge.
819        if self.state.imu_initialized {
820            curr_kf.velocity_world = self.state.velocity_world;
821            curr_kf.imu_bias = self.inertial.bias;
822        }
823        let imu_initialized = self.state.imu_initialized;
824        let is_stereo = curr_kf.frame.is_stereo();
825
826        // Neighbours are captured BEFORE publication: growing against a list
827        // that already contains the current keyframe would triangulate it
828        // against itself and drop the oldest real neighbour. Selection policy
829        // belongs to mapping.
830        let growth_plan =
831            keyframe_mapping::prepare_keyframe_growth(&self.map.lock().unwrap(), frame.idx);
832
833        let imu_edge = match (
834            self.state.last_keyframe_idx,
835            self.inertial.last_keyframe_timestamp_sec,
836        ) {
837            (Some(prev_kf_idx), Some(prev_ts)) => self.inertial.keyframe_edge(
838                self.rig.imu_noise(),
839                prev_kf_idx,
840                frame.idx,
841                prev_ts,
842                timestamp_sec,
843            ),
844            _ => None,
845        };
846        let KeyframeInsertion {
847            insertion,
848            close_points: n_close,
849        } = keyframe_mapping::keyframe_insertion(
850            curr_kf,
851            matches,
852            self.stereo_close_depth,
853            &self.rig.camera,
854            imu_edge,
855        );
856
857        let published = self.map.lock().unwrap().apply_insertion(insertion);
858        if let Err(error) = published {
859            // Nothing was written, so no tracker or IMU reference may advance to
860            // a keyframe the map does not hold.
861            self.dbg(format!(
862                "[kf] frame={} publication rejected: {error}",
863                frame.idx
864            ));
865            return false;
866        }
867        if self.stereo_close_depth.is_some() && is_stereo {
868            self.dbg(format!(
869                "[kf_stereo] frame={} close_points={}",
870                frame.idx, n_close
871            ));
872        }
873
874        self.inertial.last_keyframe_timestamp_sec = Some(timestamp_sec);
875
876        self.state.current_keyframe_idx = Some(frame.idx);
877        self.state.last_keyframe_idx = Some(frame.idx);
878
879        // Optional mapping work the accepted keyframe earns: grow against the
880        // captured neighbours, then forward SearchInNeighbors / Fuse, before
881        // local BA so BA sees the extra constraints. Mapping owns the sequence
882        // and its lock boundaries; the messages stay here.
883        let growth = keyframe_mapping::grow_keyframe(
884            &self.map,
885            growth_plan,
886            &self.rig.camera,
887            self.two_view_init_config.match_config,
888            &self.two_view_init_config.triangulation_config,
889        );
890        self.report_growth(frame.idx, &growth);
891
892        // Refinement can rotate/scale the world and update gravity. Do it before
893        // constructing the BA request so the job and its future snapshot agree.
894        if imu_initialized {
895            self.refine_inertial_init(timestamp_sec);
896        }
897
898        if !self.local_mapping.submit(KeyframeJob {
899            imu_initialized,
900            imu_t_bc: self.rig.camera_to_body(),
901            gravity_world: self.inertial.gravity_world,
902        }) {
903            self.dbg("[local_mapping] worker is unavailable".into());
904        }
905        // Synchronous mode has a completed correction available immediately;
906        // asynchronous mode will deliver it at a later frame boundary.
907        self.apply_local_mapping_results();
908
909        // Index this keyframe for appearance-based place recognition and surface
910        // any loop candidates (no-op unless a vocabulary was provided).
911        self.register_place_recognition(frame.idx);
912
913        true
914    }
915
916    fn report_growth(&mut self, frame_idx: usize, growth: &KeyframeGrowthResult) {
917        if !self.debug {
918            return;
919        }
920        for failure in &growth.pair_failures {
921            self.dbg(format!(
922                "[kf] frame={frame_idx} pair growth against {} rejected: {}",
923                failure.neighbor_keyframe_idx, failure.error
924            ));
925        }
926        self.dbg(format!(
927            "[kf] frame={frame_idx} grown={} from {} neighbor kfs",
928            growth.landmarks_added, growth.neighbor_count
929        ));
930        self.dbg(format!(
931            "[fuse] frame={frame_idx} fused={}",
932            growth.observations_added
933        ));
934        for error in &growth.fusion_failures {
935            self.dbg(format!("[fuse] frame={frame_idx} link refused: {error}"));
936        }
937    }
938
939    /// VIBA1/VIBA2: progressive re-solves with relaxed priors over the same
940    /// (now-growing) window VIBA0 used, mirroring LocalMapping.cc:200-228.
941    /// Tracking already runs on VIBA0's result, so a rejection here only means
942    /// that stage never runs, not a tracking failure.
943    fn refine_inertial_init(&mut self, timestamp_sec: f64) {
944        let (Some(start_idx), Some(refinement)) = (
945            self.inertial.schedule.start_kf_idx(),
946            self.inertial.schedule.due_refinement(timestamp_sec),
947        ) else {
948            return;
949        };
950        let stage = refinement.name();
951
952        // Drop the solve's map lock before applying its result with a new lock.
953        let init_result = self.inertial.initializer.try_initialize(
954            &self.map.lock().unwrap(),
955            self.rig.camera_to_body(),
956            self.inertial.bias,
957            start_idx,
958            refinement.prior_g(),
959            refinement.prior_a(),
960            true,
961        );
962        match init_result {
963            Some(init) => {
964                let applied = self.inertial.apply_initialization(
965                    &mut self.map.lock().unwrap(),
966                    &mut self.state,
967                    init,
968                    start_idx,
969                );
970                match applied {
971                    Ok(AppliedInitialization {
972                        scale,
973                        gyro_bias: bg,
974                        ..
975                    }) => self.dbg(format!(
976                        "[imu_init] {stage} accepted: scale_correction={scale:.4} gyro_bias=({:.4},{:.4},{:.4})",
977                        bg.x, bg.y, bg.z
978                    )),
979                    Err(error) => {
980                        self.dbg(format!("[imu_init] {stage} alignment refused: {error}"));
981                    }
982                }
983            }
984            None => {
985                self.dbg(format!(
986                    "[imu_init] {stage} rejected: solve failed or invalid scale"
987                ));
988            }
989        }
990
991        self.inertial.schedule.complete(refinement);
992    }
993
994    /// Offers a freshly inserted keyframe to place recognition and, when
995    /// configured, loop closing; applies whatever the closer reports back.
996    fn register_place_recognition(&mut self, kf_idx: usize) {
997        let context = LoopClosingContext {
998            pose_world_to_cam: self.state.pose_world_to_cam,
999            velocity_world: self.state.velocity_world,
1000            current_keyframe_idx: self.state.current_keyframe_idx,
1001            imu_initialized: self.state.imu_initialized,
1002            gravity_world: self.inertial.gravity_world,
1003        };
1004        let outcome = {
1005            let mut map = self.map.lock().unwrap();
1006            self.loop_closer
1007                .on_keyframe(&mut map, &self.rig.camera, kf_idx, context)
1008        };
1009        if let Some(message) = outcome.debug_message {
1010            self.dbg(message);
1011        }
1012        if let Some((corrected_tracking_pose, corrected_velocity)) = outcome.tracking_correction {
1013            self.state.pose_world_to_cam = corrected_tracking_pose;
1014            self.state.velocity_world = corrected_velocity;
1015        }
1016        if outcome.pgo_applied {
1017            self.tracker.reset_tracks();
1018            if self.state.imu_initialized {
1019                if !self.local_mapping.submit(KeyframeJob {
1020                    imu_initialized: true,
1021                    imu_t_bc: self.rig.camera_to_body(),
1022                    gravity_world: self.inertial.gravity_world,
1023                }) {
1024                    self.dbg("[local_mapping] worker is unavailable after PGO".into());
1025                }
1026                self.apply_local_mapping_results();
1027            }
1028        }
1029        self.loop_closure_events.extend(outcome.events);
1030    }
1031}
1032
1033fn format_imu_init_gate(
1034    start_idx: usize,
1035    first_idx: Option<usize>,
1036    last_idx: Option<usize>,
1037    keyframes: usize,
1038    min_keyframes: usize,
1039    imu_time: f64,
1040    min_time_sec: f64,
1041) -> String {
1042    format!(
1043        "[imu_init_gate] start_idx={start_idx} first_idx={first_idx:?} last_idx={last_idx:?} kfs={keyframes}/{min_keyframes} imu_time={imu_time:.2}/{min_time_sec:.1}s"
1044    )
1045}
1046
1047#[cfg(test)]
1048mod tests {
1049    use super::format_imu_init_gate;
1050
1051    #[test]
1052    fn formats_compact_imu_init_gate() {
1053        assert_eq!(
1054            format_imu_init_gate(12, Some(12), Some(32), 7, 10, 1.05, 1.0),
1055            "[imu_init_gate] start_idx=12 first_idx=Some(12) last_idx=Some(32) kfs=7/10 imu_time=1.05/1.0s"
1056        );
1057    }
1058}