VulkanPatchMatch.h

/*--------------------------------------------------------------------------------------------
Copyright (c) 2026, NDEVR LLC
tyler.parke@ndevr.org
  __    __   ____     _____ __     __  _______
 |  \  |  | |  __ \  |  ___|\ \   / / |   __  \
 |   \ |  | | |  \ \ | |___  \ \ / /  |  |__)  |
 |  . \|  | | |__/ / | |___   \ V /   |   _   /
 |  |\    |_|_____/__|_____|___\_/____|  | \  \
 |__| \__________________________________|  \__\
 
Subject to the terms of the Enterprise+ Agreement, NDEVR hereby grants
Licensee a limited, non-exclusive, non-transferable, royalty-free license
(without the right to sublicense) to use the API solely for the purpose of
Licensee's internal development efforts to develop applications for which
the API was provided.
 
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
 
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE
FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
DEALINGS IN THE SOFTWARE.
 
Library: VulkanCompute
File: VulkanPatchMatch
Included in API: True
Author(s): Tyler Parke
 *-----------------------------------------------------------------------------------------**/
#pragma once
#include "DLLInfo.h"
#include "VulkanComputeDevice.h"
#include "VulkanComputeBuffer.h"
#include "VulkanComputePipeline.h"
#include <NDEVR/Buffer.h>
#include <NDEVR/Pointer.h>
#include <NDEVR/String.h>
#include <NDEVR/File.h>
 
#if !defined(__EMSCRIPTEN__) && NDEVR_SUPPORTS_THREADING
 
namespace NDEVR
{

    struct PatchMatchCamera
    {
        float K[12];          
        float R[12];          
        float C[3];           
        int32_t camWidth;     
        int32_t camHeight;    
        uint32_t pixelOffset; 
        uint32_t _pad0;       
        uint32_t _pad1;       
    };
    static_assert(sizeof(PatchMatchCamera) == 128, "PatchMatchCamera must be 128 bytes for std430 array stride");

    struct PatchMatchPushConstants
    {
        int32_t width;                
        int32_t height;               
        int32_t nNumViews;            
        int32_t nInitTopK;            
        float fDepthMin;              
        float fDepthMax;              
        float fThresholdKeepCost;     
        int32_t iteration;            
        int32_t isRedPass;            
        int32_t nEstimationIters;     
        int32_t bLowResProcessed;     
        int32_t yOffset;              
        float fConfidenceSkipThreshold;  
        int32_t _pcPad0;              
        int32_t _pcPad1;              
        int32_t _pcPad2;              
    };

    struct DepthEstimateResult
    {
        Buffer<float> depth_map;      
        Buffer<float> normal_map;     
        Buffer<float> confidence_map; 
        int width = 0;                
        int height = 0;               
        void clear()
        {
            width = 0;
            height = 0;
            depth_map.clear();
            normal_map.clear();
            confidence_map.clear();
        }
 
    };

    struct GeometricConsistencyPushConstants
    {
        int32_t width;                  
        int32_t height;                 
        int32_t nNumNeighbors;          
        int32_t nMinConsistent;         
        float fDepthDiffThreshold;      
        float fNormalDiffThreshold;     
        int32_t yOffset;                
        int32_t _pad0;                  
    };

    struct DepthConsistencyPushConstants
    {
        int32_t width;                  
        int32_t height;                 
        int32_t nMinConsistent;         
        int32_t _pad0;                  
        float fDepthRelTolerance;       
        float _pad1;                    
        float _pad2;                    
        float _pad3;                    
        int32_t yOffset;                
        int32_t _pad4;                  
        int32_t _pad5;                  
        int32_t _pad6;                  
    };

    struct MedianFilterPushConstants
    {
        int32_t width;                  
        int32_t height;                 
        int32_t _pad0;
        int32_t _pad1;
        float _pad2;
        float _pad3;
        float _pad4;
        float _pad5;
        int32_t yOffset;
        int32_t _pad6;
        int32_t _pad7;
        int32_t _pad8;
    };

    struct SGMPushConstants
    {
        int32_t width;
        int32_t height;
        int32_t maxRange;        
        int32_t _pad0;
        float _pad1;
        float _pad2;
        float P1;                
        float P2;                
        float decay;             
        float sgmWeight;         
        float threshSmall;       
        int32_t yOffset;
    };

    struct PatchMatchGpuCache
    {
        DynamicPointer<VulkanComputeBuffer> planes_buf, costs_buf, views_buf, rand_buf;
        DynamicPointer<VulkanComputeBuffer> camera_buf, image_buf, low_depth_buf;
        uint04 pixel_count = 0;
        uint04 n_views = 0;
        bool allocated = false;
    };
 
    class VULKAN_COMPUTE_API VulkanPatchMatch
    {
    public:
        static bool IsAvailable(LogPtr log);
 
        VulkanPatchMatch();
        ~VulkanPatchMatch();

