Files
Vulkan-Samples/samples/general/mobile_nerf/mobile_nerf.h
T
2025-09-04 10:54:47 +08:00

238 lines
7.0 KiB
C++

/* Copyright (c) 2023-2025, Qualcomm Innovation Center, Inc. All rights reserved.
*
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 the "License";
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#pragma once
#include "api_vulkan_sample.h"
#include <core/acceleration_structure.h>
#include <json.hpp>
using json = nlohmann::json;
namespace vkb
{
namespace sg
{
class Scene;
class Node;
class Mesh;
class SubMesh;
class Camera;
} // namespace sg
} // namespace vkb
class MobileNerf : public ApiVulkanSample
{
public:
MobileNerf();
~MobileNerf() override;
void request_gpu_features(vkb::PhysicalDevice &gpu) override;
void render(float delta_time) override;
bool prepare(const vkb::ApplicationOptions &options) override;
bool resize(const uint32_t width, const uint32_t height) override;
private:
struct GlobalUniform
{
glm::mat4x4 model;
glm::mat4x4 view;
glm::mat4x4 proj;
alignas(16) glm::vec3 camera_position;
alignas(16) glm::vec3 camera_side;
alignas(16) glm::vec3 camera_up;
alignas(16) glm::vec3 camera_lookat;
alignas(8) glm::vec2 img_dim;
alignas(4) float tan_half_fov;
} global_uniform;
struct PushConstants
{
unsigned int weight_idx;
} push_constants;
#define WEIGHTS_0_COUNT (176)
#define WEIGHTS_1_COUNT (256)
// The third layer weights' size is changed from 48 to 64 to make sure a 16 bytes alignement
#define WEIGHTS_2_COUNT (64)
#define BIAS_0_COUNT (16)
#define BIAS_1_COUNT (16)
// The third layer bias' size is changed from 3 to 4 to make sure a 16 bytes alignement
#define BIAS_2_COUNT (4)
struct MLP_Weights
{
float data[WEIGHTS_0_COUNT + WEIGHTS_1_COUNT + WEIGHTS_2_COUNT +
BIAS_0_COUNT + BIAS_1_COUNT + BIAS_2_COUNT]; // Array of floats
};
struct Vertex
{
glm::vec3 position;
glm::vec2 tex_coord;
};
struct InstancingInfo
{
glm::ivec3 dim;
glm::vec3 interval;
};
struct InstanceData
{
glm::vec3 pos_offset;
};
struct FrameBufferAttachment
{
using ImagePtr = std::unique_ptr<vkb::core::Image>;
ImagePtr image;
VkSampler sampler;
VkImageView view;
operator bool() const
{
return image && image->get_handle() != VK_NULL_HANDLE;
}
void destroy();
};
struct Model
{
std::vector<Vertex> vertices;
std::vector<std::array<uint32_t, 3>> indices;
Texture texture_input_0, texture_input_1;
// Vulkan Buffers for each model
std::unique_ptr<vkb::core::BufferC> vertex_buffer{nullptr};
std::unique_ptr<vkb::core::BufferC> index_buffer{nullptr};
// Each model will have its own pipeline and descriptor set
VkPipeline pipeline_first_pass{VK_NULL_HANDLE};
// We make the descriptor set a vector for the forward mode
// Deferred mode will only have one set of descriptor per model
std::vector<VkDescriptorSet> descriptor_set_first_pass{VK_NULL_HANDLE};
int sub_model_num;
int model_index;
};
std::vector<Model> models;
// MLPs for each model
std::vector<MLP_Weights> mlp_weight_vector;
std::vector<std::unique_ptr<vkb::core::BufferC>> weights_buffers;
// Uniform buffer for each model
std::vector<std::unique_ptr<vkb::core::BufferC>> uniform_buffers;
// Buffer to store instance data
std::unique_ptr<vkb::core::BufferC> instance_buffer{nullptr};
// Common
void read_json_map();
void load_shaders();
void build_command_buffers() override;
void create_uniforms();
void create_static_object_buffers(int model_index, int sub_model_index, int models_entry);
void prepare_instance_data();
void load_scene(int model_index, int sub_model_index, int models_entry);
void initialize_mlp_uniform_buffers(int model_index);
void update_uniform_buffers();
void update_weights_buffers();
void create_texture(int model_index, int sub_model_index, int models_entry);
void create_texture_helper(std::string const &texturePath, Texture &texture);
VkFormat feature_map_format = VK_FORMAT_R16G16B16A16_SFLOAT;
// Helper function
void setup_attachment(VkFormat format, VkImageUsageFlags usage, FrameBufferAttachment &attachment);
void draw();
// First pass
VkDescriptorSetLayout descriptor_set_first_pass_layout{VK_NULL_HANDLE};
VkPipelineLayout pipeline_first_pass_layout{VK_NULL_HANDLE};
std::array<VkPipelineShaderStageCreateInfo, 2> shader_stages_first_pass;
std::array<VkPipelineShaderStageCreateInfo, 2> shader_stages_second_pass;
void create_descriptor_pool();
void create_pipeline_layout_fist_pass();
void create_descriptor_sets_first_pass(Model &model);
void prepare_pipelines();
unsigned int color_attach_0_idx;
unsigned int color_attach_1_idx;
unsigned int color_attach_2_idx;
unsigned int color_attach_3_idx;
unsigned int depth_attach_idx;
unsigned int swapchain_attach_idx;
struct Attachments_baseline
{
FrameBufferAttachment feature_0, feature_1, feature_2, weights_idx;
};
std::vector<Attachments_baseline> frameAttachments;
std::vector<VkFramebuffer> nerf_framebuffers;
VkRenderPass render_pass_nerf{VK_NULL_HANDLE};
// For the baseline (mlp in frag shader) second pass
VkPipeline pipeline_baseline{VK_NULL_HANDLE};
VkPipelineLayout pipeline_layout_baseline{VK_NULL_HANDLE};
VkDescriptorSetLayout descriptor_set_layout_baseline{VK_NULL_HANDLE};
std::vector<VkDescriptorSet> descriptor_set_baseline{VK_NULL_HANDLE};
void update_render_pass_nerf_baseline();
void create_pipeline_layout_baseline();
void create_descriptor_sets_baseline();
void setup_nerf_framebuffer_baseline();
void update_descriptor_sets_baseline();
void build_command_buffers_baseline();
// For creating the forward mode rendepass
void update_render_pass_nerf_forward();
// For loading nerf assets map
json asset_map;
std::vector<std::string> model_path;
std::vector<bool> using_original_nerf_models;
bool combo_mode = false;
bool use_deferred = false;
bool do_rotation = false;
glm::vec3 camera_pos = glm::vec3(-2.2f, 2.2f, 2.2f);
// Currently combo mode translation are hard-coded
glm::mat4x4 combo_model_transform[4] = {
glm::translate(glm::vec3(0.5, 0.75, 0)), glm::translate(glm::vec3(0.5, 0.25, 0)),
glm::translate(glm::vec3(0, -0.25, 0.5)), glm::translate(glm::vec3(0, -0.75, -0.5))};
// For instancing
InstancingInfo instancing_info;
// Viewport Setting
float fov = 60.0f;
uint32_t view_port_width = width;
uint32_t view_port_height = height;
bool use_native_screen_size = false;
};
std::unique_ptr<vkb::VulkanSampleC> create_mobile_nerf();