From 4472ef26d233e2789372f1b2e9d05ca479d2ff80 Mon Sep 17 00:00:00 2001
From: Powei Lin <poweilin1994@gmail.com>
Date: Mon, 9 Dec 2024 17:20:08 -0500
Subject: [PATCH] add progress bar (#5)

---
 Cargo.toml         |  3 ++-
 src/data_loader.rs | 31 +++++++++++++++++++++----------
 2 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 6cf4ffe..db2029d 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "camera-intrinsic-calibration"
-version = "0.4.1"
+version = "0.4.2"
 edition = "2021"
 authors = ["Powei Lin <poweilin1994@gmail.com>"]
 readme = "README.md"
@@ -22,6 +22,7 @@ faer = "0.19.4"
 glam = "0.29.2"
 glob = "0.3.1"
 image = "0.25.5"
+indicatif = { version = "0.17.9", features = ["rayon"] }
 log = "0.4.22"
 nalgebra = "0.33.2"
 num-dual = "0.10.3"
diff --git a/src/data_loader.rs b/src/data_loader.rs
index 091374b..76876ac 100644
--- a/src/data_loader.rs
+++ b/src/data_loader.rs
@@ -8,6 +8,7 @@ use aprilgrid::detector::TagDetector;
 use glam::Vec2;
 use glob::glob;
 use image::{DynamicImage, ImageReader};
+use indicatif::ParallelProgressIterator;
 use rayon::prelude::*;
 
 const MIN_CORNERS: usize = 24;
@@ -74,14 +75,17 @@ pub fn load_euroc(
             let img_paths =
                 glob(format!("{}/mav0/cam{}/data/*.png", root_folder, cam_idx).as_str())
                     .expect("failed");
-            let mut time_frame: Vec<_> = img_paths
-                .skip(start_idx)
-                .step_by(step)
+            let mut sorted_path: Vec<_> = img_paths.collect();
+            sorted_path.sort_by(|a, b| a.as_ref().unwrap().cmp(b.as_ref().unwrap()));
+            let new_paths: Vec<_> = sorted_path.iter().skip(start_idx).step_by(step).collect();
+            let mut time_frame: Vec<_> = new_paths
+                .iter()
                 .par_bridge()
+                .progress_count(new_paths.len() as u64)
                 .map(|path| {
-                    let path = path.unwrap();
-                    let time_ns = path_to_timestamp(&path);
-                    let img = ImageReader::open(&path).unwrap().decode().unwrap();
+                    let path = path.as_ref().unwrap();
+                    let time_ns = path_to_timestamp(path);
+                    let img = ImageReader::open(path).unwrap().decode().unwrap();
                     if let Some(recording) = recording_option {
                         recording.set_time_nanos("stable", time_ns);
                         let topic = format!("/cam{}", cam_idx);
@@ -119,15 +123,22 @@ pub fn load_others(
             let img_paths = glob(format!("{}/**/cam{}/**/*.png", root_folder, cam_idx).as_str())
                 .expect("failed");
             log::trace!("loading cam{}", cam_idx);
-            let mut time_frame: Vec<_> = img_paths
+            let mut sorted_path: Vec<_> = img_paths.collect();
+            sorted_path.sort_by(|a, b| a.as_ref().unwrap().cmp(b.as_ref().unwrap()));
+            let new_paths: Vec<_> = sorted_path
+                .iter()
                 .skip(start_idx)
                 .step_by(step)
                 .enumerate()
+                .collect();
+            let mut time_frame: Vec<_> = new_paths
+                .iter()
                 .par_bridge()
+                .progress_count(new_paths.len() as u64)
                 .map(|(idx, path)| {
-                    let path = path.unwrap();
-                    let time_ns = idx as i64 * 100000000;
-                    let img = ImageReader::open(&path).unwrap().decode().unwrap();
+                    let path = path.as_ref().unwrap();
+                    let time_ns = *idx as i64 * 100000000;
+                    let img = ImageReader::open(path).unwrap().decode().unwrap();
                     if let Some(recording) = recording_option {
                         recording.set_time_nanos("stable", time_ns);
                         let topic = format!("/cam{}", cam_idx);