Dataloading Revamp (#3216)

* initial debugging and testing works * pwais changes with RayBatchStream to alleviate training * few bugs to iron out with multiprocessing, specifically pickled collate_fn * working version of RayBatchStream * additional docstrings * cleanup * much more documentation * successfully trained AEA-script2_seq2 closed_loop without OOM * porting over aria dataset-size feature * added logic to handle eviction of a worker's cached_collated_batch * antonio's implementation of stream batches * training on a dataset with 4000 images works! * some configuration speedups, loops aren't actually needed! * quick fix adjustment to aria * removed unnecessary looping * much faster training when adding i variable to collate every 5 ray bundles * cleanup unnecssary variables in Dataloader * further cleanup * adding caching of compressed images to RAM to reduce disk bottleneck * added caching to RAM for masks * found fast way to collate - many tricks applied * quick update to aria to test on different datasets * cleaned up the accelerated pil_to_numpy function * cleaning up PR * this commit was used to generate the time metrics and profiling metrics * REAL commit used to run tests * testing with nerfacto-big * generated RayBundle collate and converting images from uint8s to float32 on GPU tests * updating nerfacto to support uint8 easily, will need to figure out a way to contain this within the datamanager API * datamanager updates, both splat and nerf * must use writeable arrays because torch requires them * cleaned up base_dataset, added pickle to utils, more code in full_image, and cleaner desc for base_datamanager * lots of process on a parallel FullImageDatamanger * can train big splats with pre-assertion hack or ROI hack and 0 workers * fixed all undistortion issues with ParallelImageDatamanager * adding some downsampling and parallel tests with splatfacto! * deleted commented code in dataloaders.py and added bugfix to shuffling * testing splatfacto-big * cleaned up base_pipeline.py * cleaned up base_pipeline.py ACTUALLY THIS TIME, forgot to save last time * cleaned up a lot of code * process_project_aria back to main branch and some cleanup in full_image_datamanager * clarifying docstrings * further PR cleanup * updating models * further cleanup * removed caching of images into bytestrings * adding caching of compressed images to RAM, forgot that hardware matters * removing oom methods, adding the ability to add a flag to dataloading * removed CacheDataloader, moved RayBatchStream to dataloaders.py, new vanilla_datamanager rewritten * fixing base_piplines, deleting a weird datamanager_configs file that was accidently created * cleaning up next_train * replaced parallel datamanager with new datamanager * reverted the original base_datamanager.py, new datamanager replaced parallel_datamanager.py * modified VanillaConfig, but VanillaDataManager is the same as before * cleaning up, 2 datamanagers now - original and new parallel one * able to train with new nerfstudio dataloader now * side by side datamanagers, moved tons of logic into dataloaders.py and created new files for our parallel datamangers * added custom ray processing API to support implementations like LERF, cleaned up FullImageDatamanager to original because of new ParallelImageDatamanger * adding functionality for ns-eval by adding FixedIndicesEvalDataloader to the setup_eval * adding both ray API and image-view API to datamanagers for custom parallelization * updating splatfacto config for 4k tests * updating docstrings to be more descriptive * new datamanager API breaks when setup_eval() has multiple workers, not sure why but single worker will have to do * adding custom_view_processor to ImageBatchStream * reverting full_images_datamanager to main branch * removing nn.Module inheritance from Datamanager class * don't need to move datamanger to device anymore since Datamanager is not a subclass of nn.Module * finished integration test with nerfacto * simplified config variables, integrated the parallelism/disk-data-loading all into one datamanager * updated the splatfacto config to be simpler with the dataloading and now uses FullImageDatamanager (which has been changed) * style checks and some cleanup * new splatfacto test, cleaning up nerfacto integration test * removing redundant parallel_full_images_datamaanger, as the OG full_image_datamanager now has full parallelized support * ruff linting and pyright fixing * further pyright fixing * another pyright fixing * fixing pyright error, camera optimization no longer part of datamanager * fixing one pyright * fixing dataloading error when camera is not undistorted with dataloader * fixing comments and updating style * undoing a style change i made * undoing another style change i made by accident * fixing slow runtime * fixing a more general camera undistortion bug * move images to device properly * minor improvements * add print statement about >500 images, cleanup method configs * make method configs consistent across nerfacto models * adding description comments * updating description * resolving some pyright issues with export.py, explained in PR desc * fixing pyright issues in base_pipeline.py * ran pyright on exporter and base_pipeline.py without issues * adding a git ignore to a clearly checked pyright issue * typo * fixing most ns-dev-test cases * cleanup, passing final ns-dev-test * oops, accidentally pushed the deletion of a docstring, undoing that * another cleanup * some fixes to eval pipeline * lint * add asserts for spawn * lint * cleaning up import statements in parallel_datamanager.py * adding new developer documentation if users would like to migrate their custom datamanagers to support new features * removing unnecessary to_device no-op * further updates to documentation * lint * more docs * docs * remove comment * add docs, fix depth dataset with parallel datamanager, fix mask sampling bug * remove profiling * more profile removal * custom_view_processor->custom_image_processor * doc clarification * datamanager doc nit * whitespace * nits * remove stuff from __post_init__, tune num workers more, add random offset in raybatchstream * removing unnecessary assertion, updating docstring because DataManager is no longer an nn.Module * clarifying configuration with num_images_to_sample_from and num_times_to_repeat_images, cleaning up functions * adding logic so that nerfacto users can load_from_disk and customize image batch sizes and repeat parameters * ruff formatting! whoops forgot to format * fixing logic, now if users set load_from_disk to true, datamanager will use 50 and 10. If users set it and specify their own values, we support that as well * adding separate datamanager config so that target can be removed in method_configs --------- Co-authored-by: Justin Kerr <justin.g.kerr@gmail.com> Co-authored-by: Brent Yi <yibrenth@gmail.com>
2026-11-13 03:14:00 +08:00 · 2025-01-31 10:40:01 -08:00
1 parent 189328ec79
commit 4a3e3e6403
@@ -0,0 +1,89 @@
+# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import os
+import subprocess
+import sys
+from typing import Literal
+
+
+def run_command_with_console_output(command, stop_on_output=None):
+    try:
+        process = subprocess.Popen(
+            command,
+            shell=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,  # Redirect stderr to stdout
+            text=True,
+            bufsize=1,  # Line buffered
+            universal_newlines=True,
+        )
+
+        # Read and print output in real-time
+        for line in process.stdout:
+            print(line, end="")
+            sys.stdout.flush()  # Ensure output is printed immediately
+
+            if stop_on_output and stop_on_output in line:
+                print(f"\nDetected '{stop_on_output}'. Stopping the process.")
+                process.terminate()
+                break
+
+        return_code = process.wait()
+        if return_code != 0:
+            print(f"Command failed with return code {return_code}")
+    except Exception as e:
+        print(f"An error occurred: {e}")
+
+
+def run_ns_download_data(scene: Literal["poster", "dozer", "desolation"]):
+    command = f"ns-download-data nerfstudio --capture-name={scene}"
+    run_command_with_console_output(command)
+
+
+def run_ns_train_nerfacto(scene: Literal["poster", "dozer", "desolation"]):
+    dataset_path = f"data/nerfstudio/{scene}"
+    command = f"ns-train nerfacto --data {dataset_path}"
+    run_command_with_console_output(command, stop_on_output="Checkpoint Directory")
+
+
+def run_ns_eval(scene: Literal["poster", "dozer", "desolation"]):
+    timestamp = sorted(os.listdir(f"outputs/{scene}/nerfacto/"))[-1]
+    config_filename = f"outputs/{scene}/nerfacto/{timestamp}/config.yml"
+    command = f"ns-eval --load-config {config_filename} --output-path nerfacto_integration_eval.json"
+    run_command_with_console_output(command)
+
+    with open("nerfacto_integration_eval.json", "r") as f:
+        results = json.load(f)
+
+    assert results["results"]["psnr"] > 20.0, "PSNR was lower than 20"
+    assert results["results"]["ssim"] > 0.7, "SSIM was lower than 0.7"
+
+
+def main():
+    scene = "dozer"  # You can change this to "poster" or "desolation"
+
+    print("Starting data download...")
+    run_ns_download_data(scene)
+
+    print("\nStarting training...")
+    run_ns_train_nerfacto(scene)
+
+    print("\nStarting evaluation...")
+    run_ns_eval(scene)
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,55 @@
+# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import os
+from typing import Literal
+
+from test_nerfacto_integration import run_command_with_console_output, run_ns_download_data
+
+
+def run_ns_train_splatfacto(scene: Literal["poster", "dozer", "desolation"]):
+    dataset_path = f"data/nerfstudio/{scene}"
+    command = f"ns-train splatfacto --data {dataset_path}"
+    run_command_with_console_output(command, stop_on_output="Checkpoint Directory")
+
+
+def run_ns_eval(scene: Literal["poster", "dozer", "desolation"]):
+    timestamp = sorted(os.listdir(f"outputs/{scene}/splatfacto/"))[-1]
+    config_filename = f"outputs/{scene}/splatfacto/{timestamp}/config.yml"
+    command = f"ns-eval --load-config {config_filename} --output-path splatfacto_integration_eval.json"
+    run_command_with_console_output(command)
+
+    with open("splatfacto_integration_eval.json", "r") as f:
+        results = json.load(f)
+
+    assert results["results"]["psnr"] > 20.0, "PSNR was lower than 20"
+    assert results["results"]["ssim"] > 0.7, "SSIM was lower than 0.7"
+
+
+def main():
+    scene = "dozer"  # You can change this to "poster" or "desolation"
+
+    print("Starting data download...")
+    run_ns_download_data(scene)
+
+    print("\nStarting training...")
+    run_ns_train_splatfacto(scene)
+
+    print("\nStarting evaluation...")
+    run_ns_eval(scene)
+
+
+if __name__ == "__main__":
+    main()