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