        bool isInitialized() const { return m_initialized; }

        bool estimateDepthMap(
            uint04 ref_index,
            const Buffer<uint04>& neighbor_indices,
            const Buffer<PatchMatchCamera>& cameras,
            const Buffer<Buffer<float>>& images,
            const Buffer<int>& image_widths,
            const Buffer<int>& image_heights,
            float depth_min,
            float depth_max,
            uint04 n_estimation_iters,
            uint04 n_init_top_k,
            float threshold_keep_cost,
            const Buffer<float>& initial_planes,  
            DepthEstimateResult& result,
            const Buffer<float>& geometric_prior = Buffer<float>(),
            PatchMatchGpuCache* cache = nullptr,
            float confidence_skip = 0.0f);

        static void FilterDepthMap(DepthEstimateResult& result,
            uint04 min_component_size = 16u,
            float depth_rel_tolerance = 0.02f);

        static bool WriteDmap(File file_path,
            const String& image_file_name,
            const Buffer<uint32_t>& view_ids,
            uint32_t image_width, uint32_t image_height,
            const double K[9], const double R[9], const double C[3],
            float depth_min, float depth_max,
            float threshold_keep_cost,
            const DepthEstimateResult& result);

        bool filterGeometricConsistency(
            const PatchMatchCamera& ref_camera,
            const Buffer<PatchMatchCamera>& neighbor_cameras,
            DepthEstimateResult& ref_result,
            const Buffer<DepthEstimateResult>& neighbor_results,
            float depth_diff_threshold,
            float normal_diff_threshold,
            uint04 min_consistent_views);

        void filterDepthMapGpu(DepthEstimateResult& result,
            uint04 min_consistent = 2u,
            float depth_rel_tolerance = 0.02f);

        bool refineDepthsInPlace(
            const Buffer<PatchMatchCamera>& cameras,
            const Buffer<Buffer<float>>& images,
            const Buffer<int>& image_widths,
            const Buffer<int>& image_heights,
            float depth_min, float depth_max,
            uint04 n_refinement_iters,
            float threshold_keep_cost,
            DepthEstimateResult& result,
            PatchMatchGpuCache* cache = nullptr);

        bool refineGeometric(
            uint04 ref_index,
            const Buffer<uint04>& neighbor_indices,
            const Buffer<PatchMatchCamera>& cameras,
            const Buffer<Buffer<float>>& images,
            const Buffer<int>& image_widths,
            const Buffer<int>& image_heights,
            float depth_min, float depth_max,
            uint04 n_refinement_iters,
            uint04 n_init_top_k,
            float threshold_keep_cost,
            DepthEstimateResult& result,
            const Buffer<PatchMatchCamera>& neighbor_cameras_for_geo,
            const Buffer<DepthEstimateResult>& neighbor_depth_maps,
            float depth_diff_threshold,
            float normal_diff_threshold,
            uint04 min_consistent_views,
            PatchMatchGpuCache* cache = nullptr);
 
    private:
        bool loadShaders();
        void createPipelines();
 
    private:
        ConstPointer<VulkanComputeDevice> m_compute_device; 
 
        // Compute pipelines (one per shader)
        DynamicPointer<VulkanComputePipeline> m_initialize_pipeline;       
        DynamicPointer<VulkanComputePipeline> m_propagate_pipeline;        
        DynamicPointer<VulkanComputePipeline> m_filter_pipeline;           
        DynamicPointer<VulkanComputePipeline> m_refine_pipeline;           
        DynamicPointer<VulkanComputePipeline> m_geo_consistency_pipeline;  
        DynamicPointer<VulkanComputePipeline> m_depth_consistency_pipeline;
        DynamicPointer<VulkanComputePipeline> m_median_pipeline;           
        DynamicPointer<VulkanComputePipeline> m_sgm_pipeline;              
 
        // SPIR-V shaders
        DynamicPointer<VulkanShader> m_initialize_shader;        
        DynamicPointer<VulkanShader> m_propagate_shader;         
        DynamicPointer<VulkanShader> m_filter_shader;            
        DynamicPointer<VulkanShader> m_refine_shader;            
        DynamicPointer<VulkanShader> m_geo_consistency_shader;   
        DynamicPointer<VulkanShader> m_depth_consistency_shader; 
        DynamicPointer<VulkanShader> m_median_shader;            
        DynamicPointer<VulkanShader> m_sgm_shader;               
 
        bool m_initialized = false;  
        LogPtr m_log;                
    public:
        void setLog(LogPtr log) { m_log = log; }
    };
}
 
#else // __EMSCRIPTEN__ or no threading
 
namespace NDEVR
{
    // Stub for platforms without Vulkan support
    class VULKAN_COMPUTE_API VulkanPatchMatch
    {
    public:
        static bool IsAvailable() { return false; }
    };
}
 
#endif // !__EMSCRIPTEN__ && NDEVR_SUPPORTS_THREADING