List intersections (#6442)

Add "ListIntersections" (C++) / "list_intersections" (python) function that returns a dictionary for all ray intersections present in a raycasting scene.

cpp/open3d/t/geometry/RaycastingScene.cpp

@@ -110,6 +110,75 @@ void CountIntersectionsFunc(const RTCFilterFunctionNArguments* args) {
  }
 }
 
+struct ListIntersectionsContext {
+ RTCIntersectContext context;
+ std::vector<std::tuple<uint32_t, uint32_t, float>>*
+ previous_geom_prim_ID_tfar;
+ unsigned int* ray_ids;
+ unsigned int* geometry_ids;
+ unsigned int* primitive_ids;
+ float* primitive_uvs;
+ float* t_hit;
+ Eigen::VectorXi cumsum;
+ unsigned int* track_intersections;
+};
+
+void ListIntersectionsFunc(const RTCFilterFunctionNArguments* args) {
+ int* valid = args->valid;
+ const ListIntersectionsContext* context =
+ reinterpret_cast<const ListIntersectionsContext*>(args->context);
+ struct RTCRayN* rayN = args->ray;
+ struct RTCHitN* hitN = args->hit;
+ const unsigned int N = args->N;
+
+ // Avoid crashing when debug visualizations are used.
+ if (context == nullptr) return;
+
+ std::vector<std::tuple<uint32_t, uint32_t, float>>*
+ previous_geom_prim_ID_tfar = context->previous_geom_prim_ID_tfar;
+ unsigned int* ray_ids = context->ray_ids;
+ unsigned int* geometry_ids = context->geometry_ids;
+ unsigned int* primitive_ids = context->primitive_ids;
+ float* primitive_uvs = context->primitive_uvs;
+ float* t_hit = context->t_hit;
+ Eigen::VectorXi cumsum = context->cumsum;
+ unsigned int* track_intersections = context->track_intersections;
+
+ // Iterate over all rays in ray packet.
+ for (unsigned int ui = 0; ui < N; ui += 1) {
+ // Calculate loop and execution mask
+ unsigned int vi = ui + 0;
+ if (vi >= N) continue;
+
+ // Ignore inactive rays.
+ if (valid[vi] != -1) continue;
+
+ // Read ray/hit from ray structure.
+ RTCRay ray = rtcGetRayFromRayN(rayN, N, ui);
+ RTCHit hit = rtcGetHitFromHitN(hitN, N, ui);
+
+ unsigned int ray_id = ray.id;
+ std::tuple<uint32_t, uint32_t, float> gpID(hit.geomID, hit.primID,
+ ray.tfar);
+ auto& prev_gpIDtfar = previous_geom_prim_ID_tfar->operator[](ray_id);
+ if (std::get<0>(prev_gpIDtfar) != hit.geomID ||
+ (std::get<1>(prev_gpIDtfar) != hit.primID &&
+ std::get<2>(prev_gpIDtfar) != ray.tfar)) {
+ size_t idx = cumsum[ray_id] + track_intersections[ray_id];
+ ray_ids[idx] = ray_id;
+ geometry_ids[idx] = hit.geomID;
+ primitive_ids[idx] = hit.primID;
+ primitive_uvs[idx * 2 + 0] = hit.u;
+ primitive_uvs[idx * 2 + 1] = hit.v;
+ t_hit[idx] = ray.tfar;
+ previous_geom_prim_ID_tfar->operator[](ray_id) = gpID;
+ ++(track_intersections[ray_id]);
+ }
+ // Always ignore hit
+ valid[ui] = 0;
+ }
+}
+
 // Adapted from common/math/closest_point.h
 inline Vec3fa closestPointTriangle(Vec3fa const& p,
  Vec3fa const& a,
@@ -482,6 +551,83 @@ struct RaycastingScene::Impl {
  }
  }
 
