Skip to main content

g2o/
g2o.rs

1use std::{env, time::Instant};
2
3#[cfg(feature = "rerun")]
4use factrs::rerun::RerunObserver;
5use factrs::{
6    core::{GaussNewton, LevenMarquardt, SE2, SE3},
7    optimizers::{BaseOptParams, GncGemanMcClure, GncParams, GraduatedNonConvexity, LevenParams},
8    traits::Optimizer,
9    utils::load_g20,
10};
11#[cfg(feature = "rerun")]
12use rerun::{Arrows2D, Arrows3D, Points2D, Points3D};
13
14// Setups rerun and a callback for iteratively sending to rerun
15// Must run with --features rerun for it to work
16#[cfg(feature = "rerun")]
17fn rerun_init(opt: &mut impl Optimizer, dim: &str, obj: &str) {
18    // Setup the rerun & the callback
19    let rec = rerun::RecordingStreamBuilder::new("factrs-g2o-example")
20        .connect_grpc()
21        .unwrap();
22
23    // Log the graph
24    let (nodes, edges) = opt.graph().into();
25    rec.log_static("graph", &[&nodes as &dyn rerun::AsComponents, &edges])
26        .expect("log failed");
27
28    let topic = "base/solution";
29
30    match (dim, obj) {
31        ("se2", "points") => {
32            let callback = RerunObserver::<SE2, Points2D>::new(rec, topic);
33            opt.observers_mut().add(callback)
34        }
35        ("se2", "arrows") => {
36            let callback = RerunObserver::<SE2, Arrows2D>::new(rec, topic);
37            opt.observers_mut().add(callback)
38        }
39        ("se3", "points") => {
40            let callback = RerunObserver::<SE3, Points3D>::new(rec, topic);
41            opt.observers_mut().add(callback)
42        }
43        ("se3", "arrows") => {
44            let callback = RerunObserver::<SE3, Arrows3D>::new(rec, topic);
45            opt.observers_mut().add(callback)
46        }
47        _ => panic!("Invalid arguments"),
48    };
49}
50
51#[cfg(not(feature = "rerun"))]
52fn rerun_init(_opt: &mut impl Optimizer, _dim: &str, _obj: &str) {}
53
54fn main() -> Result<(), Box<dyn std::error::Error>> {
55    // ---------------------- Parse Arguments & Load data ---------------------- //
56    let mut args: Vec<String> = env::args().collect();
57    match args.len() {
58        2 => {
59            args.push(String::from("gauss"));
60            args.push(String::from("arrows"));
61        }
62        3 => {
63            args.push(String::from("arrows"));
64        }
65        4 => {}
66        _ => {
67            println!(
68                "Usage: {} <g2o file> <optimizer: [gauss|leven|gnc] = gauss> <vis: [points|arrows] = arrows>",
69                args[0]
70            );
71            return Ok(());
72        }
73    }
74
75    pretty_env_logger::init();
76
77    // Load the graph from the g2o file
78    let filename = &args[1];
79    let (graph, init) = load_g20(filename);
80    println!("File loaded, {} factors", graph.len());
81
82    let obj = &args[3];
83    let dim = if init.filter::<SE2>().count() != 0 {
84        "se2"
85    } else if init.filter::<SE3>().count() != 0 {
86        "se3"
87    } else {
88        panic!("Graph doesn't have SE2 or SE3 variables");
89    };
90
91    // Make optimizer
92    let mut optimizer: Box<dyn Optimizer> = match args[2].as_str() {
93        "gauss" => {
94            let params = BaseOptParams {
95                max_iterations: 200,
96                error_tol_relative: 1e-4,
97                ..Default::default()
98            };
99            let mut opt = GaussNewton::new(params, graph);
100            rerun_init(&mut opt, dim, obj);
101            Box::new(opt)
102        }
103        "leven" => {
104            let mut params = LevenParams::default();
105            params.base.max_iterations = 200;
106            params.base.error_tol_relative = 1e-4;
107            params.min_model_fidelity = -1e4;
108            let mut opt = LevenMarquardt::new(params, graph);
109            rerun_init(&mut opt, dim, obj);
110            Box::new(opt)
111        }
112        #[allow(clippy::field_reassign_with_default)]
113        "gnc" => {
114            let mut params: GncParams<LevenMarquardt> = GncParams::default();
115            params.mu_step_size = 1.4;
116            params.base.max_iterations = 200;
117            params.inner.base.max_iterations = 5;
118            params.inner.base.error_tol_absolute = 1e-4;
119            params.inner.base.error_tol_relative = 1e-4;
120            let mut opt: GraduatedNonConvexity<GncGemanMcClure, _> =
121                GraduatedNonConvexity::new(params, graph);
122            rerun_init(&mut opt, dim, obj);
123            Box::new(opt)
124        }
125        _ => {
126            println!("Optimizer not recognized");
127            return Ok(());
128        }
129    };
130
131    // ------------------------- Optimize ------------------------- //
132    let start = Instant::now();
133    let result = optimizer.optimize(init);
134    let duration = start.elapsed();
135
136    match result {
137        Ok(_) => println!("Optimization converged!"),
138        Err(_) => println!("Optimization failed!"),
139    }
140    println!("Optimization took: {duration:?}");
141    Ok(())
142}