Clean up (#24)

* update aprilgrid

* clean up

* add tests

* clean up

* docstring

Author: powei-lin

.github/workflows/rust.yml

@@ -20,5 +20,7 @@ jobs:
       run: cargo fmt --check --verbose
     - name: Linting
       run: cargo clippy
+    - name: Test
+      run: cargo test --verbose
     - name: Build
       run: cargo build --verbose

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "camera-intrinsic-calibration"
-version = "0.9.0"
+version = "0.10.0"
 edition = "2024"
 authors = ["Powei Lin <poweilin1994@gmail.com>"]
 readme = "README.md"
@@ -21,7 +21,7 @@ exclude = [
 ]
 
 [dependencies]
-aprilgrid = "0.6.1"
+aprilgrid = "0.7.0"
 camera-intrinsic-model = "0.6.0"
 clap = { version = "4.5", features = ["derive"] }
 colorous = "1.0.16"

src/bin/camera_calibration.rs

@@ -72,17 +72,56 @@ fn main() {
 
     let cli = CCRSCli::parse();
     let detector = TagDetector::new(&cli.tag_family, None);
-    let board = if let Some(board_config_path) = cli.board_config {
-        Board::from_config(&object_from_json(&board_config_path))
+    let board = setup_board(&cli);
+    let output_folder = setup_output_folder(&cli);
+
+    let recording = rerun::RecordingStreamBuilder::new("calibration")
+        .save(format!("{}/logging.rrd", output_folder))
+        .unwrap();
+    recording
+        .log_static("/", &rerun::ViewCoordinates::RDF())
+        .unwrap();
+
+    let cams_detected_feature_frames = load_feature_data(&cli, &detector, &board, &recording);
+
+    let (calibrated_intrinsics, cam_rtvecs) =
+        calibrate_all_cameras(&cli, &cams_detected_feature_frames, &recording);
+
+    let t_cam_i_0_init = init_camera_extrinsic(&cam_rtvecs);
+
+    save_and_validate_results(
+        &cli,
+        &output_folder,
+        &cams_detected_feature_frames,
+        &calibrated_intrinsics,
+        &cam_rtvecs,
+        &t_cam_i_0_init,
+        &recording,
+    );
+}
+
+/// Loads the board configuration specified in the CLI arguments or creates a default one.
+///
+/// If a config file path is provided via `--board-config`, it loads from that file.
+/// Otherwise, it creates a default 6x6 AprilGrid configuration and saves it to `default_board_config.json`.
+fn setup_board(cli: &CCRSCli) -> Board {
+    if let Some(board_config_path) = &cli.board_config {
+        Board::from_config(&object_from_json(board_config_path))
     } else {
         let config = BoardConfig::default();
         object_to_json("default_board_config.json", &config);
         Board::from_config(&config)
-    };
-    let dataset_root = &cli.path;
-    let now = Instant::now();
-    let output_folder = if let Some(output_folder) = cli.output_folder {
-        output_folder
+    }
+}
+
+/// Sets up the output directory for calibration results.
+///
+/// If `--output-folder` is specified, uses that path.
+/// Otherwise, creates a directory named with the current timestamp under `results/`.
+/// Ensures the directory exists.
+fn setup_output_folder(cli: &CCRSCli) -> String {
+    let output_folder = if let Some(output_folder) = &cli.output_folder {
+        output_folder.clone()
     } else {
         let now = OffsetDateTime::now_local().unwrap();
         format!(
@@ -96,52 +135,84 @@ fn main() {
         )
     };
     std::fs::create_dir_all(&output_folder).expect("Valid path");
+    output_folder
+}
 
-    let recording = rerun::RecordingStreamBuilder::new("calibration")
-        .save(format!("{}/logging.rrd", output_folder))
-        .unwrap();
-    recording
-        .log_static("/", &rerun::ViewCoordinates::RDF())
-        .unwrap();
+/// Loads feature data from the dataset.
+///
+/// Supports Euroc and General dataset formats.
+/// Uses the provided tag detector and board configuration to extract features.
+/// Logs images to Rerun if enabled.
+///
+/// # Returns
+/// A vector of vectors, where each inner vector contains `Option<FrameFeature>` for a camera.
+fn load_feature_data(
+    cli: &CCRSCli,
+    detector: &TagDetector,
+    board: &Board,
+    recording: &rerun::RecordingStream,
+) -> Vec<Vec<Option<FrameFeature>>> {
     trace!("Start loading data");
     println!("Start loading images and detecting charts.");
+    let now = Instant::now();
     let mut cams_detected_feature_frames: Vec<Vec<Option<FrameFeature>>> = match cli.dataset_format
     {
         DatasetFormat::Euroc => load_euroc(
-            dataset_root,
-            &detector,
-            &board,
+            &cli.path,
+            detector,
+            board,
             cli.start_idx,
             cli.step,
             cli.cam_num,
-            Some(&recording),
+            Some(recording),
         ),
         DatasetFormat::General => load_others(
-            dataset_root,
-            &detector,
-            &board,
+            &cli.path,
+            detector,
+            board,
             cli.start_idx,
             cli.step,
             cli.cam_num,
-            Some(&recording),
+            Some(recording),
         ),
     };
     let duration_sec = now.elapsed().as_secs_f64();
     println!("detecting feature took {:.6} sec", duration_sec);
-    println!("total: {} images", cams_detected_feature_frames[0].len());
+    if !cams_detected_feature_frames.is_empty() {
+        println!("total: {} images", cams_detected_feature_frames[0].len());
+        println!(
+            "avg: {} sec",
+            duration_sec / cams_detected_feature_frames[0].len() as f64
+        );
+    }
+
     cams_detected_feature_frames
         .iter_mut()
         .for_each(|f| f.truncate(cli.max_images));
-    println!(
-        "avg: {} sec",
-        duration_sec / cams_detected_feature_frames[0].len() as f64
-    );
-    let (calibrated_intrinsics, cam_rtvecs): (Vec<_>, Vec<_>) = cams_detected_feature_frames
+
+    cams_detected_feature_frames
+}
+
+/// Calibrates all cameras individually.
+///
+/// Iterates through each camera, detecting features and running the optimization.
+/// Retries calibration up to 3 times if it fails.
+///
+/// # Returns
+/// A tuple containing:
+/// - `Vec<GenericModel<f64>>`: The calibrated intrinsic models for each camera.
+/// - `Vec<HashMap<usize, RvecTvec>>`: estimated camera poses for each frame.
+fn calibrate_all_cameras(
+    cli: &CCRSCli,
+    cams_detected_feature_frames: &[Vec<Option<FrameFeature>>],
+    recording: &rerun::RecordingStream,
+) -> (Vec<GenericModel<f64>>, Vec<HashMap<usize, RvecTvec>>) {
+    cams_detected_feature_frames
         .iter()
         .enumerate()
         .map(|(cam_idx, feature_frames)| {
             let topic = format!("/cam{}", cam_idx);
-            log_feature_frames(&recording, &topic, feature_frames);
+            log_feature_frames(recording, &topic, feature_frames);
             let mut calibrated_result: Option<(GenericModel<f64>, HashMap<usize, RvecTvec>)> = None;
             let max_trials = 3;
             let cam0_fixed_focal = if cam_idx == 0 { cli.fixed_focal } else { None };
@@ -153,9 +224,9 @@ fn main() {
             for trial in 0..max_trials {
                 calibrated_result = init_and_calibrate_one_camera(
                     cam_idx,
-                    &cams_detected_feature_frames,
+                    cams_detected_feature_frames,
                     &cli.model,
-                    &recording,
+                    recording,
                     &calib_params,
                     trial > 0,
                 );
@@ -169,19 +240,35 @@ fn main() {
                     cam_idx, max_trials
                 );
             }
-            let (final_result, rtvec_map) = calibrated_result.unwrap();
-            (final_result, rtvec_map)
+            calibrated_result.unwrap()
         })
-        .unzip();
-    let t_cam_i_0_init = init_camera_extrinsic(&cam_rtvecs);
-    for t in &t_cam_i_0_init {
+        .unzip()
+}
+
+/// Saves calibration results and performs validation.
+///
+/// saves intrinsics, extrinsics, and pose data to JSON files.
+/// Generates a validation report and logs visualization data to Rerun.
+/// If multiple cameras are present, it also attempts to calibrate extrinsics between cameras.
+#[allow(clippy::too_many_arguments)]
+fn save_and_validate_results(
+    cli: &CCRSCli,
+    output_folder: &str,
+    cams_detected_feature_frames: &[Vec<Option<FrameFeature>>],
+    intrinsics: &[GenericModel<f64>],
+    cam_rtvecs: &[HashMap<usize, RvecTvec>],
+    t_cam_i_0_init: &[RvecTvec],
+    recording: &rerun::RecordingStream,
+) {
+    for t in t_cam_i_0_init {
         println!("r {} t {}", t.na_rvec(), t.na_tvec());
     }
+
     if let Some((camera_intrinsics, t_i_0, board_rtvecs)) = calib_all_camera_with_extrinsics(
-        &calibrated_intrinsics,
-        &t_cam_i_0_init,
-        &cam_rtvecs,
-        &cams_detected_feature_frames,
+        intrinsics,
+        t_cam_i_0_init,
+        cam_rtvecs,
+        cams_detected_feature_frames,
         cli.one_focal || cli.fixed_focal.is_some(),
         cli.disabled_distortion_num,
         cli.fixed_focal.is_some(),
@@ -215,7 +302,7 @@ fn main() {
                 intrinsic,
                 &new_rtvec_map,
                 &cams_detected_feature_frames[cam_idx],
-                Some(&recording),
+                Some(recording),
             );
             rep_rms.push(rep);
             println!(
@@ -232,15 +319,13 @@ fn main() {
         );
     } else {
         let mut rep_rms = Vec::new();
-        for (cam_idx, (intrinsic, rtvec_map)) in
-            calibrated_intrinsics.iter().zip(cam_rtvecs).enumerate()
-        {
+        for (cam_idx, (intrinsic, rtvec_map)) in intrinsics.iter().zip(cam_rtvecs).enumerate() {
             let rep = validation(
                 cam_idx,
                 intrinsic,
-                &rtvec_map,
+                rtvec_map,
                 &cams_detected_feature_frames[cam_idx],
-                Some(&recording),
+                Some(recording),
             );
             rep_rms.push(rep);
             println!(

src/board.rs

@@ -2,6 +2,7 @@ use glam;
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 
+/// Configuration for an AprilGrid board.
 #[derive(Debug, Serialize, Deserialize)]
 pub struct BoardConfig {
     tag_size_meter: f32,
@@ -23,7 +24,9 @@ impl Default for BoardConfig {
     }
 }
 
+/// Represents a calibration board with known 3D points.
 pub struct Board {
+    /// Maps tag ID to its 3D position (usually top-left or other corner definition).
     pub id_to_3d: HashMap<u32, glam::Vec3>,
 }
 
@@ -37,6 +40,9 @@ impl Board {
             board_config.first_id,
         )
     }
+    /// Initializes an AprilGrid board with standard layout.
+    ///
+    /// The grid is generated starting from (0,0,0) and expanding in positive X (cols) and negative Y (rows).
     pub fn init_aprilgrid(
         tag_size_meter: f32,
         tag_spacing: f32,
@@ -89,6 +95,7 @@ impl Board {
     }
 }
 
+/// Creates a default 6x6 AprilGrid board.
 pub fn create_default_6x6_board() -> Board {
     Board::init_aprilgrid(0.088, 0.3, 6, 6, 0)
 }

src/data_loader.rs

@@ -14,6 +14,9 @@ use rerun::TimeCell;
 
 const MIN_CORNERS: usize = 24;
 
+/// Parses the timestamp from a file path.
+///
+/// Assumes the filename (without extension) is a timestamp in nanoseconds.
 fn path_to_timestamp(path: &Path) -> i64 {
     let time_ns: i64 = path
         .file_stem()
@@ -25,6 +28,11 @@ fn path_to_timestamp(path: &Path) -> i64 {
     time_ns
 }
 
+/// Detects features in an image and converts it to a `FrameFeature`.
+///
+/// Uses `aprilgrid` detector to find tags.
+/// matches detected point IDs to 3D board coordinates.
+/// Returns `None` if the number of detected corners is less than `min_corners`.
 fn image_to_option_feature_frame(
     tag_detector: &TagDetector,
     img: &DynamicImage,
@@ -72,6 +80,18 @@ fn img_filter(rp: glob::GlobResult) -> Option<std::path::PathBuf> {
     None
 }
 
+/// Loads data from a Euroc-style dataset.
+///
+/// Iterates through camera folders, loads images, detects features in parallel.
+///
+/// # Arguments
+/// * `root_folder` - Path to the dataset root.
+/// * `tag_detector` - The tag detector instance.
+/// * `board` - The calibration board configuration.
+/// * `start_idx` - Starting image index.
+/// * `step` - Step size for sampling images.
+/// * `cam_num` - Number of cameras.
+/// * `recording_option` - Optional Rerun recording stream for visualization.
 pub fn load_euroc(
     root_folder: &str,
     tag_detector: &TagDetector,
@@ -92,8 +112,7 @@ pub fn load_euroc(
             sorted_path.sort();
             let new_paths: Vec<_> = sorted_path.iter().skip(start_idx).step_by(step).collect();
             let mut time_frame: Vec<_> = new_paths
-                .iter()
-                .par_bridge()
+                .par_iter()
                 .progress_count(new_paths.len() as u64)
                 .map(|path| {
                     let time_ns = path_to_timestamp(path);
@@ -124,6 +143,20 @@ pub fn load_euroc(
         .collect()
 }
 
+/// Loads data from a general dataset structure.
+///
+/// Iterates through camera folders matching `**/cam{}/**/*`.
+/// Loads images and detects features in parallel.
+/// Timestamps are generated artificially based on index if not present in filename (though this function ignores filename timestamp logic in favor of index-based).
+///
+/// # Arguments
+/// * `root_folder` - Path to the dataset root.
+/// * `tag_detector` - The tag detector instance.
+/// * `board` - The calibration board configuration.
+/// * `start_idx` - Starting image index.
+/// * `step` - Step size for sampling images.
+/// * `cam_num` - Number of cameras.
+/// * `recording_option` - Optional Rerun recording stream for visualization.
 pub fn load_others(
     root_folder: &str,
     tag_detector: &TagDetector,
@@ -149,8 +182,7 @@ pub fn load_others(
                 .enumerate()
                 .collect();
             let mut time_frame: Vec<_> = new_paths
-                .iter()
-                .par_bridge()
+                .par_iter()
                 .progress_count(new_paths.len() as u64)
                 .map(|(idx, path)| {
                     let time_ns = *idx as i64 * 100000000;