VoxelBlockGrid.h

Go to the documentation of this file.

// ----------------------------------------------------------------------------
// -                        Open3D: www.open3d.org                            -
// ----------------------------------------------------------------------------
// The MIT License (MIT)
//
// Copyright (c) 2018-2021 www.open3d.org
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// 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.
// ----------------------------------------------------------------------------

#pragma once

#include "open3d/core/Tensor.h"
#include "open3d/core/hashmap/HashMap.h"
#include "open3d/t/geometry/Geometry.h"
#include "open3d/t/geometry/Image.h"
#include "open3d/t/geometry/PointCloud.h"
#include "open3d/t/geometry/TensorMap.h"
#include "open3d/t/geometry/TriangleMesh.h"

namespace open3d {
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 open3d