feat: Add function to transform points along a single ray

Solves issue #17
ekut-es · Feb 15, 2024 · 5ebba3a · 5ebba3a
1 parent c0ae0ec
commit 5ebba3a
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/LidarAug/cpp/include/transformations.hpp b/LidarAug/cpp/include/transformations.hpp
@@ -170,4 +170,7 @@ _delete_labels_by_min_points(at::Tensor points, at::Tensor labels,
 }
 
 void random_point_noise(torch::Tensor points, float sigma);
+
+void transform_along_ray(torch::Tensor points, float sigma);
+
 #endif // !TRANSFORMATIONS_HPP
diff --git a/LidarAug/cpp/src/transformations.cpp b/LidarAug/cpp/src/transformations.cpp
@@ -435,6 +435,23 @@ void random_point_noise(torch::Tensor points, float sigma) {
   }
 }
 
+void transform_along_ray(torch::Tensor points, float sigma) {
+  dimensions dims = {points.size(0), points.size(1), points.size(2)};
+
+  std::normal_distribution<float> dist(sigma, 0);
+
+  // TODO(tom): perf measure this
+  for (tensor_size_t i = 0; i < dims.batch_size; i++) {
+    for (tensor_size_t j = 0; j < dims.num_items; j++) {
+      const auto v = points[i][j].data_ptr<float>();
+      const auto noise = std::get<VALUE>(draw_values<float>(dist));
+      auto point_vector = f3{v[0], v[1], v[2]};
+
+      _random_point_noise(point_vector, noise);
+    }
+  }
+}
+
 #ifdef BUILD_MODULE
 #undef TEST_RNG
 #include "../include/bindings.hpp"