From 8b728d4ec8becf5b35dfe724c6d2bae774126e87 Mon Sep 17 00:00:00 2001 From: Brent Yi Date: Wed, 28 Feb 2024 11:41:03 -0800 Subject: [PATCH] Fix progress bar for full images datamanager (#2969) --- .../datamanagers/full_images_datamanager.py | 26 +++++++++++-------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/nerfstudio/data/datamanagers/full_images_datamanager.py b/nerfstudio/data/datamanagers/full_images_datamanager.py index 9b2057357..6e041b33d 100644 --- a/nerfstudio/data/datamanagers/full_images_datamanager.py +++ b/nerfstudio/data/datamanagers/full_images_datamanager.py @@ -180,26 +180,30 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): return data CONSOLE.log("Caching / undistorting train images") - with ThreadPoolExecutor() as executor: + with ThreadPoolExecutor(max_workers=2) as executor: cached_train = list( - executor.map( - process_train_data, - track( + track( + executor.map( + process_train_data, range(len(self.train_dataset)), - description="Caching / undistorting train images", - transient=True, ), + description="Caching / undistorting train images", + transient=True, + total=len(self.train_dataset), ) ) CONSOLE.log("Caching / undistorting eval images") - with ThreadPoolExecutor() as executor: + with ThreadPoolExecutor(max_workers=2) as executor: cached_eval = list( - executor.map( - process_eval_data, - track( - range(len(self.eval_dataset)), description="Caching / undistorting eval images", transient=True + track( + executor.map( + process_eval_data, + range(len(self.eval_dataset)), ), + description="Caching / undistorting eval images", + transient=True, + total=len(self.eval_dataset), ) )