Add fuse_cusfm binary to support cuSFM dataset

docker/run_docker.sh

@@ -2,7 +2,7 @@
 set -ex
 
 
-while getopts ":haid:" OPTION; do
+while getopts ":ha:i:d:" OPTION; do
     case $OPTION in
 
         a)

nvblox/CMakeLists.txt

@@ -194,7 +194,7 @@ add_subdirectory(experiments)
 include(GNUInstallDirs)
 
 install(
-  TARGETS nvblox_lib nvblox_datasets nvblox_eigen fuse_3dmatch fuse_replica
+  TARGETS nvblox_lib nvblox_datasets nvblox_gpu_hash nvblox_stdgpu nvblox_eigen fuse_3dmatch fuse_replica
   EXPORT nvbloxTargets
   ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}/${PROJECT_NAME}
   LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}

nvblox/executables/CMakeLists.txt

@@ -6,9 +6,14 @@ add_nvblox_shared_library(
   src/datasets/image_loader.cpp
   src/datasets/replica.cpp
   src/datasets/redwood.cpp
+  src/datasets/cusfm_data.cpp
   src/fuser.cpp
+  src/conversions/occupancy_conversions.cpp
+  src/conversions/esdf_slice_conversions.cu
   LINK_LIBRARIES_PUBLIC
   nvblox_lib
+  nvblox_stdgpu
+  nvblox_gpu_hash
   INCLUDE_DIRECTORIES_PUBLIC
   $<BUILD_INTERFACE:${CMAKE_CURRENT_LIST_DIR}/include>
   $<INSTALL_INTERFACE:include>)
@@ -44,3 +49,14 @@ add_nvblox_executable(
   nvblox_datasets
   nvblox_stdgpu
   nvblox_gpu_hash)
+
+# CUSFM executable
+add_nvblox_executable(
+  fuse_cusfm
+  SOURCE_FILES
+  src/fuse_cusfm.cpp
+  LINK_LIBRARIES_PUBLIC
+  nvblox_lib
+  nvblox_datasets
+  nvblox_stdgpu
+  nvblox_gpu_hash)

nvblox/executables/include/nvblox/conversions/esdf_slice_conversions.h

