From 5ebba3afa8494c73e97a2193d1b470d714a1aa7b Mon Sep 17 00:00:00 2001
From: Tom Schammo <tom@tomschammo.com>
Date: Thu, 15 Feb 2024 18:12:13 +0100
Subject: [PATCH] feat: Add function to transform points along a single ray

Solves issue #17
---
 LidarAug/cpp/include/transformations.hpp |  3 +++
 LidarAug/cpp/src/transformations.cpp     | 17 +++++++++++++++++
 2 files changed, 20 insertions(+)

diff --git a/LidarAug/cpp/include/transformations.hpp b/LidarAug/cpp/include/transformations.hpp
index 8b35e15..5cb2d79 100644
--- 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
index f3e0864..2cdf6a2 100644
--- 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"