Merge pull request #342 from NVlabs/misc-improvements

Misc improvements

.gitignore

@@ -5,11 +5,11 @@
 /out
 /results
 /tmp
+/venv
 /video
 /*.json
 *.msgpack
 *.training
 __pycache__
 .DS_Store
 imgui.ini
-/venv

configs/sdf/frequency.json

@@ -2,7 +2,10 @@
 	"parent" : "base.json",
 	"encoding": {
 		"otype": "Frequency",
-		"n_frequencies": 8
+		"n_frequencies": 10
+	},
+	"loss": {
+		"otype": "RelativeL2"
 	},
 	"optimizer": {
 		"nested": {
@@ -16,6 +19,6 @@
 		"activation": "ReLU",
 		"output_activation": "None",
 		"n_neurons": 128,
-		"n_hidden_layers": 6
+		"n_hidden_layers": 8
 	}
 }

include/neural-graphics-primitives/adam_optimizer.h

@@ -64,7 +64,7 @@ class AdamOptimizer {
 		m_state = {};
 	}
 
-	private:
+private:
 	struct State {
 		uint32_t iter = 0;
 		T first_moment = T::Zero();
@@ -130,7 +130,7 @@ class RotationAdamOptimizer {
 		m_state = {};
 	}
 
-	private:
+private:
 	struct State {
 		uint32_t iter = 0;
 		Eigen::Vector3f first_moment = Eigen::Vector3f::Zero();

include/neural-graphics-primitives/common.h

@@ -15,6 +15,7 @@
 
 #pragma once
 
+
 #include <tinylogger/tinylogger.h>
 
 // Eigen uses __device__ __host__ on a bunch of defaulted constructors.
@@ -143,13 +144,20 @@ struct Ray {
 	Eigen::Vector3f d;
 };
 
+struct TrainingXForm {
+	Eigen::Matrix<float, 3, 4> start;
+	Eigen::Matrix<float, 3, 4> end;
+};
+
+enum class ECameraDistortionMode : int {
+	None,
+	Iterative,
+	FTheta,
+};
+
 struct CameraDistortion {
-	float params[4] = {};
-#ifdef __NVCC__
-	inline __host__ __device__ bool is_zero() const {
-		return params[0] == 0.0f && params[1] == 0.0f && params[2] == 0.0f && params[3] == 0.0f;
-	}
-#endif
+	ECameraDistortionMode mode = ECameraDistortionMode::None;
+	float params[7] = {};
 };
 
 #ifdef __NVCC__

include/neural-graphics-primitives/common_device.cuh

@@ -200,7 +200,6 @@ __device__ __host__ inline void iterative_camera_undistortion(const T* params, T
 	*v = x(1);
 }
 
-
 inline __host__ __device__ Ray pixel_to_ray_pinhole(
 	uint32_t spp,
 	const Eigen::Vector2i& pixel,
@@ -225,6 +224,25 @@ inline __host__ __device__ Ray pixel_to_ray_pinhole(
 	return {origin, dir};
 }
 
+inline __host__ __device__ Eigen::Matrix<float, 3, 4> get_xform_given_rolling_shutter(const TrainingXForm &training_xform, const Eigen::Vector4f &rolling_shutter, const Eigen::Vector2f &uv, float motionblur_time) {
+	float pixel_t = rolling_shutter.x() + rolling_shutter.y() * uv.x() + rolling_shutter.z() * uv.y() + rolling_shutter.w() * motionblur_time;
+	return training_xform.start + (training_xform.end - training_xform.start) * pixel_t;
+}
+
+inline __host__ __device__ Eigen::Vector3f f_theta_undistortion(const Eigen::Vector2f &uv, const Eigen::Vector2f &screen_center, const CameraDistortion& camera_distortion, const Eigen::Vector3f& error_direction) {
+	// we take f_theta intrinsics to be: resx, resy, r0, r1, r2, r3; we rescale to whatever res the intrinsics specify.
+	float xpix = (uv.x() - screen_center.x()) * camera_distortion.params[5];
+	float ypix = (uv.y() - screen_center.y()) * camera_distortion.params[6];
+	float norm = sqrtf(xpix*xpix + ypix*ypix);
+	float alpha = camera_distortion.params[0] + norm * (camera_distortion.params[1] + norm * (camera_distortion.params[2] + norm * (camera_distortion.params[3] + norm * camera_distortion.params[4])));
+	float sin_alpha, cos_alpha;
+	sincosf(alpha, &sin_alpha, &cos_alpha);
+	if (cos_alpha <= std::numeric_limits<float>::min() || norm == 0.f)
+		return error_direction;
+	sin_alpha *= 1.f/norm;
+	return { sin_alpha * xpix, sin_alpha * ypix, cos_alpha };
+}
+
 inline __host__ __device__ Ray pixel_to_ray(
 	uint32_t spp,
 	const Eigen::Vector2i& pixel,
@@ -242,13 +260,21 @@ inline __host__ __device__ Ray pixel_to_ray(
 	Eigen::Vector2f offset = ld_random_pixel_offset(snap_to_pixel_centers ? 0 : spp, pixel.x(), pixel.y());
 	auto uv = (pixel.cast<float>() + offset).cwiseQuotient(resolution.cast<float>());
 
-	Eigen::Vector3f dir = {
-		(uv.x() - screen_center.x()) * (float)resolution.x() / focal_length.x(),
-		(uv.y() - screen_center.y()) * (float)resolution.y() / focal_length.y(),
-		1.0f
-	};
-	if (!camera_distortion.is_zero()) {
-		iterative_camera_undistortion(camera_distortion.params, &dir.x(), &dir.y());
+	Eigen::Vector3f dir;
+	if (camera_distortion.mode == ECameraDistortionMode::FTheta) {
+		dir = f_theta_undistortion(uv, screen_center, camera_distortion, {1000.f, 0.f, 0.f});
+		if (dir.x() == 1000.f) {
+			return {{1000.f, 0.f, 0.f}, {0.f, 0.f, 1.f}}; // return a point outside the aabb so the pixel is not rendered
+		}
+	} else {
+		dir = {
+			(uv.x() - screen_center.x()) * (float)resolution.x() / focal_length.x(),
+			(uv.y() - screen_center.y()) * (float)resolution.y() / focal_length.y(),
+			1.0f
+		};
+		if (camera_distortion.mode == ECameraDistortionMode::Iterative) {
+			iterative_camera_undistortion(camera_distortion.params, &dir.x(), &dir.y());
+		}
 	}
 	if (distortion_data) {
 		dir.head<2>() += read_image<2>(distortion_data, distortion_resolution, uv);

include/neural-graphics-primitives/json_binding.h

@@ -81,17 +81,51 @@ inline void from_json(const nlohmann::json& j, BoundingBox& box) {
 }
 
 inline void to_json(nlohmann::json& j, const CameraDistortion& dist) {
-	j["k1"] = dist.params[0];
-	j["k2"] = dist.params[1];
-	j["p1"] = dist.params[2];
-	j["p2"] = dist.params[3];
+	if (dist.mode == ECameraDistortionMode::Iterative) {
+		j["k1"] = dist.params[0];
+		j["k2"] = dist.params[1];
+		j["p1"] = dist.params[2];
+		j["p2"] = dist.params[3];
+	} else if (dist.mode == ECameraDistortionMode::FTheta) {
+		j["ftheta_p0"] = dist.params[0];
+		j["ftheta_p1"] = dist.params[1];
+		j["ftheta_p2"] = dist.params[2];
+		j["ftheta_p3"] = dist.params[3];
+		j["ftheta_p4"] = dist.params[4];
+		j["w"] = dist.params[5];
+		j["h"] = dist.params[6];
+	}
 }
 
 inline void from_json(const nlohmann::json& j, CameraDistortion& dist) {
-	dist.params[0] = j.at("k1");
-	dist.params[1] = j.at("k2");
-	dist.params[2] = j.at("p1");
-	dist.params[3] = j.at("p2");
+	if (j.contains("k1")) {
+		dist.mode = ECameraDistortionMode::Iterative;
+		dist.params[0] = j.at("k1");
+		dist.params[1] = j.at("k2");
+		dist.params[2] = j.at("p1");
+		dist.params[3] = j.at("p2");
+	} else if (j.contains("ftheta_p0")) {
+		dist.mode = ECameraDistortionMode::FTheta;
+		dist.params[0] = j.at("ftheta_p0");
+		dist.params[1] = j.at("ftheta_p1");
+		dist.params[2] = j.at("ftheta_p2");
+		dist.params[3] = j.at("ftheta_p3");
+		dist.params[4] = j.at("ftheta_p4");
+		dist.params[5] = j.at("w");
+		dist.params[6] = j.at("h");
+	} else {
+		dist.mode = ECameraDistortionMode::None;
+	}
+}
+
+inline void from_json(const nlohmann::json& j, TrainingXForm& x) {
+	from_json(j.at("start"), x.start);
+	from_json(j.at("end"), x.end);
+}
+
+inline void to_json(nlohmann::json& j, const TrainingXForm& x) {
+	to_json(j["start"], x.start);
+	to_json(j["end"], x.end);
 }
 
 inline void to_json(nlohmann::json& j, const NerfDataset& dataset) {
@@ -102,6 +136,7 @@ inline void to_json(nlohmann::json& j, const NerfDataset& dataset) {
 		to_json(j["metadata"].at(i)["focal_length"], dataset.metadata[i].focal_length);
 		to_json(j["metadata"].at(i)["camera_distortion"], dataset.metadata[i].camera_distortion);
 		to_json(j["metadata"].at(i)["principal_point"], dataset.metadata[i].principal_point);
+		to_json(j["metadata"].at(i)["rolling_shutter"], dataset.metadata[i].rolling_shutter);
 		to_json(j["xforms"].at(i), dataset.xforms[i]);
 	}
 	j["render_aabb"] = dataset.render_aabb;
@@ -114,6 +149,7 @@ inline void to_json(nlohmann::json& j, const NerfDataset& dataset) {
 	// global distortion removed
 	//j["camera_distortion"] = dataset.camera_distortion;
 	//to_json(j["principal_point"], dataset.principal_point);
+	//to_json(j["rolling_shutter"], dataset.rolling_shutter);
 	j["from_mitsuba"] = dataset.from_mitsuba;
 	j["is_hdr"] = dataset.is_hdr;
 	j["wants_importance_sampling"] = dataset.wants_importance_sampling;
@@ -126,6 +162,7 @@ inline void from_json(const nlohmann::json& j, NerfDataset& dataset) {
 	TrainingImageMetadata global_metadata = {};
 	if (j.contains("camera_distortion")) from_json(j.at("camera_distortion"), global_metadata.camera_distortion);
 	if (j.contains("principal_point")) from_json(j.at("principal_point"), global_metadata.principal_point);
+	if (j.contains("rolling_shutter")) from_json(j.at("rolling_shutter"), global_metadata.rolling_shutter);
 	if (j.contains("focal_length")) from_json(j.at("focal_length"), global_metadata.focal_length);
 
 	for (size_t i = 0; i < dataset.n_images; ++i) {
@@ -137,6 +174,7 @@ inline void from_json(const nlohmann::json& j, NerfDataset& dataset) {
 			auto &ji = j["metadata"].at(i);
 			from_json(ji.at("focal_length"),dataset.metadata[i].focal_length);
 			from_json(ji.at("principal_point"),dataset.metadata[i].principal_point);
+			from_json(ji.at("rolling_shutter"),dataset.metadata[i].rolling_shutter);
 			from_json(ji.at("camera_distortion"),dataset.metadata[i].camera_distortion);
 		}
 	}

include/neural-graphics-primitives/nerf_loader.h

@@ -29,18 +29,21 @@ struct TrainingImageMetadata {
 	CameraDistortion camera_distortion = {};
 	Eigen::Vector2f principal_point = Eigen::Vector2f::Constant(0.5f);
 	Eigen::Vector2f focal_length = Eigen::Vector2f::Constant(1000.f);
+	Eigen::Vector4f rolling_shutter = Eigen::Vector4f::Zero();
 
 	// TODO: replace this with more generic float[] of task-specific metadata.
 	Eigen::Vector3f light_dir = Eigen::Vector3f::Constant(0.f);
 };
 
 struct NerfDataset {
 	std::vector<TrainingImageMetadata> metadata;
-	std::vector<Eigen::Matrix<float, 3, 4>> xforms;
+	std::vector<TrainingXForm> xforms;
 	tcnn::GPUMemory<__half> images_data;
 	tcnn::GPUMemory<float> sharpness_data;
 	Eigen::Vector2i sharpness_resolution = {0, 0};
 	tcnn::GPUMemory<float> envmap_data;
+	tcnn::GPUMemory<Ray> rays_data;
+
 	BoundingBox render_aabb = {};
 	Eigen::Vector3f up = {0.0f, 1.0f, 0.0f};
 	Eigen::Vector3f offset = {0.0f, 0.0f, 0.0f};
@@ -58,8 +61,6 @@ struct NerfDataset {
 
 	void set_training_image(int frame_idx, const float *pixels);
 
-	tcnn::GPUMemory<Ray> rays_data;
-
 	Eigen::Vector3f nerf_direction_to_ngp(const Eigen::Vector3f& nerf_dir) {
 		Eigen::Vector3f result = nerf_dir;
 		if (from_mitsuba) {
@@ -108,8 +109,10 @@ struct NerfDataset {
 		return result;
 	}
 
-	void nerf_ray_to_ngp(Ray& ray) {
+	void nerf_ray_to_ngp(Ray& ray, bool scale_direction = false) {
 		ray.o = ray.o * scale + offset;
+		if (scale_direction)
+			ray.d *= scale;
 
 		float tmp = ray.o[0];
 		ray.o[0] = ray.o[1];