Add Doppler ICP in tensor registration pipeline

CHANGELOG.md

@@ -24,6 +24,7 @@
 * Fix `toString`, `CreateFromPoints` methods and improve docs in `AxisAlignedBoundingBox`. 🐛📝
 * Migrate Open3d documentation to furo theme ✨ (#6470)
 * Expose Near Clip + Far Clip parameters to setup_camera in OffscreenRenderer (#6520)
+* Add Doppler ICP in tensor registration pipeline (PR #5237)
 
 ## 0.13
 

cpp/benchmarks/t/pipelines/registration/Registration.cpp

@@ -10,6 +10,7 @@
 #include <benchmark/benchmark.h>
 
 #include "open3d/core/CUDAUtils.h"
+#include "open3d/core/EigenConverter.h"
 #include "open3d/core/nns/NearestNeighborSearch.h"
 #include "open3d/data/Dataset.h"
 #include "open3d/t/io/PointCloudIO.h"
@@ -31,7 +32,11 @@ static const double voxel_downsampling_factor = 0.02;
 // NNS parameter.
 static const double max_correspondence_distance = 0.05;
 
-// Initial transformation guess for registation.
+// Normal estimation parameters.
+static const double normals_search_radius = 10.0;
+static const int normals_max_neighbors = 30;
+
+// Initial transformation guess for registration.
 static const std::vector<float> initial_transform_flat{
         0.862, 0.011, -0.507, 0.5,  -0.139, 0.967, -0.215, 0.7,
         0.487, 0.255, 0.835,  -1.4, 0.0,    0.0,   0.0,    1.0};
@@ -118,36 +123,155 @@ static void BenchmarkICP(benchmark::State& state,
     }
 }
 
-#define ENUM_ICP_METHOD_DEVICE(METHOD_NAME, TRANSFORMATION_TYPE, DEVICE) \
-    BENCHMARK_CAPTURE(BenchmarkICP, DEVICE METHOD_NAME##_Float32,        \
-                      core::Device(DEVICE), core::Float32,               \
-                      TRANSFORMATION_TYPE)                               \
-            ->Unit(benchmark::kMillisecond);                             \
-    BENCHMARK_CAPTURE(BenchmarkICP, DEVICE METHOD_NAME##_Float64,        \
-                      core::Device(DEVICE), core::Float64,               \
-                      TRANSFORMATION_TYPE)                               \
+core::Tensor ComputeDirectionVectors(const core::Tensor& positions) {
+    core::Tensor directions = core::Tensor::Empty(
+            positions.GetShape(), positions.GetDtype(), positions.GetDevice());
+    for (int64_t i = 0; i < positions.GetLength(); ++i) {
+        // Compute the norm of the position vector.
+        core::Tensor norm = (positions[i][0] * positions[i][0] +
+                             positions[i][1] * positions[i][1] +
+                             positions[i][2] * positions[i][2])
+                                    .Sqrt();
+
+        // If the norm is zero, set the direction vector to zero.
+        if (norm.Item<float>() == 0.0) {
+            directions[i].Fill(0.0);
+        } else {
+            // Otherwise, compute the direction vector by dividing the position
+            // vector by its norm.
+            directions[i] = positions[i] / norm;
+        }
+    }
+    return directions;
+}
+
+static std::tuple<geometry::PointCloud, geometry::PointCloud>
+LoadTensorDopplerPointCloudFromFile(
+        const std::string& source_pointcloud_filename,
+        const std::string& target_pointcloud_filename,
+        const double voxel_downsample_factor,
+        const core::Dtype& dtype,
+        const core::Device& device) {
+    geometry::PointCloud source, target;
+
+    io::ReadPointCloud(source_pointcloud_filename, source,
+                       {"auto", false, false, true});
+    io::ReadPointCloud(target_pointcloud_filename, target,
+                       {"auto", false, false, true});
+
+    source.SetPointAttr("directions",
+                        ComputeDirectionVectors(source.GetPointPositions()));
+
+    source = source.To(device);
+    target = target.To(device);
+
+    // Eliminates the case of impractical values (including negative).
+    if (voxel_downsample_factor > 0.001) {
+        source = source.VoxelDownSample(voxel_downsample_factor);
+        target = target.VoxelDownSample(voxel_downsample_factor);
+    } else {
+        utility::LogWarning(
+                "VoxelDownsample: Impractical voxel size [< 0.001], skipping "
+                "downsampling.");
+    }
+
+    source.SetPointPositions(source.GetPointPositions().To(dtype));
+    source.SetPointAttr("dopplers", source.GetPointAttr("dopplers").To(dtype));
+    source.SetPointAttr("directions",
+                        source.GetPointAttr("directions").To(dtype));
+
+    target.SetPointPositions(target.GetPointPositions().To(dtype));
+    target.EstimateNormals(normals_search_radius, normals_max_neighbors);
+
+    return std::make_tuple(source, target);
+}
+
+static void BenchmarkDopplerICP(benchmark::State& state,
+                                const core::Device& device,
+                                const core::Dtype& dtype,
+                                const TransformationEstimationType& type) {
+    utility::SetVerbosityLevel(utility::VerbosityLevel::Error);
+    data::DemoDopplerICPSequence demo_sequence;
+    geometry::PointCloud source, target;
+    std::tie(source, target) = LoadTensorDopplerPointCloudFromFile(
+            demo_sequence.GetPath(0), demo_sequence.GetPath(1),
+            voxel_downsampling_factor, dtype, device);
+
+    Eigen::Matrix4d calibration{Eigen::Matrix4d::Identity()};
+    double period{0.0};
+    demo_sequence.GetCalibration(calibration, period);
+
+    const core::Tensor calibration_t =
+            core::eigen_converter::EigenMatrixToTensor(calibration)
+                    .To(device, dtype);
+
+    TransformationEstimationForDopplerICP estimation_dicp;
+    estimation_dicp.period_ = period;
+    estimation_dicp.transform_vehicle_to_sensor_ = calibration_t;
+
+    core::Tensor init_trans = core::Tensor::Eye(4, dtype, device);
+    RegistrationResult reg_result(init_trans);
+
+    // Warm up.
+    reg_result = ICP(source, target, max_correspondence_distance, init_trans,
+                     estimation_dicp,
+                     ICPConvergenceCriteria(relative_fitness, relative_rmse,
+                                            max_iterations));
+
+    for (auto _ : state) {
+        reg_result = ICP(source, target, max_correspondence_distance,
+                         init_trans, estimation_dicp,
+                         ICPConvergenceCriteria(relative_fitness, relative_rmse,
+                                                max_iterations));
+        core::cuda::Synchronize(device);
+    }
+}
+
+#define ENUM_ICP_METHOD_DEVICE(BENCHMARK_FUNCTION, METHOD_NAME,         \
+                               TRANSFORMATION_TYPE, DEVICE)             \
+    BENCHMARK_CAPTURE(BENCHMARK_FUNCTION, DEVICE METHOD_NAME##_Float32, \
+                      core::Device(DEVICE), core::Float32,              \
+                      TRANSFORMATION_TYPE)                              \
+            ->Unit(benchmark::kMillisecond);                            \
+    BENCHMARK_CAPTURE(BENCHMARK_FUNCTION, DEVICE METHOD_NAME##_Float64, \
+                      core::Device(DEVICE), core::Float64,              \
+                      TRANSFORMATION_TYPE)                              \
             ->Unit(benchmark::kMillisecond);
 
-ENUM_ICP_METHOD_DEVICE(PointToPoint,
+ENUM_ICP_METHOD_DEVICE(BenchmarkICP,
+                       PointToPoint,
                        TransformationEstimationType::PointToPoint,
                        "CPU:0")
-ENUM_ICP_METHOD_DEVICE(PointToPlane,
+ENUM_ICP_METHOD_DEVICE(BenchmarkICP,
+                       PointToPlane,
                        TransformationEstimationType::PointToPlane,
                        "CPU:0")
-ENUM_ICP_METHOD_DEVICE(ColoredICP,
+ENUM_ICP_METHOD_DEVICE(BenchmarkICP,
+                       ColoredICP,
                        TransformationEstimationType::ColoredICP,
                        "CPU:0")
+ENUM_ICP_METHOD_DEVICE(BenchmarkDopplerICP,
+                       DopplerICP,
+                       TransformationEstimationType::DopplerICP,
+                       "CPU:0")
 
 #ifdef BUILD_CUDA_MODULE
-ENUM_ICP_METHOD_DEVICE(PointToPoint,
+ENUM_ICP_METHOD_DEVICE(BenchmarkICP,
+                       PointToPoint,
                        TransformationEstimationType::PointToPoint,
                        "CUDA:0")
-ENUM_ICP_METHOD_DEVICE(PointToPlane,
+ENUM_ICP_METHOD_DEVICE(BenchmarkICP,
+                       PointToPlane,
                        TransformationEstimationType::PointToPlane,
                        "CUDA:0")
-ENUM_ICP_METHOD_DEVICE(ColoredICP,
+ENUM_ICP_METHOD_DEVICE(BenchmarkICP,
+                       ColoredICP,
                        TransformationEstimationType::ColoredICP,
                        "CUDA:0")
+ENUM_ICP_METHOD_DEVICE(BenchmarkDopplerICP,
+                       DopplerICP,
+                       TransformationEstimationType::DopplerICP,
+                       "CUDA:0")
 #endif
 
 }  // namespace registration

cpp/open3d/core/linalg/kernel/Matrix.h

@@ -36,9 +36,9 @@ static OPEN3D_DEVICE OPEN3D_FORCE_INLINE void matmul3x3_3x1(const scalar_t& m00,
 }
 
 template <typename scalar_t>
-OPEN3D_DEVICE OPEN3D_FORCE_INLINE void matmul3x3_3x1(const scalar_t* A_3x3,
-                                                     const scalar_t* B_3x1,
-                                                     scalar_t* C_3x1) {
+OPEN3D_HOST_DEVICE OPEN3D_FORCE_INLINE void matmul3x3_3x1(const scalar_t* A_3x3,
+                                                          const scalar_t* B_3x1,
+                                                          scalar_t* C_3x1) {
     C_3x1[0] = A_3x3[0] * B_3x1[0] + A_3x3[1] * B_3x1[1] + A_3x3[2] * B_3x1[2];
     C_3x1[1] = A_3x3[3] * B_3x1[0] + A_3x3[4] * B_3x1[1] + A_3x3[5] * B_3x1[2];
     C_3x1[2] = A_3x3[6] * B_3x1[0] + A_3x3[7] * B_3x1[1] + A_3x3[8] * B_3x1[2];

cpp/open3d/data/CMakeLists.txt

@@ -14,6 +14,7 @@ target_sources(data PRIVATE
     dataset/DemoColoredICPPointClouds.cpp
     dataset/DemoCropPointCloud.cpp
     dataset/DemoCustomVisualization.cpp
+    dataset/DemoDopplerICPSequence.cpp
     dataset/DemoFeatureMatchingPointClouds.cpp
     dataset/DemoICPPointClouds.cpp
     dataset/DemoPoseGraphOptimization.cpp

cpp/open3d/data/Dataset.h

@@ -12,6 +12,8 @@
 #include <utility>
 #include <vector>
 
+#include "open3d/utility/Eigen.h"
+
 namespace open3d {
 namespace data {
 
@@ -343,6 +345,39 @@ class DemoCustomVisualization : public DownloadDataset {
     std::string render_option_path_;
 };
 
+/// \class DemoDopplerICPSequence
+/// \brief Data class for `DemoDopplerICPSequence` contains an example sequence
+/// of 100 point clouds with Doppler velocity channel and corresponding ground
+/// truth poses. The sequence was generated using the CARLA simulator.
+class DemoDopplerICPSequence : public DownloadDataset {
+public:
+    DemoDopplerICPSequence(const std::string& data_root = "");
+
+    /// \brief Returns the list of the point cloud paths in the sequence.
+    std::vector<std::string> GetPaths() const { return paths_; }
+    /// \brief Path to the point cloud at index.
+    std::string GetPath(std::size_t index) const;
+    /// \brief Path to the calibration metadata file, containing transformation
+    /// between the vehicle and sensor frames and the time period.
+    std::string GetCalibrationPath() const { return calibration_path_; }
+    /// \brief Path to the ground truth poses for the entire sequence.
+    std::string GetTrajectoryPath() const { return trajectory_path_; }
+    /// \brief Returns the vehicle to sensor calibration transformation and the
+    /// time period (in secs) between sequential point cloud scans.
+    bool GetCalibration(Eigen::Matrix4d& calibration, double& period) const;
+    /// \brief Returns a list of (timestamp, pose) representing the ground truth
+    /// trajectory of the sequence.
+    std::vector<std::pair<double, Eigen::Matrix4d>> GetTrajectory() const;
+
+private:
+    /// List of paths to the point clouds.
+    std::vector<std::string> paths_;
+    /// Path to the calibration JSON file.
+    std::string calibration_path_;
+    /// Path to the TUM ground truth trajectory text file.
+    std::string trajectory_path_;
+};
+
 /// \class DemoFeatureMatchingPointClouds
 /// \brief Data class for `DemoFeatureMatchingPointClouds` contains 2
 /// point cloud fragments and their respective FPFH features and L32D features.

cpp/open3d/data/dataset/DemoDopplerICPSequence.cpp

@@ -0,0 +1,121 @@
+// ----------------------------------------------------------------------------
+// -                        Open3D: www.open3d.org                            -
+// ----------------------------------------------------------------------------
+// Copyright (c) 2018-2023 www.open3d.org
+// SPDX-License-Identifier: MIT
+// ----------------------------------------------------------------------------
+
+#include <json/json.h>
+
+#include <Eigen/Geometry>
+#include <fstream>
+#include <iomanip>
+#include <sstream>
+#include <string>
+#include <vector>
+
+#include "open3d/data/Dataset.h"
+#include "open3d/utility/IJsonConvertible.h"
+#include "open3d/utility/Logging.h"
+
+namespace open3d {
+namespace data {
+
+namespace {
+
+bool ReadJSONFromFile(const std::string& path, Json::Value& json) {
+    std::ifstream file(path);
+    if (!file.is_open()) {
+        utility::LogWarning("Failed to open: {}", path);
+        return false;
+    }
+
+    Json::CharReaderBuilder builder;
+    builder["collectComments"] = false;
+    Json::String errs;
+    bool is_parse_successful = parseFromStream(builder, file, &json, &errs);
+    if (!is_parse_successful) {
+        utility::LogWarning("Read JSON failed: {}.", errs);
+        return false;
+    }
+    return true;
+}
+
+std::vector<std::pair<double, Eigen::Matrix4d>> LoadTUMTrajectory(
+        const std::string& filename) {
+    std::vector<std::pair<double, Eigen::Matrix4d>> trajectory;
+
+    std::ifstream file(filename);
+    if (!file.is_open()) {
+        utility::LogError("Failed to open: {}", filename);
+    }
+
+    std::string line;
+    while (std::getline(file, line)) {
+        std::istringstream iss(line);
+        double timestamp, x, y, z, qx, qy, qz, qw;
+        if (!(iss >> timestamp >> x >> y >> z >> qx >> qy >> qz >> qw)) {
+            utility::LogError("Error parsing line: {}", line);
+        }
+
+        Eigen::Affine3d transform = Eigen::Affine3d::Identity();
+        transform.translate(Eigen::Vector3d(x, y, z));
+        transform.rotate(Eigen::Quaterniond(qw, qx, qy, qz));
+
+        trajectory.emplace_back(timestamp, transform.matrix());
+    }
+
+    return trajectory;
+}
+
+}  // namespace
+
+const static DataDescriptor data_descriptor = {
+        Open3DDownloadsPrefix() +
+                "doppler-icp-data/carla-town05-curved-walls.zip",
+        "73a9828fb7790481168124c02398ee01"};
+
+DemoDopplerICPSequence::DemoDopplerICPSequence(const std::string& data_root)
+    : DownloadDataset("DemoDopplerICPSequence", data_descriptor, data_root) {
+    for (int i = 1; i <= 100; ++i) {
+        std::stringstream ss;
+        ss << std::setw(5) << std::setfill('0') << i;
+        paths_.push_back(GetExtractDir() + "/xyzd_sequence/" + ss.str() +
+                         ".xyzd");
+    }
+
+    calibration_path_ = GetExtractDir() + "/calibration.json";
+    trajectory_path_ = GetExtractDir() + "/ground_truth_poses.txt";
+}
+
+std::string DemoDopplerICPSequence::GetPath(std::size_t index) const {
+    if (index > 99) {
+        utility::LogError(
+                "Invalid index. Expected index between 0 to 99 but got {}.",
+                index);
+    }
+    return paths_[index];
+}
+
+bool DemoDopplerICPSequence::GetCalibration(Eigen::Matrix4d& calibration,
+                                            double& period) const {
+    Json::Value calibration_data;
+    Eigen::Matrix4d calibration_temp;
+    if (ReadJSONFromFile(calibration_path_, calibration_data) &&
+        utility::IJsonConvertible::EigenMatrix4dFromJsonArray(
+                calibration_temp,
+                calibration_data["transform_vehicle_to_sensor"])) {
+        calibration = calibration_temp.transpose();
+        period = calibration_data["period"].asDouble();
+        return true;
+    }
+    return false;
+}
+
+std::vector<std::pair<double, Eigen::Matrix4d>>
+DemoDopplerICPSequence::GetTrajectory() const {
+    return LoadTUMTrajectory(trajectory_path_);
+}
+
+}  // namespace data
+}  // namespace open3d