-
Notifications
You must be signed in to change notification settings - Fork 6
/
nerf_loader.h
124 lines (100 loc) · 4.21 KB
/
nerf_loader.h
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
/*
* 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 nerf_loader.h
* @author Alex Evans, NVIDIA
* @brief Ability to load nerf datasets.
*/
#pragma once
#include <neural-graphics-primitives/bounding_box.cuh>
#include <neural-graphics-primitives/common.h>
#include <filesystem/path.h>
#include <vector>
NGP_NAMESPACE_BEGIN
// how much to scale the scene by vs the original nerf dataset; we want to fit the thing in the unit cube
static constexpr float NERF_SCALE = 0.33f;
struct TrainingImageMetadata {
// Camera intrinsics and additional data associated with a NeRF training image
// the memory to back the pixels and rays is held by GPUMemory objects in the NerfDataset and copied here.
const void* pixels = nullptr;
EImageDataType image_data_type = EImageDataType::Half;
const float* depth = nullptr;
const Ray* rays = nullptr;
CameraDistortion camera_distortion = {};
Eigen::Vector2i resolution = Eigen::Vector2i::Zero();
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();
Eigen::Vector3f light_dir = Eigen::Vector3f::Constant(0.f); // TODO: replace this with more generic float[] of task-specific metadata.
};
inline size_t image_type_size(EImageDataType type) {
switch (type) {
case EImageDataType::None: return 0;
case EImageDataType::Byte: return 1;
case EImageDataType::Half: return 2;
case EImageDataType::Float: return 4;
default: return 0;
}
}
inline size_t depth_type_size(EDepthDataType type) {
switch (type) {
case EDepthDataType::UShort: return 2;
case EDepthDataType::Float: return 4;
default: return 0;
}
}
struct NerfDataset {
std::vector<tcnn::GPUMemory<Ray>> raymemory;
std::vector<tcnn::GPUMemory<uint8_t>> pixelmemory;
std::vector<tcnn::GPUMemory<float>> depthmemory;
std::vector<TrainingImageMetadata> metadata;
std::vector<TrainingXForm> xforms;
tcnn::GPUMemory<float> sharpness_data;
Eigen::Vector2i sharpness_resolution = {0, 0};
tcnn::GPUMemory<float> envmap_data;
BoundingBox render_aabb = {};
Eigen::Vector3f up = {0.0f, 1.0f, 0.0f};
Eigen::Vector3f offset = {0.0f, 0.0f, 0.0f};
size_t n_images = 0;
Eigen::Vector2i envmap_resolution = {0, 0};
float scale = 1.0f;
int aabb_scale = 1;
bool from_mitsuba = false;
bool is_hdr = false;
bool wants_importance_sampling = true;
bool has_rays = false;
uint32_t n_extra_learnable_dims = 0;
bool has_light_dirs = false;
uint32_t n_extra_dims() const {
return (has_light_dirs ? 3u : 0u) + n_extra_learnable_dims;
}
void set_training_image(int frame_idx, const Eigen::Vector2i& image_resolution, const void* pixels, const void* depth_pixels, float depth_scale, bool image_data_on_gpu, EImageDataType image_type, EDepthDataType depth_type, float sharpen_amount = 0.f, bool white_transparent = false, bool black_transparent = false, uint32_t mask_color = 0, const Ray *rays = nullptr);
Eigen::Vector3f nerf_direction_to_ngp(const Eigen::Vector3f& nerf_dir) {
Eigen::Vector3f result = nerf_dir;
return result;
}
Eigen::Matrix<float, 3, 4> nerf_matrix_to_ngp(const Eigen::Matrix<float, 3, 4>& nerf_matrix) {
Eigen::Matrix<float, 3, 4> result = nerf_matrix;
result.col(3) = result.col(3) * scale + offset;
return result;
}
Eigen::Matrix<float, 3, 4> ngp_matrix_to_nerf(const Eigen::Matrix<float, 3, 4>& ngp_matrix) {
Eigen::Matrix<float, 3, 4> result = ngp_matrix;
result.col(3) = (result.col(3) - offset) / scale;
return result;
}
void nerf_ray_to_ngp(Ray& ray, bool scale_direction = false) {
ray.o = ray.o * scale + offset;
if (scale_direction)
ray.d *= scale;
}
};
NerfDataset load_nerf(const std::vector<filesystem::path>& jsonpaths, float sharpen_amount = 0.f);
NerfDataset create_empty_nerf_dataset(size_t n_images, int aabb_scale = 1, bool is_hdr = false);
NGP_NAMESPACE_END