From 8c813f83490ff983daa1a440be4b2318a43394fe Mon Sep 17 00:00:00 2001
From: Thomas Rouch <thom01.rouch@gmail.com>
Date: Thu, 13 Jul 2023 10:06:39 +0200
Subject: [PATCH] :sparkles: avoid early saturation when rendering GT depth of
 large scenes

---
 include/neural-graphics-primitives/render_buffer.h | 1 +
 include/neural-graphics-primitives/testbed.h       | 1 +
 src/render_buffer.cu                               | 5 ++++-
 src/testbed.cu                                     | 2 ++
 4 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/include/neural-graphics-primitives/render_buffer.h b/include/neural-graphics-primitives/render_buffer.h
index 6bbb4bad3..b1aff4986 100644
--- a/include/neural-graphics-primitives/render_buffer.h
+++ b/include/neural-graphics-primitives/render_buffer.h
@@ -251,6 +251,7 @@ class CudaRenderBuffer {
 		float alpha,
 		const float* __restrict__ depth,
 		float depth_scale,
+		float max_depth,
 		const ivec2& resolution,
 		int fov_axis,
 		float zoom,
diff --git a/include/neural-graphics-primitives/testbed.h b/include/neural-graphics-primitives/testbed.h
index 75aa028dd..6695a1127 100644
--- a/include/neural-graphics-primitives/testbed.h
+++ b/include/neural-graphics-primitives/testbed.h
@@ -542,6 +542,7 @@ class Testbed {
 	bool m_render_ground_truth = false;
 	EGroundTruthRenderMode m_ground_truth_render_mode = EGroundTruthRenderMode::Shade;
 	float m_ground_truth_alpha = 1.0f;
+	float m_ground_truth_max_depth = 1.0f;
 
 	bool m_train = false;
 	bool m_training_data_available = false;
diff --git a/src/render_buffer.cu b/src/render_buffer.cu
index 0ba1fad1e..462001f56 100644
--- a/src/render_buffer.cu
+++ b/src/render_buffer.cu
@@ -434,6 +434,7 @@ __global__ void overlay_depth_kernel(
 	float alpha,
 	const float* __restrict__ depth,
 	float depth_scale,
+	float max_depth,
 	ivec2 image_resolution,
 	int fov_axis,
 	float zoom,
@@ -467,7 +468,7 @@ __global__ void overlay_depth_kernel(
 		color = {0.0f, 0.0f, 0.0f, 0.0f};
 	} else {
 		float depth_value = depth[srcidx] * depth_scale;
-		vec3 c = colormap_turbo(depth_value);
+		vec3 c = colormap_turbo(depth_value/max_depth);
 		color = {c[0], c[1], c[2], 1.0f};
 	}
 
@@ -731,6 +732,7 @@ void CudaRenderBuffer::overlay_depth(
 	float alpha,
 	const float* __restrict__ depth,
 	float depth_scale,
+	float max_depth,
 	const ivec2& image_resolution,
 	int fov_axis,
 	float zoom,
@@ -745,6 +747,7 @@ void CudaRenderBuffer::overlay_depth(
 		alpha,
 		depth,
 		depth_scale,
+		max_depth,
 		image_resolution,
 		fov_axis,
 		zoom,
diff --git a/src/testbed.cu b/src/testbed.cu
index 78b000cc6..469f2827d 100644
--- a/src/testbed.cu
+++ b/src/testbed.cu
@@ -1286,6 +1286,7 @@ void Testbed::imgui() {
 
 			accum_reset |= ImGui::Combo("Groundtruth render mode", (int*)&m_ground_truth_render_mode, GroundTruthRenderModeStr);
 			accum_reset |= ImGui::SliderFloat("Groundtruth alpha", &m_ground_truth_alpha, 0.0f, 1.0f, "%.02f", ImGuiSliderFlags_AlwaysClamp);
+			accum_reset |= ImGui::SliderFloat("Groundtruth Max Depth", &m_ground_truth_max_depth, 1e-3f, 10.0f, "%.02f", ImGuiSliderFlags_AlwaysClamp);
 
 			bool lens_changed = ImGui::Checkbox("Apply lens distortion", &m_nerf.render_with_lens_distortion);
 			if (m_nerf.render_with_lens_distortion) {
@@ -4520,6 +4521,7 @@ void Testbed::render_frame_epilogue(
 					m_ground_truth_alpha,
 					metadata.depth,
 					1.0f/m_nerf.training.dataset.scale,
+					m_ground_truth_max_depth,
 					metadata.resolution,
 					m_fov_axis,
 					m_zoom,