Merge pull request #1578 from franz96521/save-raw-volumes-from-python-terminal

save raw volumes python terminal

Author: Tom94

Diff for: 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 for: 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<const fs::path&>(&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 for: 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<<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<vec4> 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<vec4> 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) {