-
Thomas Pickles authored
Options can now be selected at command line without requiring separate builds of the code
4f593d08
testbed.h 40.30 KiB
/*
* Copyright (c) 2020-2022, NVIDIA CORPORATION. All rights reserved.
*
* NVIDIA CORPORATION and its licensors retain all intellectual property
* and proprietary rights in and to this software, related documentation
* and any modifications thereto. Any use, reproduction, disclosure or
* distribution of this software and related documentation without an express
* license agreement from NVIDIA CORPORATION is strictly prohibited.
*/
/** @file testbed.h
* @author Thomas Müller & Alex Evans, NVIDIA
*/
#pragma once
#include <neural-graphics-primitives/adam_optimizer.h>
#include <neural-graphics-primitives/camera_path.h>
#include <neural-graphics-primitives/common.h>
#include <neural-graphics-primitives/discrete_distribution.h>
#include <neural-graphics-primitives/nerf.h>
#include <neural-graphics-primitives/nerf_loader.h>
#include <neural-graphics-primitives/render_buffer.h>
#include <neural-graphics-primitives/sdf.h>
#include <neural-graphics-primitives/shared_queue.h>
#include <neural-graphics-primitives/thread_pool.h>
#include <neural-graphics-primitives/trainable_buffer.cuh>
#ifdef NGP_GUI
# include <neural-graphics-primitives/openxr_hmd.h>
#endif
#include <tiny-cuda-nn/multi_stream.h>
#include <tiny-cuda-nn/random.h>
#include <json/json.hpp>
#ifdef NGP_PYTHON
# include <pybind11/pybind11.h>
# include <pybind11/numpy.h>
#endif
#include <thread>
struct GLFWwindow;
TCNN_NAMESPACE_BEGIN
template <typename T> class Loss;
template <typename T> class Optimizer;
template <typename T> class Encoding;
template <typename T, typename PARAMS_T> class Network;
template <typename T, typename PARAMS_T, typename COMPUTE_T> class Trainer;
template <uint32_t N_DIMS, uint32_t RANK, typename T> class TrainableBuffer;
TCNN_NAMESPACE_END
NGP_NAMESPACE_BEGIN
template <typename T> class NerfNetwork;
class TriangleOctree;
class TriangleBvh;
struct Triangle;
class GLTexture;
class Testbed {
public:
EIGEN_MAKE_ALIGNED_OPERATOR_NEW
Testbed(ETestbedMode mode = ETestbedMode::None);
~Testbed();
Testbed(ETestbedMode mode, const fs::path& data_path) : Testbed(mode) { load_training_data(data_path); } Testbed(ETestbedMode mode, const fs::path& data_path, const fs::path& network_config_path) : Testbed(mode, data_path) { reload_network_from_file(network_config_path); }
Testbed(ETestbedMode mode, const fs::path& data_path, const nlohmann::json& network_config) : Testbed(mode, data_path) { reload_network_from_json(network_config); }
bool clear_tmp_dir();
void update_imgui_paths();
void load_training_data(const fs::path& path);
void reload_training_data();
void clear_training_data();
void set_mode(ETestbedMode mode);
using distance_fun_t = std::function<void(uint32_t, const Eigen::Vector3f*, float*, cudaStream_t)>;
using normals_fun_t = std::function<void(uint32_t, const Eigen::Vector3f*, Eigen::Vector3f*, cudaStream_t)>;
class SphereTracer {
public:
SphereTracer() {}
void init_rays_from_camera(
uint32_t spp,
const Eigen::Vector2i& resolution,
const Eigen::Vector2f& focal_length,
const Eigen::Matrix<float, 3, 4>& camera_matrix,
const Eigen::Vector2f& screen_center,
const Eigen::Vector3f& parallax_shift,
bool snap_to_pixel_centers,
const BoundingBox& aabb,
float floor_y,
float near_distance,
float plane_z,
float aperture_size,
const Foveation& foveation,
const Buffer2DView<const Eigen::Array4f>& envmap,
Eigen::Array4f* frame_buffer,
float* depth_buffer,
const Buffer2DView<const uint8_t>& hidden_area_mask,
const TriangleOctree* octree,
uint32_t n_octree_levels,
cudaStream_t stream
);
void init_rays_from_data(uint32_t n_elements, const RaysSdfSoa& data, cudaStream_t stream);
uint32_t trace_bvh(TriangleBvh* bvh, const Triangle* triangles, cudaStream_t stream);
uint32_t trace(
const distance_fun_t& distance_function,
float zero_offset,
float distance_scale,
float maximum_distance,
const BoundingBox& aabb,
const float floor_y,
const TriangleOctree* octree,
uint32_t n_octree_levels,
cudaStream_t stream
);
void enlarge(size_t n_elements, cudaStream_t stream);
RaysSdfSoa& rays_hit() { return m_rays_hit; }
RaysSdfSoa& rays_init() { return m_rays[0]; }
uint32_t n_rays_initialized() const { return m_n_rays_initialized; }
void set_trace_shadow_rays(bool val) { m_trace_shadow_rays = val; }
void set_shadow_sharpness(float val) { m_shadow_sharpness = val; }
private:
RaysSdfSoa m_rays[2];
RaysSdfSoa m_rays_hit;
uint32_t* m_hit_counter;
uint32_t* m_alive_counter;
uint32_t m_n_rays_initialized = 0;
float m_shadow_sharpness = 2048.f;
bool m_trace_shadow_rays = false;
tcnn::GPUMemoryArena::Allocation m_scratch_alloc;
};
class NerfTracer {
public:
NerfTracer() {}
void init_rays_from_camera(
uint32_t spp,
uint32_t padded_output_width,
uint32_t n_extra_dims,
const Eigen::Vector2i& resolution,
const Eigen::Vector2f& focal_length,
const Eigen::Matrix<float, 3, 4>& camera_matrix0,
const Eigen::Matrix<float, 3, 4>& camera_matrix1,
const Eigen::Vector4f& rolling_shutter,
const Eigen::Vector2f& screen_center,
const Eigen::Vector3f& parallax_shift,
bool snap_to_pixel_centers,
const BoundingBox& render_aabb,
const Eigen::Matrix3f& render_aabb_to_local,
float near_distance,
float plane_z,
float aperture_size,
const Foveation& foveation,
const Lens& lens,
const Buffer2DView<const Eigen::Array4f>& envmap,
const Buffer2DView<const Eigen::Vector2f>& distortion,
Eigen::Array4f* frame_buffer,
float* depth_buffer,
const Buffer2DView<const uint8_t>& hidden_area_mask,
const uint8_t* grid,
int show_accel,
float cone_angle_constant,
ERenderMode render_mode,
cudaStream_t stream
);
uint32_t trace(
NerfNetwork<precision_t>& network,
const BoundingBox& render_aabb,
const Eigen::Matrix3f& render_aabb_to_local,
const BoundingBox& train_aabb,
const Eigen::Vector2f& focal_length,
float cone_angle_constant,
const uint8_t* grid,
ERenderMode render_mode,
const Eigen::Matrix<float, 3, 4> &camera_matrix,
float depth_scale,
int visualized_layer,
int visualized_dim,
ENerfActivation rgb_activation,
ENerfActivation density_activation,
int render_no_attenuation,
int show_accel,
float min_transmittance,
float glow_y_cutoff,
int glow_mode,
const float* extra_dims_gpu,
cudaStream_t stream
);
void enlarge(size_t n_elements, uint32_t padded_output_width, uint32_t n_extra_dims, cudaStream_t stream);
RaysNerfSoa& rays_hit() { return m_rays_hit; }
RaysNerfSoa& rays_init() { return m_rays[0]; }
uint32_t n_rays_initialized() const { return m_n_rays_initialized; }
private:
RaysNerfSoa m_rays[2];
RaysNerfSoa m_rays_hit; precision_t* m_network_output;
float* m_network_input;
uint32_t* m_hit_counter;
uint32_t* m_alive_counter;
uint32_t m_n_rays_initialized = 0;
tcnn::GPUMemoryArena::Allocation m_scratch_alloc;
};
class FiniteDifferenceNormalsApproximator {
public:
void enlarge(uint32_t n_elements, cudaStream_t stream);
void normal(uint32_t n_elements, const distance_fun_t& distance_function, const Eigen::Vector3f* pos, Eigen::Vector3f* normal, float epsilon, cudaStream_t stream);
private:
Eigen::Vector3f* dx;
Eigen::Vector3f* dy;
Eigen::Vector3f* dz;
float* dist_dx_pos;
float* dist_dy_pos;
float* dist_dz_pos;
float* dist_dx_neg;
float* dist_dy_neg;
float* dist_dz_neg;
tcnn::GPUMemoryArena::Allocation m_scratch_alloc;
};
struct LevelStats {
float mean() { return count ? (x / (float)count) : 0.f; }
float variance() { return count ? (xsquared - (x * x) / (float)count) / (float)count : 0.f; }
float sigma() { return sqrtf(variance()); }
float fraczero() { return (float)numzero / float(count + numzero); }
float fracquant() { return (float)numquant / float(count); }
float x;
float xsquared;
float min;
float max;
int numzero;
int numquant;
int count;
};
// Due to mixed-precision training, small loss values can lead to
// underflow (round to zero) in the gradient computations. Hence,
// scale the loss (and thereby gradients) up by this factor and
// divide it out in the optimizer later on.
static constexpr float LOSS_SCALE = 128.0f;
struct NetworkDims {
uint32_t n_input;
uint32_t n_output;
uint32_t n_pos;
};
NetworkDims network_dims_volume() const;
NetworkDims network_dims_sdf() const;
NetworkDims network_dims_image() const;
NetworkDims network_dims_nerf() const;
NetworkDims network_dims() const;
void train_volume(size_t target_batch_size, bool get_loss_scalar, cudaStream_t stream);
void training_prep_volume(uint32_t batch_size, cudaStream_t stream) {}
void load_volume(const fs::path& data_path);
class CudaDevice;
const float* get_inference_extra_dims(cudaStream_t stream) const;
void render_nerf(
cudaStream_t stream,
const CudaRenderBufferView& render_buffer,
NerfNetwork<precision_t>& nerf_network,
const uint8_t* density_grid_bitfield,
const Eigen::Vector2f& focal_length,
const Eigen::Matrix<float, 3, 4>& camera_matrix0,
const Eigen::Matrix<float, 3, 4>& camera_matrix1,
const Eigen::Vector4f& rolling_shutter,
const Eigen::Vector2f& screen_center,
const Foveation& foveation,
int visualized_dimension
);
void render_sdf(
cudaStream_t stream,
const distance_fun_t& distance_function,
const normals_fun_t& normals_function,
const CudaRenderBufferView& render_buffer,
const Eigen::Vector2f& focal_length,
const Eigen::Matrix<float, 3, 4>& camera_matrix,
const Eigen::Vector2f& screen_center,
const Foveation& foveation,
int visualized_dimension
);
void render_image(
cudaStream_t stream,
const CudaRenderBufferView& render_buffer,
const Eigen::Vector2f& focal_length,
const Eigen::Matrix<float, 3, 4>& camera_matrix,
const Eigen::Vector2f& screen_center,
const Foveation& foveation,
int visualized_dimension
);
void render_volume(
cudaStream_t stream,
const CudaRenderBufferView& render_buffer,
const Eigen::Vector2f& focal_length,
const Eigen::Matrix<float, 3, 4>& camera_matrix,
const Eigen::Vector2f& screen_center,
const Foveation& foveation
);
void render_frame(
cudaStream_t stream,
const Eigen::Matrix<float, 3, 4>& camera_matrix0,
const Eigen::Matrix<float, 3, 4>& camera_matrix1,
const Eigen::Matrix<float, 3, 4>& prev_camera_matrix,
const Eigen::Vector2f& screen_center,
const Eigen::Vector2f& relative_focal_length,
const Eigen::Vector4f& nerf_rolling_shutter,
const Foveation& foveation,
const Foveation& prev_foveation,
int visualized_dimension,
CudaRenderBuffer& render_buffer,
bool to_srgb = true,
CudaDevice* device = nullptr
);
void render_frame_main(
CudaDevice& device,
const Eigen::Matrix<float, 3, 4>& camera_matrix0,
const Eigen::Matrix<float, 3, 4>& camera_matrix1,
const Eigen::Vector2f& screen_center,
const Eigen::Vector2f& relative_focal_length,
const Eigen::Vector4f& nerf_rolling_shutter,
const Foveation& foveation,
int visualized_dimension
);
void render_frame_epilogue(
cudaStream_t stream, const Eigen::Matrix<float, 3, 4>& camera_matrix0,
const Eigen::Matrix<float, 3, 4>& prev_camera_matrix,
const Eigen::Vector2f& screen_center,
const Eigen::Vector2f& relative_focal_length,
const Foveation& foveation,
const Foveation& prev_foveation,
CudaRenderBuffer& render_buffer,
bool to_srgb = true
);
void visualize_nerf_cameras(ImDrawList* list, const Eigen::Matrix<float, 4, 4>& world2proj);
fs::path find_network_config(const fs::path& network_config_path);
nlohmann::json load_network_config(const fs::path& network_config_path);
void reload_network_from_file(const fs::path& path = "");
void reload_network_from_json(const nlohmann::json& json, const std::string& config_base_path=""); // config_base_path is needed so that if the passed in json uses the 'parent' feature, we know where to look... be sure to use a filename, or if a directory, end with a trailing slash
void reset_accumulation(bool due_to_camera_movement = false, bool immediate_redraw = true);
void redraw_next_frame() {
m_render_skip_due_to_lack_of_camera_movement_counter = 0;
}
bool reprojection_available() { return m_dlss; }
static ELossType string_to_loss_type(const std::string& str);
void reset_network(bool clear_density_grid = true);
void create_empty_nerf_dataset(size_t n_images, int aabb_scale = 1, bool is_hdr = false);
void load_nerf(const fs::path& data_path);
void load_nerf_post();
void load_mesh(const fs::path& data_path);
void set_exposure(float exposure) { m_exposure = exposure; }
void set_max_level(float maxlevel);
void set_min_level(float minlevel);
void set_visualized_dim(int dim);
void set_visualized_layer(int layer);
void translate_camera(const Eigen::Vector3f& rel, const Eigen::Matrix3f& rot, bool allow_up_down = true);
Eigen::Matrix3f rotation_from_angles(const Eigen::Vector2f& angles) const;
void mouse_drag();
void mouse_wheel();
void load_file(const fs::path& path);
void set_nerf_camera_matrix(const Eigen::Matrix<float, 3, 4>& cam);
Eigen::Vector3f look_at() const;
void set_look_at(const Eigen::Vector3f& pos);
float scale() const { return m_scale; }
void set_scale(float scale);
Eigen::Vector3f view_pos() const { return m_camera.col(3); }
Eigen::Vector3f view_dir() const { return m_camera.col(2); }
Eigen::Vector3f view_up() const { return m_camera.col(1); }
Eigen::Vector3f view_side() const { return m_camera.col(0); }
void set_view_dir(const Eigen::Vector3f& dir);
void first_training_view();
void last_training_view();
void previous_training_view();
void next_training_view();
void set_camera_to_training_view(int trainview);
void reset_camera();
bool keyboard_event();
void generate_training_samples_sdf(Eigen::Vector3f* positions, float* distances, uint32_t n_to_generate, cudaStream_t stream, bool uniform_only);
void update_density_grid_nerf(float decay, uint32_t n_uniform_density_grid_samples, uint32_t n_nonuniform_density_grid_samples, cudaStream_t stream);
void update_density_grid_mean_and_bitfield(cudaStream_t stream);
void mark_density_grid_in_sphere_empty(const Eigen::Vector3f& pos, float radius, cudaStream_t stream);
struct NerfCounters {
tcnn::GPUMemory<uint32_t> numsteps_counter; // number of steps each ray took
tcnn::GPUMemory<uint32_t> numsteps_counter_compacted; // number of steps each ray took
tcnn::GPUMemory<float> loss;
uint32_t rays_per_batch = 1<<12;
uint32_t n_rays_total = 0;
uint32_t measured_batch_size = 0;
uint32_t measured_batch_size_before_compaction = 0;
void prepare_for_training_steps(cudaStream_t stream);
float update_after_training(uint32_t target_batch_size, bool get_loss_scalar, cudaStream_t stream);
};
void train_nerf(uint32_t target_batch_size, bool get_loss_scalar, cudaStream_t stream);
void train_nerf_step(uint32_t target_batch_size, NerfCounters& counters, cudaStream_t stream);
void train_sdf(size_t target_batch_size, bool get_loss_scalar, cudaStream_t stream);
void train_image(size_t target_batch_size, bool get_loss_scalar, cudaStream_t stream);
void set_train(bool mtrain);
template <typename T>
void dump_parameters_as_images(const T* params, const std::string& filename_base);
void prepare_next_camera_path_frame();
void imgui();
void training_prep_nerf(uint32_t batch_size, cudaStream_t stream);
void training_prep_sdf(uint32_t batch_size, cudaStream_t stream);
void training_prep_image(uint32_t batch_size, cudaStream_t stream) {}
void train(uint32_t batch_size);
Eigen::Vector2f calc_focal_length(const Eigen::Vector2i& resolution, const Eigen::Vector2f& relative_focal_length, int fov_axis, float zoom) const;
Eigen::Vector2f render_screen_center(const Eigen::Vector2f& screen_center) const;
void optimise_mesh_step(uint32_t N_STEPS);
void compute_mesh_vertex_colors();
tcnn::GPUMemory<float> get_density_on_grid(Eigen::Vector3i res3d, const BoundingBox& aabb, const Eigen::Matrix3f& render_aabb_to_local); // network version (nerf or sdf)
tcnn::GPUMemory<float> get_sdf_gt_on_grid(Eigen::Vector3i res3d, const BoundingBox& aabb, const Eigen::Matrix3f& render_aabb_to_local); // sdf gt version (sdf only)
tcnn::GPUMemory<Eigen::Array4f> get_rgba_on_grid(Eigen::Vector3i res3d, Eigen::Vector3f ray_dir, bool voxel_centers, float depth, bool density_as_alpha = false);
int marching_cubes(Eigen::Vector3i res3d, const BoundingBox& render_aabb, const Eigen::Matrix3f& render_aabb_to_local, float thresh);
float get_depth_from_renderbuffer(const CudaRenderBuffer& render_buffer, const Eigen::Vector2f& uv);
Eigen::Vector3f get_3d_pos_from_pixel(const CudaRenderBuffer& render_buffer, const Eigen::Vector2i& focus_pixel);
void autofocus();
size_t n_params();
size_t first_encoder_param();
size_t n_encoding_params();
#ifdef NGP_PYTHON
pybind11::dict compute_marching_cubes_mesh(Eigen::Vector3i res3d = Eigen::Vector3i::Constant(128), BoundingBox aabb = BoundingBox{Eigen::Vector3f::Zero(), Eigen::Vector3f::Ones()}, float thresh=2.5f);
pybind11::array_t<float> render_to_cpu(int width, int height, int spp, bool linear, float start_t, float end_t, float fps, float shutter_fraction);
pybind11::array_t<float> screenshot(bool linear) const;
void override_sdf_training_data(pybind11::array_t<float> points, pybind11::array_t<float> distances);
#endif
double calculate_iou(uint32_t n_samples=128*1024*1024, float scale_existing_results_factor=0.0, bool blocking=true, bool force_use_octree = true);
void draw_visualizations(ImDrawList* list, const Eigen::Matrix<float, 3, 4>& camera_matrix);
void train_and_render(bool skip_rendering);
fs::path training_data_path() const;
void init_window(int resw, int resh, bool hidden = false, bool second_window = false);
void destroy_window();
void init_vr();
void update_vr_performance_settings();
void apply_camera_smoothing(float elapsed_ms);
int find_best_training_view(int default_view);
bool begin_frame();
void handle_user_input();
Eigen::Vector3f vr_to_world(const Eigen::Vector3f& pos) const;
void begin_vr_frame_and_handle_vr_input();
void gather_histograms();
void draw_gui();
bool frame();
bool want_repl();
void load_image(const fs::path& data_path);
void load_exr_image(const fs::path& data_path);
void load_stbi_image(const fs::path& data_path);
void load_binary_image(const fs::path& data_path);
uint32_t n_dimensions_to_visualize() const;
float fov() const ;
void set_fov(float val) ;
Eigen::Vector2f fov_xy() const ;
void set_fov_xy(const Eigen::Vector2f& val);
void save_snapshot(const fs::path& path, bool include_optimizer_state, bool compress);
void load_snapshot(const fs::path& path);
CameraKeyframe copy_camera_to_keyframe() const;
void set_camera_from_keyframe(const CameraKeyframe& k); void set_camera_from_time(float t);
void update_loss_graph();
void load_camera_path(const fs::path& path);
bool loop_animation();
void set_loop_animation(bool value);
float compute_image_mse(bool quantize_to_byte);
void compute_and_save_marching_cubes_mesh(const char* filename, Eigen::Vector3i res3d = Eigen::Vector3i::Constant(128), BoundingBox aabb = {}, float thresh = 2.5f, bool unwrap_it = false);
Eigen::Vector3i compute_and_save_png_slices(const char* filename, int res, BoundingBox aabb = {}, float thresh = 2.5f, float density_range = 4.f, bool flip_y_and_z_axes = false);
fs::path root_dir();
////////////////////////////////////////////////////////////////
// marching cubes related state
struct MeshState {
float thresh = 2.5f;
int res = 256;
bool unwrap = false;
float smooth_amount = 2048.f;
float density_amount = 128.f;
float inflate_amount = 1.f;
bool optimize_mesh = false;
tcnn::GPUMemory<Eigen::Vector3f> verts;
tcnn::GPUMemory<Eigen::Vector3f> vert_normals;
tcnn::GPUMemory<Eigen::Vector3f> vert_colors;
tcnn::GPUMemory<Eigen::Vector4f> verts_smoothed; // homogenous
tcnn::GPUMemory<uint32_t> indices;
tcnn::GPUMemory<Eigen::Vector3f> verts_gradient;
std::shared_ptr<TrainableBuffer<3, 1, float>> trainable_verts;
std::shared_ptr<tcnn::Optimizer<float>> verts_optimizer;
void clear() {
indices={};
verts={};
vert_normals={};
vert_colors={};
verts_smoothed={};
verts_gradient={};
trainable_verts=nullptr;
verts_optimizer=nullptr;
}
};
MeshState m_mesh;
bool m_want_repl = false;
bool m_render_window = false;
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;
bool m_train = false;
bool m_training_data_available = false;
bool m_render = true;
int m_max_spp = 0;
ETestbedMode m_testbed_mode = ETestbedMode::None;
bool m_max_level_rand_training = false;
// Rendering stuff
Eigen::Vector2i m_window_res = Eigen::Vector2i::Constant(0);
bool m_dynamic_res = true;
float m_dynamic_res_target_fps = 20.0f;
int m_fixed_res_factor = 8;
float m_scale = 1.0;
float m_aperture_size = 0.0f;
Eigen::Vector2f m_relative_focal_length = Eigen::Vector2f::Ones(); uint32_t m_fov_axis = 1;
float m_zoom = 1.f; // 2d zoom factor (for insets?)
Eigen::Vector2f m_screen_center = Eigen::Vector2f::Constant(0.5f); // center of 2d zoom
float m_ndc_znear = 1.0f / 32.0f;
float m_ndc_zfar = 128.0f;
Eigen::Matrix<float, 3, 4> m_camera = Eigen::Matrix<float, 3, 4>::Zero();
Eigen::Matrix<float, 3, 4> m_smoothed_camera = Eigen::Matrix<float, 3, 4>::Zero();
size_t m_render_skip_due_to_lack_of_camera_movement_counter = 0;
bool m_fps_camera = false;
bool m_camera_smoothing = false;
bool m_autofocus = false;
Eigen::Vector3f m_autofocus_target = Eigen::Vector3f::Constant(0.5f);
CameraPath m_camera_path = {};
Eigen::Vector3f m_up_dir = {0.0f, 1.0f, 0.0f};
Eigen::Vector3f m_sun_dir = Eigen::Vector3f::Ones().normalized();
float m_bounding_radius = 1;
float m_exposure = 0.f;
ERenderMode m_render_mode = ERenderMode::Shade;
EMeshRenderMode m_mesh_render_mode = EMeshRenderMode::VertexNormals;
uint32_t m_seed = 1337;
#ifdef NGP_GUI
GLFWwindow* m_glfw_window = nullptr;
struct SecondWindow {
GLFWwindow* window = nullptr;
GLuint program = 0;
GLuint vao = 0, vbo = 0;
void draw(GLuint texture);
} m_second_window;
float m_drag_depth = 1.0f;
// The VAO will be empty, but we need a valid one for attribute-less rendering
GLuint m_blit_vao = 0;
GLuint m_blit_program = 0;
void init_opengl_shaders();
void blit_texture(const Foveation& foveation, GLint rgba_texture, GLint rgba_filter_mode, GLint depth_texture, GLint framebuffer, const Eigen::Vector2i& offset, const Eigen::Vector2i& resolution);
void create_second_window();
std::unique_ptr<OpenXRHMD> m_hmd;
OpenXRHMD::FrameInfoPtr m_vr_frame_info;
bool m_vr_use_depth_reproject = false;
bool m_vr_use_hidden_area_mask = true;
void set_n_views(size_t n_views);
std::function<bool()> m_keyboard_event_callback;
std::shared_ptr<GLTexture> m_pip_render_texture;
std::vector<std::shared_ptr<GLTexture>> m_rgba_render_textures;
std::vector<std::shared_ptr<GLTexture>> m_depth_render_textures;
#endif
std::unique_ptr<CudaRenderBuffer> m_pip_render_buffer;
SharedQueue<std::unique_ptr<ICallable>> m_task_queue;
void redraw_gui_next_frame() {
m_gui_redraw = true;
}
bool m_gui_redraw = true;
struct Nerf {
struct Training {
NerfDataset dataset;
int n_images_for_training = 0; // how many images to train from, as a high watermark compared to the dataset size
int n_images_for_training_prev = 0; // how many images we saw last time we updated the density grid
struct ErrorMap {
tcnn::GPUMemory<float> data;
tcnn::GPUMemory<float> cdf_x_cond_y;
tcnn::GPUMemory<float> cdf_y;
tcnn::GPUMemory<float> cdf_img;
std::vector<float> pmf_img_cpu;
Eigen::Vector2i resolution = {16, 16};
Eigen::Vector2i cdf_resolution = {16, 16};
bool is_cdf_valid = false;
} error_map;
std::vector<TrainingXForm> transforms;
tcnn::GPUMemory<TrainingXForm> transforms_gpu;
std::vector<Eigen::Vector3f> cam_pos_gradient;
tcnn::GPUMemory<Eigen::Vector3f> cam_pos_gradient_gpu;
std::vector<Eigen::Vector3f> cam_rot_gradient;
tcnn::GPUMemory<Eigen::Vector3f> cam_rot_gradient_gpu;
tcnn::GPUMemory<Eigen::Array3f> cam_exposure_gpu;
std::vector<Eigen::Array3f> cam_exposure_gradient;
tcnn::GPUMemory<Eigen::Array3f> cam_exposure_gradient_gpu;
Eigen::Vector2f cam_focal_length_gradient = Eigen::Vector2f::Zero();
tcnn::GPUMemory<Eigen::Vector2f> cam_focal_length_gradient_gpu;
std::vector<AdamOptimizer<Eigen::Array3f>> cam_exposure;
std::vector<AdamOptimizer<Eigen::Vector3f>> cam_pos_offset;
std::vector<RotationAdamOptimizer> cam_rot_offset;
AdamOptimizer<Eigen::Vector2f> cam_focal_length_offset = AdamOptimizer<Eigen::Vector2f>(0.f);
tcnn::GPUMemory<float> extra_dims_gpu; // if the model demands a latent code per training image, we put them in here.
tcnn::GPUMemory<float> extra_dims_gradient_gpu;
std::vector<AdamOptimizer<Eigen::ArrayXf>> extra_dims_opt;
void reset_extra_dims(default_rng_t &rng);
float extrinsic_l2_reg = 1e-4f;
float extrinsic_learning_rate = 1e-3f;
float intrinsic_l2_reg = 1e-4f;
float exposure_l2_reg = 0.0f;
NerfCounters counters_rgb;
bool random_bg_color = true;
bool grey_loss = false;
bool linear_colors = false;
ELossType loss_type = ELossType::L2;
ELossType depth_loss_type = ELossType::L1;
bool snap_to_pixel_centers = true;
bool train_envmap = false;
bool optimize_distortion = false;
bool optimize_extrinsics = false;
bool optimize_extra_dims = false;
bool optimize_focal_length = false;
bool optimize_exposure = false;
bool render_error_overlay = false;
float error_overlay_brightness = 0.125f; uint32_t n_steps_between_cam_updates = 16;
uint32_t n_steps_since_cam_update = 0;
bool sample_focal_plane_proportional_to_error = false;
bool sample_image_proportional_to_error = false;
bool include_sharpness_in_error = false;
uint32_t n_steps_between_error_map_updates = 128;
uint32_t n_steps_since_error_map_update = 0;
uint32_t n_rays_since_error_map_update = 0;
float near_distance = 0.1f;
float density_grid_decay = 0.95f;
default_rng_t density_grid_rng;
int view = 0;
float depth_supervision_lambda = 0.f;
tcnn::GPUMemory<float> sharpness_grid;
void set_camera_intrinsics(int frame_idx, float fx, float fy = 0.0f, float cx = -0.5f, float cy = -0.5f, float k1 = 0.0f, float k2 = 0.0f, float p1 = 0.0f, float p2 = 0.0f, float k3 = 0.0f, float k4 = 0.0f, bool is_fisheye = false);
void set_camera_extrinsics_rolling_shutter(int frame_idx, Eigen::Matrix<float, 3, 4> camera_to_world_start, Eigen::Matrix<float, 3, 4> camera_to_world_end, const Eigen::Vector4f& rolling_shutter, bool convert_to_ngp = true);
void set_camera_extrinsics(int frame_idx, Eigen::Matrix<float, 3, 4> camera_to_world, bool convert_to_ngp = true);
Eigen::Matrix<float, 3, 4> get_camera_extrinsics(int frame_idx);
void update_transforms(int first = 0, int last = -1);
#ifdef NGP_PYTHON
void set_image(int frame_idx, pybind11::array_t<float> img, pybind11::array_t<float> depth_img, float depth_scale);
#endif
void reset_camera_extrinsics();
void export_camera_extrinsics(const fs::path& path, bool export_extrinsics_in_quat_format = true);
} training = {};
tcnn::GPUMemory<float> density_grid; // NERF_GRIDSIZE()^3 grid of EMA smoothed densities from the network
tcnn::GPUMemory<uint8_t> density_grid_bitfield;
uint8_t* get_density_grid_bitfield_mip(uint32_t mip);
tcnn::GPUMemory<float> density_grid_mean;
uint32_t density_grid_ema_step = 0;
uint32_t max_cascade = 0;
ENerfActivation rgb_activation = ENerfActivation::Exponential;
ENerfActivation density_activation = ENerfActivation::Exponential;
Eigen::Vector3f light_dir = Eigen::Vector3f::Constant(0.5f);
uint32_t extra_dim_idx_for_inference = 0; // which training image's latent code should be presented at inference time
int show_accel = -1;
float sharpen = 0.f;
float cone_angle_constant = 1.f/256.f;
bool visualize_cameras = false;
bool render_with_lens_distortion = false;
Lens render_lens = {};
float render_min_transmittance = 0.01f;
bool render_no_attenuation = false;
bool train_no_attenuation = false;
float glow_y_cutoff = 0.f;
int glow_mode = 0;
} m_nerf;
struct Sdf {
float shadow_sharpness = 2048.0f;
float maximum_distance = 0.00005f;
float fd_normals_epsilon = 0.0005f;
ESDFGroundTruthMode groundtruth_mode = ESDFGroundTruthMode::RaytracedMesh;
BRDFParams brdf;
// Mesh data
EMeshSdfMode mesh_sdf_mode = EMeshSdfMode::Raystab;
float mesh_scale;
tcnn::GPUMemory<Triangle> triangles_gpu;
std::vector<Triangle> triangles_cpu;
std::vector<float> triangle_weights;
DiscreteDistribution triangle_distribution;
tcnn::GPUMemory<float> triangle_cdf;
std::shared_ptr<TriangleBvh> triangle_bvh; // unique_ptr
bool uses_takikawa_encoding = false;
bool use_triangle_octree = false;
int octree_depth_target = 0; // we duplicate this state so that you can waggle the slider without triggering it immediately
std::shared_ptr<TriangleOctree> triangle_octree;
tcnn::GPUMemory<float> brick_data;
uint32_t brick_res = 0;
uint32_t brick_level = 10;
uint32_t brick_quantise_bits = 0;
bool brick_smooth_normals = false; // if true, then we space the central difference taps by one voxel
bool analytic_normals = false;
float zero_offset = 0;
float distance_scale = 0.95f;
double iou = 0.0;
float iou_decay = 0.0f;
bool calculate_iou_online = false;
tcnn::GPUMemory<uint32_t> iou_counter;
struct Training {
size_t idx = 0;
size_t size = 0;
size_t max_size = 1 << 24;
bool did_generate_more_training_data = false;
bool generate_sdf_data_online = true;
float surface_offset_scale = 1.0f;
tcnn::GPUMemory<Eigen::Vector3f> positions;
tcnn::GPUMemory<Eigen::Vector3f> positions_shuffled;
tcnn::GPUMemory<float> distances;
tcnn::GPUMemory<float> distances_shuffled;
tcnn::GPUMemory<Eigen::Vector3f> perturbations;
} training = {};
} m_sdf;
enum EDataType {
Float,
Half,
};
struct Image {
tcnn::GPUMemory<char> data;
EDataType type = EDataType::Float;
Eigen::Vector2i resolution = Eigen::Vector2i::Constant(0.0f);
tcnn::GPUMemory<Eigen::Vector2f> render_coords;
tcnn::GPUMemory<Eigen::Array3f> render_out;
struct Training {
tcnn::GPUMemory<float> positions_tmp;
tcnn::GPUMemory<Eigen::Vector2f> positions;
tcnn::GPUMemory<Eigen::Array3f> targets;
bool snap_to_pixel_centers = true; bool linear_colors = false;
} training = {};
ERandomMode random_mode = ERandomMode::Stratified;
} m_image;
struct VolPayload {
Eigen::Vector3f dir;
Eigen::Array4f col;
uint32_t pixidx;
};
struct Volume {
float albedo = 0.95f;
float scattering = 0.f;
float inv_distance_scale = 100.f;
tcnn::GPUMemory<char> nanovdb_grid;
tcnn::GPUMemory<uint8_t> bitgrid;
float global_majorant = 1.f;
Eigen::Vector3f world2index_offset = {0, 0, 0};
float world2index_scale = 1.f;
struct Training {
tcnn::GPUMemory<Eigen::Vector3f> positions = {};
tcnn::GPUMemory<Eigen::Array4f> targets = {};
} training = {};
// tracing state
tcnn::GPUMemory<Eigen::Vector3f> pos[2] = {};
tcnn::GPUMemory<VolPayload> payload[2] = {};
tcnn::GPUMemory<uint32_t> hit_counter = {};
tcnn::GPUMemory<Eigen::Array4f> radiance_and_density;
} m_volume;
float m_camera_velocity = 1.0f;
EColorSpace m_color_space = EColorSpace::Linear;
ETonemapCurve m_tonemap_curve = ETonemapCurve::Identity;
bool m_dlss = false;
std::shared_ptr<IDlssProvider> m_dlss_provider;
float m_dlss_sharpening = 0.0f;
// 3D stuff
float m_render_near_distance = 0.0f;
float m_slice_plane_z = 0.0f;
bool m_floor_enable = false;
inline float get_floor_y() const { return m_floor_enable ? m_aabb.min.y() + 0.001f : -10000.f; }
BoundingBox m_raw_aabb;
BoundingBox m_aabb;
BoundingBox m_render_aabb;
Eigen::Matrix3f m_render_aabb_to_local;
Eigen::Matrix<float, 3, 4> crop_box(bool nerf_space) const;
std::vector<Eigen::Vector3f> crop_box_corners(bool nerf_space) const;
void set_crop_box(Eigen::Matrix<float, 3, 4> m, bool nerf_space);
// Rendering/UI bookkeeping
Ema m_training_prep_ms = {EEmaType::Time, 100};
Ema m_training_ms = {EEmaType::Time, 100};
Ema m_render_ms = {EEmaType::Time, 100};
// The frame contains everything, i.e. training + rendering + GUI and buffer swapping
Ema m_frame_ms = {EEmaType::Time, 100};
std::chrono::time_point<std::chrono::steady_clock> m_last_frame_time_point;
std::chrono::time_point<std::chrono::steady_clock> m_last_gui_draw_time_point;
std::chrono::time_point<std::chrono::steady_clock> m_training_start_time_point;
Eigen::Array4f m_background_color = {0.0f, 0.0f, 0.0f, 1.0f};
bool m_vsync = false;
bool m_render_transparency_as_checkerboard = false;
// Visualization of neuron activations int m_visualized_dimension = -1;
int m_visualized_layer = 0;
struct View {
std::shared_ptr<CudaRenderBuffer> render_buffer;
Eigen::Vector2i full_resolution = {1, 1};
int visualized_dimension = 0;
Eigen::Matrix<float, 3, 4> camera0 = Eigen::Matrix<float, 3, 4>::Zero();
Eigen::Matrix<float, 3, 4> camera1 = Eigen::Matrix<float, 3, 4>::Zero();
Eigen::Matrix<float, 3, 4> prev_camera = Eigen::Matrix<float, 3, 4>::Zero();
Foveation foveation;
Foveation prev_foveation;
Eigen::Vector2f relative_focal_length;
Eigen::Vector2f screen_center;
CudaDevice* device = nullptr;
};
std::vector<View> m_views;
Eigen::Vector2i m_n_views = {1, 1};
bool m_single_view = true;
float m_picture_in_picture_res = 0.f; // if non zero, requests a small second picture :)
struct ImGuiVars {
static const uint32_t MAX_PATH_LEN = 1024;
bool enabled = true; // tab to toggle
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.ingp";
char video_path[MAX_PATH_LEN] = "video.mp4";
} m_imgui;
fs::path m_root_dir = "";
bool m_visualize_unit_cube = false;
bool m_edit_render_aabb = false;
bool m_edit_world_transform = true;
bool m_snap_to_pixel_centers = false;
Eigen::Vector3f m_parallax_shift = {0.0f, 0.0f, 0.0f}; // to shift the viewer's origin by some amount in camera space
// CUDA stuff
tcnn::StreamAndEvent m_stream;
// Hashgrid encoding analysis
float m_quant_percent = 0.f;
std::vector<LevelStats> m_level_stats;
std::vector<LevelStats> m_first_layer_column_stats;
int m_num_levels = 0;
int m_histo_level = 0; // collect a histogram for this level
uint32_t m_base_grid_resolution;
float m_per_level_scale;
float m_histo[257] = {};
float m_histo_scale = 1.f;
uint32_t m_training_step = 0;
uint32_t m_training_batch_size = 1 << 18;
Ema m_loss_scalar = {EEmaType::Time, 100};
std::vector<float> m_loss_graph = std::vector<float>(256, 0.0f);
size_t m_loss_graph_samples = 0;
bool m_train_encoding = true; bool m_train_network = true;
class CudaDevice {
public:
struct Data {
tcnn::GPUMemory<uint8_t> density_grid_bitfield;
uint8_t* density_grid_bitfield_ptr;
tcnn::GPUMemory<precision_t> params;
std::shared_ptr<Buffer2D<uint8_t>> hidden_area_mask;
};
CudaDevice(int id, bool is_primary) : m_id{id}, m_is_primary{is_primary} {
auto guard = device_guard();
m_stream = std::make_unique<tcnn::StreamAndEvent>();
m_data = std::make_unique<Data>();
m_render_worker = std::make_unique<ThreadPool>(is_primary ? 0u : 1u);
}
CudaDevice(const CudaDevice&) = delete;
CudaDevice& operator=(const CudaDevice&) = delete;
CudaDevice(CudaDevice&&) = default;
CudaDevice& operator=(CudaDevice&&) = default;
tcnn::ScopeGuard device_guard() {
int prev_device = tcnn::cuda_device();
if (prev_device == m_id) {
return {};
}
tcnn::set_cuda_device(m_id);
return tcnn::ScopeGuard{[prev_device]() {
tcnn::set_cuda_device(prev_device);
}};
}
int id() const {
return m_id;
}
bool is_primary() const {
return m_is_primary;
}
std::string name() const {
return tcnn::cuda_device_name(m_id);
}
int compute_capability() const {
return tcnn::cuda_compute_capability(m_id);
}
cudaStream_t stream() const {
return m_stream->get();
}
void wait_for(cudaStream_t stream) const {
CUDA_CHECK_THROW(cudaEventRecord(m_primary_device_event.event, stream));
m_stream->wait_for(m_primary_device_event.event);
}
void signal(cudaStream_t stream) const {
m_stream->signal(stream);
}
const CudaRenderBufferView& render_buffer_view() const {
return m_render_buffer_view;
}
void set_render_buffer_view(const CudaRenderBufferView& view) {
m_render_buffer_view = view;
}
Data& data() const {
return *m_data;
}
bool dirty() const {
return m_dirty;
}
void set_dirty(bool value) {
m_dirty = value;
}
void set_network(const std::shared_ptr<tcnn::Network<float, precision_t>>& network) {
m_network = network;
}
void set_nerf_network(const std::shared_ptr<NerfNetwork<precision_t>>& nerf_network);
const std::shared_ptr<tcnn::Network<float, precision_t>>& network() const {
return m_network;
}
const std::shared_ptr<NerfNetwork<precision_t>>& nerf_network() const {
return m_nerf_network;
}
void clear() {
m_data = std::make_unique<Data>();
m_render_buffer_view = {};
m_network = {};
m_nerf_network = {};
set_dirty(true);
}
template <class F>
auto enqueue_task(F&& f) -> std::future<std::result_of_t <F()>> {
if (is_primary()) {
return std::async(std::launch::deferred, std::forward<F>(f));
} else {
return m_render_worker->enqueue_task(std::forward<F>(f));
}
}
private:
int m_id;
bool m_is_primary;
std::unique_ptr<tcnn::StreamAndEvent> m_stream;
struct Event {
Event() {
CUDA_CHECK_THROW(cudaEventCreate(&event));
}
~Event() {
cudaEventDestroy(event);
}
Event(const Event&) = delete;
Event& operator=(const Event&) = delete;
Event(Event&& other) { *this = std::move(other); }
Event& operator=(Event&& other) {
std::swap(event, other.event);
return *this;
}
cudaEvent_t event = {};
}; Event m_primary_device_event;
std::unique_ptr<Data> m_data;
CudaRenderBufferView m_render_buffer_view = {};
std::shared_ptr<tcnn::Network<float, precision_t>> m_network;
std::shared_ptr<NerfNetwork<precision_t>> m_nerf_network;
bool m_dirty = true;
std::unique_ptr<ThreadPool> m_render_worker;
};
void sync_device(CudaRenderBuffer& render_buffer, CudaDevice& device);
tcnn::ScopeGuard use_device(cudaStream_t stream, CudaRenderBuffer& render_buffer, CudaDevice& device);
void set_all_devices_dirty();
std::vector<CudaDevice> m_devices;
CudaDevice& primary_device() {
return m_devices.front();
}
ThreadPool m_thread_pool;
std::vector<std::future<void>> m_render_futures;
bool m_use_aux_devices = false;
bool m_foveated_rendering = false;
float m_foveated_rendering_max_scaling = 2.0f;
fs::path m_data_path;
fs::path m_network_config_path = "base.json";
nlohmann::json m_network_config;
default_rng_t m_rng;
CudaRenderBuffer m_windowless_render_surface{std::make_shared<CudaSurface2D>()};
uint32_t network_width(uint32_t layer) const;
uint32_t network_num_forward_activations() const;
// Network & training stuff
std::shared_ptr<tcnn::Loss<precision_t>> m_loss;
std::shared_ptr<tcnn::Optimizer<precision_t>> m_optimizer;
std::shared_ptr<tcnn::Encoding<precision_t>> m_encoding;
std::shared_ptr<tcnn::Network<float, precision_t>> m_network;
std::shared_ptr<tcnn::Trainer<float, precision_t, precision_t>> m_trainer;
struct TrainableEnvmap {
std::shared_ptr<tcnn::Optimizer<float>> optimizer;
std::shared_ptr<TrainableBuffer<4, 2, float>> envmap;
std::shared_ptr<tcnn::Trainer<float, float, float>> trainer;
Eigen::Vector2i resolution;
ELossType loss_type;
Buffer2DView<const Eigen::Array4f> inference_view() const {
if (!envmap) {
return {};
}
return {(const Eigen::Array4f*)envmap->inference_params(), resolution};
}
Buffer2DView<const Eigen::Array4f> view() const {
if (!envmap) {
return {};
}
return {(const Eigen::Array4f*)envmap->params(), resolution}; }
} m_envmap;
struct TrainableDistortionMap {
std::shared_ptr<tcnn::Optimizer<float>> optimizer;
std::shared_ptr<TrainableBuffer<2, 2, float>> map;
std::shared_ptr<tcnn::Trainer<float, float, float>> trainer;
Eigen::Vector2i resolution;
Buffer2DView<const Eigen::Vector2f> inference_view() const {
if (!map) {
return {};
}
return {(const Eigen::Vector2f*)map->inference_params(), resolution};
}
Buffer2DView<const Eigen::Vector2f> view() const {
if (!map) {
return {};
}
return {(const Eigen::Vector2f*)map->params(), resolution};
}
} m_distortion;
std::shared_ptr<NerfNetwork<precision_t>> m_nerf_network;
};
NGP_NAMESPACE_END