@@ -0,0 +1,54 @@
+/*
+Copyright 2024 NVIDIA CORPORATION
+
+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 <nvblox/nvblox.h>
+#include <memory>
+
+namespace nvblox {
+namespace conversions {
+
+class EsdfSliceConverter {
+ public:
+  EsdfSliceConverter();
+  explicit EsdfSliceConverter(std::shared_ptr<CudaStream> cuda_stream);
+
+  // ------------- WRAPPING ESDF SLICER FUNCTIONS -------------
+
+  // Slicing an esdf layer (using EsdfSlicer)
+  void sliceLayerToDistanceImage(
+    const EsdfLayer& layer, float slice_height, float unobserved_value,
+    Image<float>* output_image,
+    AxisAlignedBoundingBox* aabb);
+
+  // Convert slice image to occupancy grid
+  void occupancyGridFromSliceImage(
+    const Image<float>& slice_image, signed char* occupancy_grid_data,
+    float unknown_value);
+
+ private:
+  // Slicer that does the work
+  EsdfSlicer esdf_slicer_;
+
+  std::shared_ptr<CudaStream> cuda_stream_;
+
+  // Buffers
+  device_vector<int8_t> occupancy_grid_device_;
+};
+
+}  // namespace conversions
+}  // namespace nvblox

nvblox/executables/include/nvblox/conversions/occupancy_conversions.h

@@ -0,0 +1,58 @@
+/*
+Copyright 2024 NVIDIA CORPORATION
+
+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 <nvblox/nvblox.h>
+#include <vector>
+#include <string>
+
+namespace nvblox {
+namespace conversions {
+
+constexpr uint8_t kOccupancyGridPngUnknownValue = 205;
+constexpr uint8_t kOccupancyGridPngFreeValue = 254;
+constexpr uint8_t kOccupancyGridPngOccupiedValue = 0;
+
+/// Save the given occupancy grid to png file specified by png_path.
+///
+/// @param png_path           Path to the png file
+/// @param free_thresh        Threshold at which the occupancy grid is considered free
+/// @param occupied_thresh    Threshold at which the occupancy grid is considered occuppied
+/// @param height             Height of the occupancy grid
+/// @param width              Width of the occupancy grid
+/// @param occupancy_grid     The occupancy grid to save as png
+void saveOccupancyGridAsPng(
+  const std::string png_path, const float free_thresh,
+  const float occupied_thresh, const size_t height, const size_t width,
+  const std::vector<signed char>& occupancy_grid);
+
+/// Save the given occupancy grid metadata to yaml_path.
+///
+/// @param yaml_path          Path of yaml file to be saved
+/// @param image_name         Name of the image file this yaml file refers to
+/// @param voxel_size         Size of each voxel
+/// @param origin_x           Origin of the png file (X axis)
+/// @param origin_y           Origin of the png file (Y axis)
+/// @param free_thresh        Threshold at which the occupancy grid is considered free
+/// @param occupied_thresh    Threshold at which the occupancy grid is considered occuppied
+void saveOccupancyGridYaml(
+  const std::string yaml_path, const std::string image_name,
+  const float voxel_size, const float origin_x, const float origin_y,
+  const float free_thresh, const float occupied_thresh);
+
+}  // namespace conversions
+}  // namespace nvblox

nvblox/executables/include/nvblox/datasets/cusfm_data.h

@@ -0,0 +1,153 @@
+/*
+Copyright 2024 NVIDIA CORPORATION
+
+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 <fstream>
+#include <memory>
+#include <string>
+#include <vector>
+#include <unordered_map>
+
+#include "nvblox/core/types.h"
+#include "nvblox/datasets/data_loader.h"
+#include "nvblox/datasets/image_loader.h"
+#include "nvblox/executables/fuser.h"
+
+// Forward declaration for simple JSON parser
+class SimpleJson;
+
+namespace nvblox {
+namespace datasets {
+namespace cusfm_data {
+
+// Build a Fuser for the mapping data
+std::unique_ptr<Fuser> createFuser(
+  const std::string& color_image_dir,
+  const std::string& depth_image_dir,
+  const std::string& frames_meta_file,
+  bool init_from_gflags,
+  bool fit_to_z_plane = false,
+  const std::string& output_dir = "");
+
+struct KeyframeMetadata {
+  uint64_t timestamp_microseconds = 0;
+  std::string color_image_path = "";
+  std::string depth_image_path = "";
+  uint32_t camera_params_id = 0;
+  Transform camera_to_world;
+
+  void fromJson(const SimpleJson& json);
+};
+
+/// @brief A class for loading mapping data comprising of color images,
+/// depth images and a metadata json file that specifies the image paths
+/// under color_image_dir and depth_image_dir, camera poses for each image,
+/// and projection matrix for each camera sensor. The image paths under
+/// color_image_dir and depth_image_dir must match.
+class DataLoader : public RgbdDataLoaderInterface {
+ public:
+  /// Constructors not intended to be called directly, use factor
+  /// DataLoader::create();
+  DataLoader(
+    const std::string& color_image_dir,
+    const std::string& depth_image_dir,
+    const std::vector<KeyframeMetadata>& keyframe_metadatas,
+    const std::unordered_map<uint32_t, Camera>& cameras,
+    bool multithreaded = true,
+    bool fit_to_z_plane = false,
+    const std::string& output_dir = "");
+  virtual ~DataLoader() = default;
+
+  /// Builds a DatasetLoader
+  /// @param color_image_dir Path to the color image folder.
+  /// @param depth_image_dir Path to the depth image folder.
+  /// @param keyframe_metadata_file Path to the frames_meta file.
+  /// @param multithreaded Whether or not to multi-thread image loading
+  /// @param fit_to_z_plane Whether to fit all poses to z=0 plane
+  /// @param output_dir Output directory for saving transform JSON
+  /// @return std::unique_ptr<DataLoader> The dataset loader. May be nullptr if
+  /// construction fails.
+  static std::unique_ptr<DataLoader> create(
+    const std::string& color_image_dir,
+    const std::string& depth_image_dir,
+    const std::string& keyframe_metadata_file,
+    bool multithreaded = true,
+    bool fit_to_z_plane = false,
+    const std::string& output_dir = "");
+
+  /// Interface for a function that loads the next frames in a dataset
+  /// @param[out] depth_frame_ptr The loaded depth frame.
+  /// @param[out] T_L_C_ptr Transform from Camera to the Layer frame.
+  /// @param[out] camera_ptr The intrinsic camera model.
+  /// @param[out] color_frame_ptr Optional, load color frame.
+  /// @return Whether loading succeeded.
+  DataLoadResult loadNext(
+    DepthImage* depth_frame_ptr,                       // NOLINT
+    Transform* T_L_C_ptr,                              // NOLINT
+    Camera* camera_ptr,                                // NOLINT
+    ColorImage* color_frame_ptr = nullptr) override;
+
+  /// Interface for a function that loads the next frames in a dataset.
+  /// This is the version of the function for different depth and color cameras.
+  /// @param[out] depth_frame_ptr The loaded depth frame.
+  /// @param[out] T_L_D_ptr Transform from depth camera to the Layer frame.
+  /// @param[out] depth_camera_ptr The intrinsic depth camera model.
+  /// @param[out] color_frame_ptr The loaded color frame.
+  /// @param[out] T_L_C_ptr Transform from color camera to the Layer frame.
+  /// @param[out] color_camera_ptr The intrinsic color camera model.
+  /// @return Whether loading succeeded.
+  DataLoadResult loadNext(
+    DepthImage* depth_frame_ptr,                               // NOLINT
+    Transform* T_L_D_ptr,                                      // NOLINT
+    Camera* depth_camera_ptr,                                  // NOLINT
+    ColorImage* color_frame_ptr,                               // NOLINT
+    Transform* T_L_C_ptr,                                      // NOLINT
+    Camera* color_camera_ptr) override;
+
+ protected:
+  // Mapping data directory
+  const std::string base_path_;
+
+  // Keyframes metadatas
+  const std::vector<KeyframeMetadata> keyframe_metadatas_;
+
+  // Camera intrinsics
+  const std::unordered_map<uint32_t, Camera> cameras_;
+
+  // The next frame to be loaded
+  int frame_number_ = 0;
+
+  std::unique_ptr<ImageLoader<DepthImage>> depth_image_loader_;
+  std::unique_ptr<ImageLoader<ColorImage>> color_image_loader_;
+
+ private:
+  // Z-plane alignment functionality
+  bool fit_to_z_plane_;
+  std::string output_dir_;
+  Transform T_world_to_z0_plane_;
+  bool has_z_plane_transform_;
+
+  /// Compute transform to align all poses to z=0 plane
+  void computeZPlaneTransform();
+
+  /// Save the z-plane transform to JSON file
+  bool saveTransformToJson() const;
+};
+
+}  // namespace cusfm_data
+}  // namespace datasets
+}  // namespace nvblox