visual_odometry_rs/core/track/
inverse_compositional.rs

1// This Source Code Form is subject to the terms of the Mozilla Public
2// License, v. 2.0. If a copy of the MPL was not distributed with this
3// file, You can obtain one at http://mozilla.org/MPL/2.0/.
4
5//! Types and functions to implement an inverse compositional tracking algorithm.
6//!
7//! Implementation of "Lucas-kanade 20 years on: A unifying framework"
8//! in the inverse compositional case.
9//! The warping function is parameterized by the Lie Algebra of twists se(3).
10
11use itertools::izip;
12use nalgebra::DMatrix;
13
14use crate::core::{
15    camera::Intrinsics,
16    candidates::coarse_to_fine as candidates,
17    gradient,
18    inverse_depth::{self, InverseDepth},
19    multires,
20    track::lm_optimizer::{self, LMOptimizerState},
21};
22use crate::math::optimizer::State as _;
23use crate::misc::helper;
24use crate::misc::type_aliases::{Float, Iso3, Mat6, Point2, Vec6};
25
26/// Type alias to easily spot vectors that are indexed over multi-resolution levels.
27pub type Levels<T> = Vec<T>;
28
29/// Struct used for tracking the camera at each frame.
30/// Can only be constructed by initialization from a `Config`.
31pub struct Tracker {
32    config: Config,
33    state: State,
34}
35
36/// Configuration of the Tracker.
37pub struct Config {
38    /// Number of levels in the multi-resolution pyramids of images.
39    pub nb_levels: usize,
40    /// Threshold for the candidates selection algorithm.
41    pub candidates_diff_threshold: u16,
42    /// Scale of the depth 16 bit images.
43    /// This is 5000.0 for the TUM RGB-D dataset.
44    pub depth_scale: Float,
45    /// Camera intrinsic parameters.
46    pub intrinsics: Intrinsics,
47    /// Default variance of the inverse depth values coming from the depth map.
48    pub idepth_variance: Float,
49}
50
51/// Internal state of the tracker.
52struct State {
53    keyframe_multires_data: MultiresData,
54    keyframe_depth_timestamp: f64,
55    keyframe_img_timestamp: f64,
56    keyframe_pose: Iso3,
57    current_frame_depth_timestamp: f64,
58    current_frame_img_timestamp: f64,
59    current_frame_pose: Iso3,
60}
61
62/// Mostly multi-resolution data related to the frame.
63#[allow(clippy::type_complexity)]
64struct MultiresData {
65    intrinsics_multires: Levels<Intrinsics>,
66    img_multires: Levels<DMatrix<u8>>,
67    usable_candidates_multires: Levels<(Vec<(usize, usize)>, Vec<Float>)>,
68    jacobians_multires: Levels<Vec<Vec6>>,
69    hessians_multires: Levels<Vec<Mat6>>,
70}
71
72impl Config {
73    /// Initialize a tracker with the first RGB-D frame.
74    pub fn init(
75        self,
76        keyframe_depth_timestamp: f64,
77        depth_map: &DMatrix<u16>,
78        keyframe_img_timestamp: f64,
79        img: DMatrix<u8>,
80    ) -> Tracker {
81        // Precompute multi-resolution first frame data.
82        let intrinsics_multires = self.intrinsics.clone().multi_res(self.nb_levels);
83        let img_multires = multires::mean_pyramid(self.nb_levels, img);
84        let keyframe_multires_data =
85            precompute_multires_data(&self, depth_map, intrinsics_multires, img_multires);
86
87        // Regroup everything under the returned Tracker.
88        Tracker {
89            state: State {
90                keyframe_multires_data,
91                keyframe_depth_timestamp,
92                keyframe_img_timestamp,
93                keyframe_pose: Iso3::identity(),
94                current_frame_depth_timestamp: keyframe_depth_timestamp,
95                current_frame_img_timestamp: keyframe_img_timestamp,
96                current_frame_pose: Iso3::identity(),
97            },
98            config: self,
99        }
100    }
101} // impl Config
102
103/// Precompute the multi-resolution data of a frame.
104#[allow(clippy::used_underscore_binding)]
105fn precompute_multires_data(
106    config: &Config,
107    depth_map: &DMatrix<u16>,
108    intrinsics_multires: Levels<Intrinsics>,
109    img_multires: Levels<DMatrix<u8>>,
110) -> MultiresData {
111    // Precompute multi-resolution of keyframe gradients.
112    let mut gradients_multires = multires::gradients_xy(&img_multires);
113    gradients_multires.insert(0, gradient::centered(&img_multires[0]));
114    let gradients_squared_norm_multires: Vec<_> = gradients_multires
115        .iter()
116        .map(|(gx, gy)| gradient::squared_norm(gx, gy))
117        .collect();
118
119    // Precompute mask of candidate points for tracking.
120    let candidates_points = candidates::select(
121        config.candidates_diff_threshold,
122        &gradients_squared_norm_multires,
123    )
124    .pop()
125    .unwrap();
126
127    // Only keep the "usable" points, i.e. those with a known depth information.
128    let from_depth = |z| inverse_depth::from_depth(config.depth_scale, z, config.idepth_variance);
129    let idepth_candidates = helper::zip_mask_map(
130        depth_map,
131        &candidates_points,
132        InverseDepth::Unknown,
133        from_depth,
134    );
135    let fuse = |a, b, c, d| inverse_depth::fuse(a, b, c, d, inverse_depth::strategy_dso_mean);
136    let idepth_multires = multires::limited_sequence(config.nb_levels, idepth_candidates, |m| {
137        multires::halve(m, fuse)
138    });
139    let usable_candidates_multires: Levels<_> = idepth_multires.iter().map(extract_z).collect();
140
141    // Precompute the Jacobians.
142    let jacobians_multires: Levels<Vec<Vec6>> = izip!(
143        &intrinsics_multires,
144        &usable_candidates_multires,
145        &gradients_multires,
146    )
147    .map(|(intrinsics, (coord, _z), (gx, gy))| warp_jacobians(intrinsics, coord, _z, gx, gy))
148    .collect();
149
150    // Precompute the Hessians.
151    let hessians_multires: Levels<_> = jacobians_multires.iter().map(hessians_vec).collect();
152
153    // Regroup everything under a MultiresData.
154    MultiresData {
155        intrinsics_multires,
156        img_multires,
157        usable_candidates_multires,
158        jacobians_multires,
159        hessians_multires,
160    }
161}
162
163impl Tracker {
164    /// Track a new frame.
165    /// Internally mutates the tracker state.
166    ///
167    /// You can use `tracker.current_frame()` after tracking to retrieve the new frame pose.
168    #[allow(clippy::used_underscore_binding)]
169    #[allow(clippy::cast_precision_loss)]
170    pub fn track(
171        &mut self,
172        depth_time: f64,
173        depth_map: &DMatrix<u16>,
174        img_time: f64,
175        img: DMatrix<u8>,
176    ) {
177        let mut lm_model = self.state.current_frame_pose.inverse() * self.state.keyframe_pose;
178        let img_multires = multires::mean_pyramid(self.config.nb_levels, img);
179        let keyframe_data = &self.state.keyframe_multires_data;
180        let mut optimization_went_well = true;
181        for lvl in (0..self.config.nb_levels).rev() {
182            let obs = lm_optimizer::Obs {
183                intrinsics: &keyframe_data.intrinsics_multires[lvl],
184                template: &keyframe_data.img_multires[lvl],
185                image: &img_multires[lvl],
186                coordinates: &keyframe_data.usable_candidates_multires[lvl].0,
187                _z_candidates: &keyframe_data.usable_candidates_multires[lvl].1,
188                jacobians: &keyframe_data.jacobians_multires[lvl],
189                hessians: &keyframe_data.hessians_multires[lvl],
190            };
191            match LMOptimizerState::iterative_solve(&obs, lm_model) {
192                Ok((lm_state, _)) => {
193                    lm_model = lm_state.eval_data.model;
194                }
195                Err(err) => {
196                    eprintln!("{}", err);
197                    optimization_went_well = false;
198                    break;
199                }
200            }
201        }
202
203        // Update current frame info in tracker.
204        self.state.current_frame_depth_timestamp = depth_time;
205        self.state.current_frame_img_timestamp = img_time;
206        if optimization_went_well {
207            self.state.current_frame_pose = self.state.keyframe_pose * lm_model.inverse();
208        }
209
210        // Check if we need to change the keyframe.
211        let (coordinates, _z_candidates) = keyframe_data.usable_candidates_multires.last().unwrap();
212        let intrinsics = keyframe_data.intrinsics_multires.last().unwrap();
213        let optical_flow_sum: Float = _z_candidates
214            .iter()
215            .zip(coordinates.iter())
216            .map(|(&_z, &(x, y))| {
217                let (u, v) = warp(&lm_model, x as Float, y as Float, _z, intrinsics);
218                (x as Float - u).abs() + (y as Float - v).abs()
219            })
220            .sum();
221        let optical_flow = optical_flow_sum / _z_candidates.len() as Float;
222        eprintln!("Optical_flow: {}", optical_flow);
223
224        let change_keyframe = optical_flow >= 1.0;
225
226        // In case of keyframe change, update all keyframe info with current frame.
227        if change_keyframe {
228            let delta_time = depth_time - self.state.keyframe_depth_timestamp;
229            eprintln!("Changing keyframe after: {} seconds", delta_time);
230            self.state.keyframe_multires_data = precompute_multires_data(
231                &self.config,
232                depth_map,
233                keyframe_data.intrinsics_multires.clone(),
234                img_multires,
235            );
236            self.state.keyframe_depth_timestamp = depth_time;
237            self.state.keyframe_img_timestamp = img_time;
238            self.state.keyframe_pose = self.state.current_frame_pose;
239        }
240    } // track
241
242    /// Retrieve the current frame timestamp (of depth image) and pose.
243    pub fn current_frame(&self) -> (f64, Iso3) {
244        (
245            self.state.current_frame_depth_timestamp,
246            self.state.current_frame_pose,
247        )
248    }
249} // impl Tracker
250
251// Helper ######################################################################
252
253// fn angle(uq: UnitQuaternion<Float>) -> Float {
254//     let w = uq.into_inner().scalar();
255//     2.0 * uq.into_inner().vector().norm().atan2(w)
256// }
257
258/// Extract known inverse depth values (and coordinates) into vectorized data.
259#[allow(clippy::used_underscore_binding)]
260fn extract_z(idepth_mat: &DMatrix<InverseDepth>) -> (Vec<(usize, usize)>, Vec<Float>) {
261    let mut u = 0;
262    let mut v = 0;
263    // TODO: can allocating with a known max size improve performances?
264    let mut coordinates = Vec::new();
265    let mut _z_vec = Vec::new();
266    let (nb_rows, _) = idepth_mat.shape();
267    for idepth in idepth_mat.iter() {
268        if let InverseDepth::WithVariance(_z, _) = *idepth {
269            coordinates.push((u, v));
270            _z_vec.push(_z);
271        }
272        v += 1;
273        if v >= nb_rows {
274            u += 1;
275            v = 0;
276        }
277    }
278    (coordinates, _z_vec)
279}
280
281/// Precompute jacobians for each candidate.
282#[allow(clippy::used_underscore_binding)]
283#[allow(clippy::cast_precision_loss)]
284fn warp_jacobians(
285    intrinsics: &Intrinsics,
286    coordinates: &[(usize, usize)],
287    _z_candidates: &[Float],
288    grad_x: &DMatrix<i16>,
289    grad_y: &DMatrix<i16>,
290) -> Vec<Vec6> {
291    // Bind intrinsics to shorter names
292    let (cu, cv) = intrinsics.principal_point;
293    let (fu, fv) = intrinsics.focal;
294    let s = intrinsics.skew;
295
296    // Iterate on inverse depth candidates
297    coordinates
298        .iter()
299        .zip(_z_candidates.iter())
300        .map(|(&(u, v), &_z)| {
301            let gu = Float::from(grad_x[(v, u)]);
302            let gv = Float::from(grad_y[(v, u)]);
303            warp_jacobian_at(gu, gv, u as Float, v as Float, _z, cu, cv, fu, fv, s)
304        })
305        .collect()
306}
307
308/// Jacobian of the warping function for the inverse compositional algorithm.
309#[allow(clippy::too_many_arguments)]
310#[allow(clippy::many_single_char_names)]
311#[allow(clippy::used_underscore_binding)]
312#[allow(clippy::similar_names)]
313fn warp_jacobian_at(
314    gu: Float,
315    gv: Float,
316    u: Float,
317    v: Float,
318    _z: Float,
319    cu: Float,
320    cv: Float,
321    fu: Float,
322    fv: Float,
323    s: Float,
324) -> Vec6 {
325    // Intermediate computations
326    let a = u - cu;
327    let b = v - cv;
328    let c = a * fv - s * b;
329    let _fv = 1.0 / fv;
330    let _fuv = 1.0 / (fu * fv);
331
332    // Jacobian of the warp
333    Vec6::new(
334        gu * _z * fu,                                       //
335        _z * (gu * s + gv * fv),                            //  linear velocity terms
336        -_z * (gu * a + gv * b),                            //  ___
337        gu * (-a * b * _fv - s) + gv * (-b * b * _fv - fv), //
338        gu * (a * c * _fuv + fu) + gv * (b * c * _fuv),     //  angular velocity terms
339        gu * (-fu * fu * b + s * c) * _fuv + gv * (c / fu), //
340    )
341}
342
343/// Compute hessians components for each candidate point.
344#[allow(clippy::ptr_arg)] // TODO: Applying clippy lint here results in compilation error.
345fn hessians_vec(jacobians: &Vec<Vec6>) -> Vec<Mat6> {
346    // TODO: might be better to inline this within the function computing the jacobians.
347    jacobians.iter().map(|j| j * j.transpose()).collect()
348}
349
350/// Warp a point from an image to another by a given rigid body motion.
351#[allow(clippy::used_underscore_binding)]
352fn warp(model: &Iso3, x: Float, y: Float, _z: Float, intrinsics: &Intrinsics) -> (Float, Float) {
353    // TODO: maybe move into the camera module?
354    let x1 = intrinsics.back_project(Point2::new(x, y), 1.0 / _z);
355    let x2 = model * x1;
356    let uvz2 = intrinsics.project(x2);
357    (uvz2.x / uvz2.z, uvz2.y / uvz2.z)
358}
visual_odometry_rs/core/track/inverse_compositional.rs

visual_odometry_rs/core/track/
inverse_compositional.rs
