diff --git a/.gitignore b/.gitignore index 6c8fa161399aaa051d390801de0ce2bf8c796d1c..6761425ebaf6889408575f830d14cf3c9139da97 100644 --- a/.gitignore +++ b/.gitignore @@ -12,6 +12,7 @@ /*.dll /*.exe /*.json +*.ingp *.msgpack *.training __pycache__ diff --git a/.gitmodules b/.gitmodules index b3d6c796fffaf79376ad3ccc4a338bc6de5c930c..37d13bdac9a6f3ffedcf2fb1299600a66bf0adcf 100644 --- a/.gitmodules +++ b/.gitmodules @@ -22,3 +22,9 @@ [submodule "dependencies/dlss"] path = dependencies/dlss url = https://github.com/NVIDIA/DLSS +[submodule "dependencies/zstr"] + path = dependencies/zstr + url = https://github.com/Tom94/zstr +[submodule "dependencies/zlib"] + path = dependencies/zlib + url = https://github.com/Tom94/zlib diff --git a/CMakeLists.txt b/CMakeLists.txt index d638a2a8e7c7c64316ceea64a0b7534c95233b5c..c35a8ccd3a367398e6c41c170c824c25252770d8 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -193,6 +193,26 @@ if (Python_FOUND) add_subdirectory("dependencies/pybind11") endif() +# Compile zlib (only on Windows) +if (WIN32) + set(ZLIB_USE_STATIC_LIBS ON CACHE BOOL " " FORCE) + set(ZLIB_BUILD_STATIC_LIBS ON CACHE BOOL " " FORCE) + set(ZLIB_BUILD_SHARED_LIBS OFF CACHE BOOL " " FORCE) + set(SKIP_INSTALL_ALL ON CACHE BOOL " " FORCE) + add_subdirectory("dependencies/zlib") + set_property(TARGET zlibstatic PROPERTY FOLDER "dependencies") + + set(ZLIB_INCLUDE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/dependencies/zlib" CACHE PATH " " FORCE) + set(ZLIB_LIBRARY zlibstatic) + + include_directories(${ZLIB_INCLUDE_DIR} "${CMAKE_CURRENT_BINARY_DIR}/dependencies/zlib") + + list(APPEND NGP_LIBRARIES zlibstatic) +endif() + +add_subdirectory("dependencies/zstr") +list(APPEND NGP_LIBRARIES zstr::zstr) + ############################################################################### # Program diff --git a/dependencies/zlib b/dependencies/zlib new file mode 160000 index 0000000000000000000000000000000000000000..fafd714d8b3ce85a6decf566e3b5b611352eadea --- /dev/null +++ b/dependencies/zlib @@ -0,0 +1 @@ +Subproject commit fafd714d8b3ce85a6decf566e3b5b611352eadea diff --git a/dependencies/zstr b/dependencies/zstr new file mode 160000 index 0000000000000000000000000000000000000000..007f97607865934fb52e08248f152712e428688a --- /dev/null +++ b/dependencies/zstr @@ -0,0 +1 @@ +Subproject commit 007f97607865934fb52e08248f152712e428688a diff --git a/include/neural-graphics-primitives/testbed.h b/include/neural-graphics-primitives/testbed.h index 284f54e45ead18ab08077cf8b5625002ee01c5e9..c766a84a19521949b58e2030113c7aa0d52eab06 100644 --- a/include/neural-graphics-primitives/testbed.h +++ b/include/neural-graphics-primitives/testbed.h @@ -412,7 +412,7 @@ public: void set_fov(float val) ; Eigen::Vector2f fov_xy() const ; void set_fov_xy(const Eigen::Vector2f& val); - void save_snapshot(const std::string& filepath_string, bool include_optimizer_state); + void save_snapshot(const std::string& filepath_string, bool include_optimizer_state, bool compress); void load_snapshot(const std::string& filepath_string); CameraKeyframe copy_camera_to_keyframe() const; void set_camera_from_keyframe(const CameraKeyframe& k); @@ -464,6 +464,7 @@ public: bool m_gather_histograms = false; bool m_include_optimizer_state_in_snapshot = false; + bool m_compress_snapshot = true; bool m_render_ground_truth = false; EGroundTruthRenderMode m_ground_truth_render_mode = EGroundTruthRenderMode::Shade; float m_ground_truth_alpha = 1.0f; @@ -831,7 +832,7 @@ public: char cam_path_path[MAX_PATH_LEN] = "cam.json"; char extrinsics_path[MAX_PATH_LEN] = "extrinsics.json"; char mesh_path[MAX_PATH_LEN] = "base.obj"; - char snapshot_path[MAX_PATH_LEN] = "base.msgpack"; + char snapshot_path[MAX_PATH_LEN] = "base.ingp"; char video_path[MAX_PATH_LEN] = "video.mp4"; } m_imgui; diff --git a/notebooks/instant_ngp.ipynb b/notebooks/instant_ngp.ipynb index 4d2cda757d7d326b57df3cdeabc4197d611034db..02be3e5aabc1bec812b31470488a31fee05ec04f 100644 --- a/notebooks/instant_ngp.ipynb +++ b/notebooks/instant_ngp.ipynb @@ -5846,7 +5846,7 @@ ], "source": [ "train_steps = 2000 #@param {type:\"integer\"}\n", - "snapshot_path = os.path.join(scene_path, f\"{train_steps}.msgpack\")\n", + "snapshot_path = os.path.join(scene_path, f\"{train_steps}.ingp\")\n", "!python ./scripts/run.py {scene_path} --n_steps {train_steps} --save_snapshot {snapshot_path}" ] }, @@ -5863,7 +5863,7 @@ "\n", "Example command:\n", "```\n", - "./build/instant-ngp /data/nerf/fox/2000.msgpack\n", + "./build/instant-ngp /data/nerf/fox/2000.ingp\n", "```\n", "\n", "After you're done, **upload `base_cam.json` to the root folder of your scene.**" diff --git a/scripts/run.py b/scripts/run.py index cbcf4a3fba0b40c00b9cb91b3aba366f69cda956..72cc9f947795d01fbd06a6d0229fe5827154c22b 100644 --- a/scripts/run.py +++ b/scripts/run.py @@ -33,10 +33,10 @@ def parse_args(): parser.add_argument("--mode", default="", type=str, help=argparse.SUPPRESS) # deprecated parser.add_argument("--network", default="", help="Path to the network config. Uses the scene's default if unspecified.") - parser.add_argument("--load_snapshot", "--snapshot", default="", help="Load this snapshot before training. recommended extension: .msgpack") - parser.add_argument("--save_snapshot", default="", help="Save this snapshot after training. recommended extension: .msgpack") + parser.add_argument("--load_snapshot", "--snapshot", default="", help="Load this snapshot before training. recommended extension: .ingp/.msgpack") + parser.add_argument("--save_snapshot", default="", help="Save this snapshot after training. recommended extension: .ingp/.msgpack") - parser.add_argument("--nerf_compatibility", action="store_true", help="Matches parameters with original NeRF. Can cause slowness and worse results on some scenes.") + parser.add_argument("--nerf_compatibility", action="store_true", help="Matches parameters with original NeRF. Can cause slowness and worse results on some scenes, but helps with high PSNR on synthetic scenes.") parser.add_argument("--test_transforms", default="", help="Path to a nerf style transforms json from which we will compute PSNR.") parser.add_argument("--near_distance", default=-1, type=float, help="Set the distance from the camera at which training rays start for nerf. <0 means use ngp default") parser.add_argument("--exposure", default=0.0, type=float, help="Controls the brightness of the image. Positive numbers increase brightness, negative numbers decrease it.") diff --git a/scripts/scenes.py b/scripts/scenes.py index ca9e30c39ca87b5f99819b3a7cbfecf4ee554eda..2cc47088ec4aa36d5b67f599c277c68634bdb224 100644 --- a/scripts/scenes.py +++ b/scripts/scenes.py @@ -222,7 +222,7 @@ def setup_colored_sdf(testbed, scene, softshadow=True): testbed.scale = testbed.scale * 1.13 def default_snapshot_filename(scene_info): - filename = "base.msgpack" + filename = f"base.ingp" if scene_info["dataset"]: filename = f"{os.path.splitext(scene_info['dataset'])[0]}_{filename}" return os.path.join(scene_info["data_dir"], filename) diff --git a/src/testbed.cu b/src/testbed.cu index 20e5c6538852bd3ec0cb5b2265c97a2ffcd227b7..11c9acdb8de1856eb56c0fa7efa8129ead93cdca 100644 --- a/src/testbed.cu +++ b/src/testbed.cu @@ -41,6 +41,8 @@ #include <stb_image/stb_image.h> #include <stb_image/stb_image_write.h> +#include <zstr.hpp> + #include <fstream> #include <set> #include <unordered_set> @@ -86,7 +88,7 @@ NGP_NAMESPACE_BEGIN std::atomic<size_t> g_total_n_bytes_allocated{0}; -json merge_parent_network_config(const json &child, const fs::path &child_filename) { +json merge_parent_network_config(const json& child, const fs::path& child_filename) { if (!child.contains("parent")) { return child; } @@ -121,7 +123,7 @@ void Testbed::update_imgui_paths() { snprintf(m_imgui.cam_path_path, sizeof(m_imgui.cam_path_path), "%s", get_filename_in_data_path_with_suffix(m_data_path, m_network_config_path, "_cam.json").c_str()); snprintf(m_imgui.extrinsics_path, sizeof(m_imgui.extrinsics_path), "%s", get_filename_in_data_path_with_suffix(m_data_path, m_network_config_path, "_extrinsics.json").c_str()); snprintf(m_imgui.mesh_path, sizeof(m_imgui.mesh_path), "%s", get_filename_in_data_path_with_suffix(m_data_path, m_network_config_path, ".obj").c_str()); - snprintf(m_imgui.snapshot_path, sizeof(m_imgui.snapshot_path), "%s", get_filename_in_data_path_with_suffix(m_data_path, m_network_config_path, ".msgpack").c_str()); + snprintf(m_imgui.snapshot_path, sizeof(m_imgui.snapshot_path), "%s", get_filename_in_data_path_with_suffix(m_data_path, m_network_config_path, ".ingp").c_str()); snprintf(m_imgui.video_path, sizeof(m_imgui.video_path), "%s", get_filename_in_data_path_with_suffix(m_data_path, m_network_config_path, "_video.mp4").c_str()); } @@ -215,7 +217,7 @@ fs::path Testbed::find_network_config(const fs::path& network_config_path) { } json Testbed::load_network_config(const fs::path& network_config_path) { - bool is_snapshot = equals_case_insensitive(network_config_path.extension(), "msgpack"); + bool is_snapshot = equals_case_insensitive(network_config_path.extension(), "msgpack") || equals_case_insensitive(network_config_path.extension(), "ingp"); if (network_config_path.empty() || !network_config_path.exists()) { throw std::runtime_error{fmt::format("Network {} '{}' does not exist.", is_snapshot ? "snapshot" : "config", network_config_path.str())}; } @@ -223,14 +225,20 @@ json Testbed::load_network_config(const fs::path& network_config_path) { tlog::info() << "Loading network " << (is_snapshot ? "snapshot" : "config") << " from: " << network_config_path; json result; - if (equals_case_insensitive(network_config_path.extension(), "json")) { + if (is_snapshot) { + if (equals_case_insensitive(network_config_path.extension(), "ingp")) { + // zstr::ifstream applies zlib compression. + zstr::ifstream f{network_config_path.str(), std::ios::in | std::ios::binary}; + result = json::from_msgpack(f); + } else { + std::ifstream f{network_config_path.str(), std::ios::in | std::ios::binary}; + result = json::from_msgpack(f); + } + // we assume parent pointers are already resolved in snapshots. + } else if (equals_case_insensitive(network_config_path.extension(), "json")) { std::ifstream f{network_config_path.str()}; result = json::parse(f, nullptr, true, true); result = merge_parent_network_config(result, network_config_path); - } else if (equals_case_insensitive(network_config_path.extension(), "msgpack")) { - std::ifstream f{network_config_path.str(), std::ios::in | std::ios::binary}; - result = json::from_msgpack(f); - // we assume parent pointers are already resolved in snapshots. } return result; @@ -289,7 +297,7 @@ void Testbed::load_file(const std::string& file_path) { return; } - if (ends_with_case_insensitive(file_path, ".msgpack")) { + if (ends_with_case_insensitive(file_path, ".ingp") || ends_with_case_insensitive(file_path, ".msgpack")) { load_snapshot(file_path); return; } @@ -1218,7 +1226,7 @@ void Testbed::imgui() { ImGui::Text("Snapshot"); ImGui::SameLine(); if (ImGui::Button("Save")) { - save_snapshot(m_imgui.snapshot_path, m_include_optimizer_state_in_snapshot); + save_snapshot(m_imgui.snapshot_path, m_include_optimizer_state_in_snapshot, m_compress_snapshot); } ImGui::SameLine(); if (ImGui::Button("Load")) { @@ -1237,6 +1245,16 @@ void Testbed::imgui() { ImGui::SameLine(); ImGui::Checkbox("w/ optimizer state", &m_include_optimizer_state_in_snapshot); ImGui::InputText("File##Snapshot file path", m_imgui.snapshot_path, sizeof(m_imgui.snapshot_path)); + ImGui::SameLine(); + + bool can_compress = ends_with_case_insensitive(m_imgui.snapshot_path, ".ingp"); + + if (!can_compress) { + ImGui::BeginDisabled(); + m_compress_snapshot = false; + } + ImGui::Checkbox("Compress", &m_compress_snapshot); + if (!can_compress) ImGui::EndDisabled(); } if (m_testbed_mode == ETestbedMode::Nerf || m_testbed_mode == ETestbedMode::Sdf) { @@ -3481,7 +3499,7 @@ void Testbed::gather_histograms() { // Increment this number when making a change to the snapshot format static const size_t SNAPSHOT_FORMAT_VERSION = 1; -void Testbed::save_snapshot(const std::string& filepath_string, bool include_optimizer_state) { +void Testbed::save_snapshot(const std::string& filepath_string, bool include_optimizer_state, bool compress) { fs::path filepath = filepath_string; m_network_config["snapshot"] = m_trainer->serialize(include_optimizer_state); @@ -3516,8 +3534,16 @@ void Testbed::save_snapshot(const std::string& filepath_string, bool include_opt } m_network_config_path = filepath; - std::ofstream f(m_network_config_path.str(), std::ios::out | std::ios::binary); - json::to_msgpack(m_network_config, f); + if (equals_case_insensitive(m_network_config_path.extension(), "ingp")) { + // zstr::ofstream applies zlib compression. + zstr::ofstream f{m_network_config_path.str(), std::ios::out | std::ios::binary, compress ? Z_DEFAULT_COMPRESSION : Z_NO_COMPRESSION}; + json::to_msgpack(m_network_config, f); + } else { + std::ofstream f{m_network_config_path.str(), std::ios::out | std::ios::binary}; + json::to_msgpack(m_network_config, f); + } + + tlog::success() << "Saved snapshot '" << filepath_string << "'"; } void Testbed::load_snapshot(const std::string& filepath_string) { @@ -3526,10 +3552,7 @@ void Testbed::load_snapshot(const std::string& filepath_string) { throw std::runtime_error{fmt::format("File '{}' does not contain a snapshot.", filepath_string)}; } - m_network_config_path = filepath_string; - const auto& snapshot = config["snapshot"]; - if (snapshot.value("version", 0) < SNAPSHOT_FORMAT_VERSION) { throw std::runtime_error{"Snapshot uses an old format and can not be loaded."}; }