diff --git a/include/neural-graphics-primitives/testbed.h b/include/neural-graphics-primitives/testbed.h index 83e224207..278e9cc42 100644 --- a/include/neural-graphics-primitives/testbed.h +++ b/include/neural-graphics-primitives/testbed.h @@ -485,6 +485,7 @@ class Testbed { vec2 fov_xy() const ; void set_fov_xy(const vec2& val); void save_snapshot(const fs::path& path, bool include_optimizer_state, bool compress); + void save_raw_volumes(const fs::path &filename, int res, BoundingBox aabb, bool flip_y_and_z_axes); void load_snapshot(nlohmann::json config); void load_snapshot(const fs::path& path); void load_snapshot(std::istream& stream, bool is_compressed = true); diff --git a/src/python_api.cu b/src/python_api.cu index da62ecca1..7a7acc3e7 100644 --- a/src/python_api.cu +++ b/src/python_api.cu @@ -443,6 +443,7 @@ PYBIND11_MODULE(pyngp, m) { .def("n_params", &Testbed::n_params, "Number of trainable parameters") .def("n_encoding_params", &Testbed::n_encoding_params, "Number of trainable parameters in the encoding") .def("save_snapshot", &Testbed::save_snapshot, py::arg("path"), py::arg("include_optimizer_state")=false, py::arg("compress")=true, "Save a snapshot of the currently trained model. Optionally compressed (only when saving '.ingp' files).") + .def("save_raw_volumes", &Testbed::save_raw_volumes) .def("load_snapshot", py::overload_cast(&Testbed::load_snapshot), py::arg("path"), "Load a previously saved snapshot") .def("load_camera_path", &Testbed::load_camera_path, py::arg("path"), "Load a camera path") .def("load_file", &Testbed::load_file, py::arg("path"), "Load a file and automatically determine how to handle it. Can be a snapshot, dataset, network config, or camera path.") diff --git a/src/testbed.cu b/src/testbed.cu index a42734175..01d216db0 100644 --- a/src/testbed.cu +++ b/src/testbed.cu @@ -1624,24 +1624,7 @@ void Testbed::imgui() { save_rgba_grid_to_png_sequence(rgba, dir, res3d, flip_y_and_z_axes); } if (imgui_colored_button("Save raw volumes", 0.4f)) { - auto effective_view_dir = flip_y_and_z_axes ? vec3{0.0f, 1.0f, 0.0f} : vec3{0.0f, 0.0f, 1.0f}; - auto old_local = m_render_aabb_to_local; - auto old_aabb = m_render_aabb; - m_render_aabb_to_local = mat3::identity(); - auto dir = m_data_path.is_directory() || m_data_path.empty() ? (m_data_path / "volume_raw") : (m_data_path.parent_path() / fmt::format("{}_volume_raw", m_data_path.filename())); - if (!dir.exists()) { - fs::create_directory(dir); - } - - for (int cascade = 0; (1< rgba = get_rgba_on_grid(res3d, effective_view_dir, true, 0.0f, true); - save_rgba_grid_to_raw_file(rgba, dir, res3d, flip_y_and_z_axes, cascade); - } - m_render_aabb_to_local = old_local; - m_render_aabb = old_aabb; + save_raw_volumes(m_data_path, m_mesh.res, {}, flip_y_and_z_axes); } } @@ -4754,6 +4737,41 @@ void Testbed::save_snapshot(const fs::path& path, bool include_optimizer_state, tlog::success() << "Saved snapshot '" << path.str() << "'"; } +void Testbed::save_raw_volumes(const fs::path &filename, int res, BoundingBox aabb, bool flip_y_and_z_axes) +{ + auto effective_view_dir = flip_y_and_z_axes ? vec3{0.0f, 1.0f, 0.0f} : vec3{0.0f, 0.0f, 1.0f}; + mat3 render_aabb_to_local = mat3(1.0f); + + if (aabb.is_empty()) + { + aabb = m_testbed_mode == ETestbedMode::Nerf ? m_render_aabb : m_aabb; + render_aabb_to_local = m_render_aabb_to_local; + } + + if (m_testbed_mode != ETestbedMode::Nerf) + { + throw std::runtime_error{"Raw volume export is only supported for NeRF."}; + } + + auto res3d = get_marching_cubes_res(res, aabb); + + std::string flipped = flip_y_and_z_axes ? "_flipedYZ" : ""; + auto dir =( filename.is_directory() || filename.empty() ? (filename / fmt::format("volume_raw{}",flipped)) : (filename.parent_path() / fmt::format("{}_volume_raw{}", filename.filename(),flipped))) ; + if (!dir.exists()) + { + fs::create_directory(dir); + } + + for (int cascade = 0; (1 << cascade) <= m_aabb.diag().x + 0.5f; ++cascade) + { + float radius = (1 << cascade) * 0.5f; + m_render_aabb = BoundingBox(vec3(0.5f - radius), vec3(0.5f + radius)); + // Dump raw density values that the user can then convert to alpha as they please. + GPUMemory rgba = get_rgba_on_grid(res3d, effective_view_dir, true, 0.0f, true); + save_rgba_grid_to_raw_file(rgba, dir, res3d, flip_y_and_z_axes, cascade); + } +} + void Testbed::load_snapshot(nlohmann::json config) { const auto& snapshot = config["snapshot"]; if (snapshot.value("version", 0) < SNAPSHOT_FORMAT_VERSION) {