+ void ListIntersections(const float* const rays,
+ const size_t num_rays,
+ const size_t num_intersections,
+ const Eigen::VectorXi& cumsum,
+ unsigned int* track_intersections,
+ unsigned int* ray_ids,
+ unsigned int* geometry_ids,
+ unsigned int* primitive_ids,
+ float* primitive_uvs,
+ float* t_hit,
+ const int nthreads) {
+ CommitScene();
+
+ memset(track_intersections, 0, sizeof(uint32_t) * num_rays);
+ memset(ray_ids, 0, sizeof(uint32_t) * num_intersections);
+ memset(geometry_ids, 0, sizeof(uint32_t) * num_intersections);
+ memset(primitive_ids, 0, sizeof(uint32_t) * num_intersections);
+ memset(primitive_uvs, 0, sizeof(float) * num_intersections * 2);
+ memset(t_hit, 0, sizeof(float) * num_intersections);
+
+ std::vector<std::tuple<uint32_t, uint32_t, float>>
+ previous_geom_prim_ID_tfar(
+ num_rays,
+ std::make_tuple(uint32_t(RTC_INVALID_GEOMETRY_ID),
+ uint32_t(RTC_INVALID_GEOMETRY_ID),
+ 0.f));
+
+ ListIntersectionsContext context;
+ rtcInitIntersectContext(&context.context);
+ context.context.filter = ListIntersectionsFunc;
+ context.previous_geom_prim_ID_tfar = &previous_geom_prim_ID_tfar;
+ context.ray_ids = ray_ids;
+ context.geometry_ids = geometry_ids;
+ context.primitive_ids = primitive_ids;
+ context.primitive_uvs = primitive_uvs;
+ context.t_hit = t_hit;
+ context.cumsum = cumsum;
+ context.track_intersections = track_intersections;
+
+ auto LoopFn = [&](const tbb::blocked_range<size_t>& range) {
+ std::vector<RTCRayHit> rayhits(range.size());
+
+ for (size_t i = range.begin(); i < range.end(); ++i) {
+ RTCRayHit* rh = &rayhits[i - range.begin()];
+ const float* r = &rays[i * 6];
+ rh->ray.org_x = r[0];
+ rh->ray.org_y = r[1];
+ rh->ray.org_z = r[2];
+ rh->ray.dir_x = r[3];
+ rh->ray.dir_y = r[4];
+ rh->ray.dir_z = r[5];
+ rh->ray.tnear = 0;
+ rh->ray.tfar = std::numeric_limits<float>::infinity();
+ rh->ray.mask = 0;
+ rh->ray.flags = 0;
+ rh->ray.id = i;
+ rh->hit.geomID = RTC_INVALID_GEOMETRY_ID;
+ rh->hit.instID[0] = RTC_INVALID_GEOMETRY_ID;
+ }
+ rtcIntersect1M(scene_, &context.context, &rayhits[0], range.size(),
+ sizeof(RTCRayHit));
+ };
+
+ if (nthreads > 0) {
+ tbb::task_arena arena(nthreads);
+ arena.execute([&]() {
+ tbb::parallel_for(
+ tbb::blocked_range<size_t>(0, num_rays, BATCH_SIZE),
+ LoopFn);
+ });
+ } else {
+ tbb::parallel_for(
+ tbb::blocked_range<size_t>(0, num_rays, BATCH_SIZE),
+ LoopFn);
+ }
+ }
+
  void ComputeClosestPoints(const float* const query_points,
  const size_t num_query_points,
  float* closest_points,
@@ -691,6 +837,64 @@ core::Tensor RaycastingScene::CountIntersections(const core::Tensor& rays,
  return intersections;
 }
 
+std::unordered_map<std::string, core::Tensor>
+RaycastingScene::ListIntersections(const core::Tensor& rays,
+ const int nthreads) {
+ AssertTensorDtypeLastDimDeviceMinNDim<float>(rays, "rays", 6,
+ impl_->tensor_device_);
+
+ auto shape = rays.GetShape();
+ shape.pop_back(); // Remove last dim, we want to use this shape for the
+ // results.
+ size_t num_rays = shape.NumElements();
+
+ // determine total number of intersections
+ core::Tensor intersections(shape, core::Dtype::FromType<int>());
+ core::Tensor track_intersections(shape, core::Dtype::FromType<uint32_t>());
+ auto data = rays.Contiguous();
+ impl_->CountIntersections(data.GetDataPtr<float>(), num_rays,
+ intersections.GetDataPtr<int>(), nthreads);
+
+ // prepare shape with that number of elements
+ Eigen::Map<Eigen::VectorXi> intersections_vector(
+ intersections.GetDataPtr<int>(), num_rays);
+ size_t num_intersections = intersections_vector.sum();
+
+ // prepare ray allocations (cumsum)
+ Eigen::VectorXi cumsum = Eigen::MatrixXi::Zero(num_rays, 1);
+ std::partial_sum(intersections_vector.begin(),
+ intersections_vector.end() - 1, cumsum.begin() + 1,
+ std::plus<int>());
+
+ // generate results structure
+ std::unordered_map<std::string, core::Tensor> result;
+ shape.clear();
+ shape.push_back(num_rays + 1);
+ result["ray_splits"] = core::Tensor(shape, core::UInt32);
+ uint32_t* ptr = result["ray_splits"].GetDataPtr<uint32_t>();
+ for (int i = 0; i < cumsum.size(); ++i) {
+ ptr[i] = cumsum[i];
+ }
+ ptr[num_rays] = num_intersections;
+ shape[0] = intersections_vector.sum();
+ result["ray_ids"] = core::Tensor(shape, core::UInt32);
+ result["geometry_ids"] = core::Tensor(shape, core::UInt32);
+ result["primitive_ids"] = core::Tensor(shape, core::UInt32);
+ result["t_hit"] = core::Tensor(shape, core::Float32);
+ shape.push_back(2);
+ result["primitive_uvs"] = core::Tensor(shape, core::Float32);
+
+ impl_->ListIntersections(data.GetDataPtr<float>(), num_rays,
+ num_intersections, cumsum,
+ track_intersections.GetDataPtr<uint32_t>(),
+ result["ray_ids"].GetDataPtr<uint32_t>(),
+ result["geometry_ids"].GetDataPtr<uint32_t>(),
+ result["primitive_ids"].GetDataPtr<uint32_t>(),
+ result["primitive_uvs"].GetDataPtr<float>(),
+ result["t_hit"].GetDataPtr<float>(), nthreads);
+ return result;
+}
+
 std::unordered_map<std::string, core::Tensor>
 RaycastingScene::ComputeClosestPoints(const core::Tensor& query_points,
  const int nthreads) {
@@ -961,4 +1165,4 @@ uint32_t RaycastingScene::INVALID_ID() { return RTC_INVALID_GEOMETRY_ID; }
 
 } // namespace geometry
 } // namespace t
