VoxelBlockGrid.h

Go to the documentation of this file.

// ----------------------------------------------------------------------------
// -                        CloudViewer: www.cloudViewer.org                  -
// ----------------------------------------------------------------------------
// Copyright (c) 2018-2024 www.cloudViewer.org
// SPDX-License-Identifier: MIT
// ----------------------------------------------------------------------------
 
#pragma once
 
#include "cloudViewer/core/Tensor.h"
#include "cloudViewer/core/hashmap/HashMap.h"
#include "cloudViewer/t/geometry/Geometry.h"
#include "cloudViewer/t/geometry/Image.h"
#include "cloudViewer/t/geometry/PointCloud.h"
#include "cloudViewer/t/geometry/TensorMap.h"
#include "cloudViewer/t/geometry/TriangleMesh.h"
 
namespace cloudViewer {
namespace t {
namespace geometry {
 
class VoxelBlockGrid {
public:
    VoxelBlockGrid() = default;
 
    VoxelBlockGrid(const std::vector<std::string> &attr_names,
                   const std::vector<core::Dtype> &attr_dtypes,
                   const std::vector<core::SizeVector> &attr_channels,
                   float voxel_size = 0.0058,
                   int64_t block_resolution = 16,
                   int64_t block_count = 10000,
                   const core::Device &device = core::Device("CPU:0"),
                   const core::HashBackendType &backend =
                           core::HashBackendType::Default);
 
    core::HashMap GetHashMap() { return *block_hashmap_; }
 
    core::Tensor GetAttribute(const std::string &attr_name) const;
 
    core::Tensor GetVoxelIndices(const core::Tensor &buf_indices) const;
 
    core::Tensor GetVoxelIndices() const;
 
    core::Tensor GetVoxelCoordinates(const core::Tensor &voxel_indices) const;
 
    std::pair<core::Tensor, core::Tensor>
    GetVoxelCoordinatesAndFlattenedIndices(const core::Tensor &buf_indices);
 
    std::pair<core::Tensor, core::Tensor>
    GetVoxelCoordinatesAndFlattenedIndices();
 
    core::Tensor GetUniqueBlockCoordinates(const Image &depth,
                                           const core::Tensor &intrinsic,
                                           const core::Tensor &extrinsic,
                                           float depth_scale = 1000.0f,
                                           float depth_max = 3.0f,
                                           float trunc_voxel_multiplier = 8.0);
 
    core::Tensor GetUniqueBlockCoordinates(const PointCloud &pcd,
                                           float trunc_voxel_multiplier = 8.0);
 
    void Integrate(const core::Tensor &block_coords,
                   const Image &depth,
                   const Image &color,
                   const core::Tensor &depth_intrinsic,
                   const core::Tensor &color_intrinsic,
                   const core::Tensor &extrinsic,
                   float depth_scale = 1000.0f,
                   float depth_max = 3.0f,
                   float trunc_voxel_multiplier = 8.0f);
 
    void Integrate(const core::Tensor &block_coords,
                   const Image &depth,
                   const Image &color,
                   const core::Tensor &intrinsic,
                   const core::Tensor &extrinsic,
                   float depth_scale = 1000.0f,
                   float depth_max = 3.0f,
                   float trunc_voxel_multiplier = 8.0f);
 
    void Integrate(const core::Tensor &block_coords,
                   const Image &depth,
                   const core::Tensor &intrinsic,
                   const core::Tensor &extrinsic,
                   float depth_scale = 1000.0f,
                   float depth_max = 3.0f,
                   float trunc_voxel_multiplier = 8.0f);
 
    TensorMap RayCast(const core::Tensor &block_coords,
                      const core::Tensor &intrinsic,
                      const core::Tensor &extrinsic,
                      int width,
                      int height,
                      const std::vector<std::string> attrs = {"depth", "color"},
                      float depth_scale = 1000.0f,
                      float depth_min = 0.1f,
                      float depth_max = 3.0f,
                      float weight_threshold = 3.0f,
                      float trunc_voxel_multiplier = 8.0f,
                      int range_map_down_factor = 8);
 
    PointCloud ExtractPointCloud(float weight_threshold = 3.0f,
                                 int estimated_point_number = -1);
 
    TriangleMesh ExtractTriangleMesh(float weight_threshold = 3.0f,
                                     int estimated_vertex_numer = -1);
 
    void Save(const std::string &file_name) const;
 
    static VoxelBlockGrid Load(const std::string &file_name);
 
    VoxelBlockGrid To(const core::Device &device, bool copy = false) const;
 
private:
    void AssertInitialized() const;
 
    VoxelBlockGrid(float voxelSize,
                   int64_t blockResolution,
                   const std::shared_ptr<core::HashMap> &blockHashmap,
                   const std::unordered_map<std::string, int> &nameAttrMap)
        : voxel_size_(voxelSize),
          block_resolution_(blockResolution),
          block_hashmap_(blockHashmap),
          name_attr_map_(nameAttrMap) {}
 
    float voxel_size_ = -1;
    int64_t block_resolution_ = -1;
 
    // Global hash map: 3D coords -> voxel blocks in SoA.
    std::shared_ptr<core::HashMap> block_hashmap_;
 
    // Local hash map: 3D coords -> indices in block_hashmap_.
    std::shared_ptr<core::HashMap> frustum_hashmap_;
 
    // Map: attribute name -> index to access the attribute in SoA.
    std::unordered_map<std::string, int> name_attr_map_;
 
    // Allocated fragment buffer for reuse in depth estimation
    core::Tensor fragment_buffer_;
};
}  // namespace geometry
}  // namespace t
}  // namespace cloudViewer