(fix): Bar init lag

(adjust): Moving resize cameras logic
AsherJingkongChen · Dec 10, 2024 · 184346b · 184346b
1 parent d486b96
commit 184346b
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 19 deletions.
diff --git a/crates/gausplat-renderer b/crates/gausplat-renderer
diff --git a/examples/gausplat-scepter/src/main.rs b/examples/gausplat-scepter/src/main.rs
@@ -8,7 +8,7 @@ pub fn main() -> Result<(), Report> {
     use ModelCommand::*;
 
     init()?;
-    log::info!(target: "gausplat::scepter::main", "init");
+    log::info!(target: "gausplat::scepter::main", "start");
 
     let args = GausplatArguments::parse()?;
     let args = match &args.model {

diff --git a/examples/gausplat-scepter/src/runner/gaussian_3d/mod.rs b/examples/gausplat-scepter/src/runner/gaussian_3d/mod.rs
@@ -25,7 +25,7 @@ use gausplat::trainer::{
     metric::{MeanStructuralSimilarity, Metric, Psnr},
     train::gaussian_3d::{Tensor, WgpuDevice},
 };
-use rayon::slice::ParallelSliceMut;
+use rayon::{iter::ParallelIterator, slice::ParallelSliceMut};
 
 /// ## Returns
 ///
@@ -160,3 +160,18 @@ pub fn get_mssim_and_psnr(
 
     Ok((mssim_mean, psnr_mean))
 }
+
+/// Resize `cameras` to proper sizes.
+pub fn resize_cameras(cameras: &mut Cameras) -> Result<(), Report> {
+    const IMAGE_SIZE_MIN: u32 = 320;
+    const IMAGE_SIZE_MAX: u32 = IMAGE_SIZE_MIN * 5;
+
+    cameras.par_values_mut().try_for_each(|camera| {
+        let size_source = camera.size_max();
+        let size_target = size_source.clamp(IMAGE_SIZE_MIN, IMAGE_SIZE_MAX);
+        if size_source != size_target {
+            camera.resize_max(size_target)?;
+        }
+        Ok(())
+    })
+}
diff --git a/examples/gausplat-scepter/src/runner/gaussian_3d/render.rs b/examples/gausplat-scepter/src/runner/gaussian_3d/render.rs
@@ -95,6 +95,9 @@ impl RenderRunner {
         let should_show_progress = !bar.disable && size != 0;
 
         bar.reset(Some(size));
+        if should_show_progress {
+            bar.refresh()?;
+        }
 
         let result = cameras
             .into_values()

diff --git a/examples/gausplat-scepter/src/runner/gaussian_3d/train.rs b/examples/gausplat-scepter/src/runner/gaussian_3d/train.rs
@@ -13,7 +13,6 @@ use gausplat::trainer::{
         Autodiff, AutodiffModule, Gaussian3dTrainer, RefinerConfig, SEED,
     },
 };
-use rayon::iter::ParallelIterator;
 use std::{
     fmt, fs,
     path::{Path, PathBuf},
@@ -86,9 +85,6 @@ impl TrainRunner {
 
 impl Runner for TrainRunner {
     fn run(mut self) -> Result<(), Report> {
-        const IMAGE_SIZE_MIN: u32 = 320;
-        const IMAGE_SIZE_MAX: u32 = IMAGE_SIZE_MIN * 5;
-
         // Specifying the parameters
 
         let device = self.scene.device();
@@ -123,22 +119,16 @@ impl Runner for TrainRunner {
         // Rescaling down the images at initialization
 
         let time = Instant::now();
-        self.cameras_train.par_values_mut().try_for_each(|camera| {
-            let size_source = camera.size_max();
-            let size_target = size_source.clamp(IMAGE_SIZE_MIN, IMAGE_SIZE_MAX);
-            if size_source != size_target {
-                camera.resize_max(size_target)?;
-            }
-            Ok::<_, Report>(())
-        })?;
+        resize_cameras(&mut self.cameras_train)?;
         log::info!(
             target: "gausplat::scepter::gaussian_3d::train",
             "may rescale in {:.03?}", time.elapsed(),
         );
 
         // Optimizing the scene iteratively
 
-        self.cameras_train
+        let result = self
+            .cameras_train
             .seed(SEED)
             .random_values()
             .take(iterations)
@@ -204,14 +194,14 @@ impl Runner for TrainRunner {
                     }
                 }
 
-                Ok::<_, Report>(())
-            })?;
+                Ok(())
+            });
 
         if !bar.disable {
             eprintln!();
         }
 
-        Ok(())
+        result
     }
 }