-} // namespace open3d
+} // namespace open3d

cpp/open3d/t/geometry/RaycastingScene.h

@@ -104,6 +104,33 @@ class RaycastingScene {
  core::Tensor CountIntersections(const core::Tensor &rays,
  const int nthreads = 0);
 
+ /// \brief Lists the intersections of the rays with the scene
+ /// \param rays A tensor with >=2 dims, shape {.., 6}, and Dtype Float32
+ /// describing the rays; {..} can be any number of dimensions.
+ /// The last dimension must be 6 and has the format [ox, oy, oz, dx, dy, dz]
+ /// with [ox,oy,oz] as the origin and [dx,dy,dz] as the direction. It is not
+ /// necessary to normalize the direction although it should be normalised if
+ /// t_hit is to be calculated in coordinate units.
+ /// \param nthreads The number of threads to use. Set to 0 for automatic.
+ /// \return The returned dictionary contains: ///
+ /// - \b ray_splits A tensor with ray intersection splits. Can be
+ /// used to iterate over all intersections for each ray. The shape
+ /// is {num_rays + 1}.
+ /// - \b ray_ids A tensor with ray IDs. The shape is
+ /// {num_intersections}.
+ /// - \b t_hit A tensor with the distance to the hit. The shape is
+ /// {num_intersections}.
+ /// - \b geometry_ids A tensor with the geometry IDs. The shape is
+ /// {num_intersections}.
+ /// - \b primitive_ids A tensor with the primitive IDs, which
+ /// corresponds to the triangle index. The shape is
+ /// {num_intersections}.
+ /// - \b primitive_uvs A tensor with the barycentric coordinates of
+ /// the intersection points within the triangles. The shape is
+ /// {num_intersections, 2}.
+ std::unordered_map<std::string, core::Tensor> ListIntersections(
+ const core::Tensor &rays, const int nthreads = 0);
+
  /// \brief Computes the closest points on the surfaces of the scene.
  /// \param query_points A tensor with >=2 dims, shape {.., 3} and Dtype
  /// Float32 describing the query points. {..} can be any number of

cpp/pybind/t/geometry/raycasting_scene.cpp

@@ -177,6 +177,102 @@ Computes the number of intersection of the rays with the scene.
 
 Returns:
  A tensor with the number of intersections. The shape is {..}.
+)doc");
+
+ raycasting_scene.def("list_intersections",
+ &RaycastingScene::ListIntersections, "rays"_a,
+ "nthreads"_a = 0, R"doc(
+Lists the intersections of the rays with the scene::
+
+ import open3d as o3d
+ import numpy as np
+
+ # Create scene and add the monkey model.
+ scene = o3d.t.geometry.RaycastingScene()
+ d = o3d.data.MonkeyModel()
+ mesh = o3d.t.io.read_triangle_mesh(d.path)
+ mesh_id = scene.add_triangles(mesh)
+
+ # Create a grid of rays covering the bounding box
+ bb_min = mesh.vertex['positions'].min(dim=0).numpy()
+ bb_max = mesh.vertex['positions'].max(dim=0).numpy()
+ x,y = np.linspace(bb_min, bb_max, num=10)[:,:2].T
+ xv, yv = np.meshgrid(x,y)
+ orig = np.stack([xv, yv, np.full_like(xv, bb_min[2]-1)], axis=-1).reshape(-1,3)
+ dest = orig + np.full(orig.shape, (0,0,2+bb_max[2]-bb_min[2]),dtype=np.float32)
+ rays = np.concatenate([orig, dest-orig], axis=-1).astype(np.float32)
+
+ # Compute the ray intersections.
+ lx = scene.list_intersections(rays)
+ lx = {k:v.numpy() for k,v in lx.items()}
+
+ # Calculate intersection coordinates using the primitive uvs and the mesh
+ v = mesh.vertex['positions'].numpy()
+ t = mesh.triangle['indices'].numpy()
+ tidx = lx['primitive_ids']
+ uv = lx['primitive_uvs']
+ w = 1 - np.sum(uv, axis=1)
+ c = \
+ v[t[tidx, 1].flatten(), :] * uv[:, 0][:, None] + \
+ v[t[tidx, 2].flatten(), :] * uv[:, 1][:, None] + \
+ v[t[tidx, 0].flatten(), :] * w[:, None]
+
+ # Calculate intersection coordinates using ray_ids
+ c = rays[lx['ray_ids']][:,:3] + rays[lx['ray_ids']][:,3:]*lx['t_hit'][...,None]
+ 
+ # Visualize the rays and intersections.
+ lines = o3d.t.geometry.LineSet()
+ lines.point.positions = np.hstack([orig,dest]).reshape(-1,3)
+ lines.line.indices = np.arange(lines.point.positions.shape[0]).reshape(-1,2)
+ lines.line.colors = np.full((lines.line.indices.shape[0],3), (1,0,0))
+ x = o3d.t.geometry.PointCloud(positions=c)
+ o3d.visualization.draw([mesh, lines, x], point_size=8)
+
+
+Args:
+ rays (open3d.core.Tensor): A tensor with >=2 dims, shape {.., 6}, and Dtype
+ Float32 describing the rays; {..} can be any number of dimensions.
+ The last dimension must be 6 and has the format [ox, oy, oz, dx, dy, dz]
+ with [ox,oy,oz] as the origin and [dx,dy,dz] as the direction. It is not
+ necessary to normalize the direction although it should be normalised if 
+ t_hit is to be calculated in coordinate units.
+
+ nthreads (int): The number of threads to use. Set to 0 for automatic.
+
+Returns:
+ The returned dictionary contains
+ 
+ ray_splits
+ A tensor with ray intersection splits. Can be used to iterate over all intersections for each ray. The shape is {num_rays + 1}.
+ 
+ ray_ids
+ A tensor with ray IDs. The shape is {num_intersections}.
+ 
+ t_hit
+ A tensor with the distance to the hit. The shape is {num_intersections}. 
+ 
+ geometry_ids
+ A tensor with the geometry IDs. The shape is {num_intersections}.
+
+ primitive_ids
+ A tensor with the primitive IDs, which corresponds to the triangle
+ index. The shape is {num_intersections}.
+ 
+ primitive_uvs 
+ A tensor with the barycentric coordinates of the intersection points within 
+ the triangles. The shape is {num_intersections, 2}.
+ 
+ 
+An example of using ray_splits::
+
+ ray_splits: [0, 2, 3, 6, 6, 8] # note that the length of this is num_rays+1
+ t_hit: [t1, t2, t3, t4, t5, t6, t7, t8]
+
+ for ray_id, (start, end) in enumerate(zip(ray_splits[:-1], ray_splits[1:])):
+ for i,t in enumerate(t_hit[start:end]):
+ print(f'ray {ray_id}, intersection {i} at {t}')
+ 
+ 
 )doc");
 
  raycasting_scene.def("compute_closest_points",
@@ -350,4 +446,4 @@ The value for invalid IDs
 }
 } // namespace geometry
 } // namespace t
-} // namespace open3d
+} // namespace open3d