further corrections

cpp/open3d/t/geometry/RaycastingScene.cpp

@@ -554,7 +554,7 @@ struct RaycastingScene::Impl {
     void ListIntersections(const float* const rays,
                            const size_t num_rays,
                            const size_t num_intersections,
-                           const Eigen::VectorXi cumsum,
+                           const Eigen::VectorXi& cumsum,
                            unsigned int* track_intersections,
                            unsigned int* ray_ids,
                            unsigned int* geometry_ids,
@@ -842,6 +842,7 @@ 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.
@@ -867,12 +868,18 @@ RaycastingScene::ListIntersections(const core::Tensor& rays,
 
     // generate results structure
     std::unordered_map<std::string, core::Tensor> result;
-    shape = {intersections_vector.sum(), 1};
+    result["ray_splits"] = core::Tensor({cumsum.size() + 1}, core::UInt32);
+    uint32_t* ptr = result["ray_splits"].GetDataPtr<uint32_t>();
+    ptr[0] = 0;
+    for (int i = 1; i < cumsum.size() + 1; ++i) {
+        ptr[i] = cumsum[i - 1];
+    }
+    shape = {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.back() = 2;
+    shape.push_back(2);
     result["primitive_uvs"] = core::Tensor(shape, core::Float32);
 
     impl_->ListIntersections(data.GetDataPtr<float>(), num_rays,

cpp/open3d/t/geometry/RaycastingScene.h

@@ -134,6 +134,9 @@ class RaycastingScene {
     /// \param nthreads The number of threads to use. Set to 0 for automatic.
     /// \return The returned dictionary contains:
     ///         - \b ray_ids A tensor with ray IDs. The shape is {..}.
+    ///         - \b ray_splits A tensor with ray intersection splits. Can be
+    ///         used to iterate over all intersections for each ray. The shape
+    ///         is {..}.
     ///         - \b geometry_ids A tensor with the geometry IDs. The shape is
     ///           {..}.
     ///         - \b primitive_ids A tensor with the primitive IDs, which

cpp/pybind/t/geometry/raycasting_scene.cpp

@@ -246,6 +246,9 @@ Lists the intersections of the rays with the scene::
     ray_ids
         A tensor with ray IDs. The shape is {..}.
         
+    ray_splits
+        A tensor with ray intersection splits. Can be used to iterate over all intersections for each ray. The shape is {..}.        
+        
     geometry_ids
         A tensor with the geometry IDs. The shape is {..}.
 
@@ -260,6 +263,16 @@ Lists the intersections of the rays with the scene::
     t_hit
         A tensor with the distance to the hit. The shape is {..}. 
         
+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",

python/test/t/geometry/test_raycasting_scene.py

@@ -137,6 +137,39 @@ def test_count_lots_of_intersections():
     _ = scene.count_intersections(rays)
 
 
+def test_list_intersections():
+    cube = o3d.t.geometry.TriangleMesh.from_legacy(
+        o3d.geometry.TriangleMesh.create_box())
+
+    scene = o3d.t.geometry.RaycastingScene()
+    scene.add_triangles(cube)
+
+    rays = o3d.core.Tensor([[0.5, 0.5, -1, 0, 0, 1], [0.5, 0.5, 0.5, 0, 0, 1],
+                            [10, 10, 10, 1, 0, 0]],
+                           dtype=o3d.core.float32)
+    ans = scene.list_intersections(rays)
+
+    np.testing.assert_allclose(ans['t_hit'].numpy(),
+                               np.array([[1.0], [2.0], [0.5]]),
+                               rtol=1e-6,
+                               atol=1e-6)
+
+
+# list lots of random ray intersections to test the internal batching
+# we expect no errors for this test
+def test_list_lots_of_intersections():
+    cube = o3d.t.geometry.TriangleMesh.from_legacy(
+        o3d.geometry.TriangleMesh.create_box())
+
+    scene = o3d.t.geometry.RaycastingScene()
+    scene.add_triangles(cube)
+
+    rs = np.random.RandomState(123)
+    rays = o3d.core.Tensor.from_numpy(rs.rand(123456, 6).astype(np.float32))
+
+    _ = scene.list_intersections(rays)
+
+
 def test_compute_closest_points():
     vertices = o3d.core.Tensor([[0, 0, 0], [1, 0, 0], [1, 1, 0]],
                                dtype=o3d.core.float32)
@@ -248,7 +281,9 @@ def test_output_shapes(shape):
         'primitive_ids': [],
         'primitive_uvs': [2],
         'primitive_normals': [3],
-        'points': [3]
+        'points': [3],
+        'ray_ids': [],
+        'ray_splits': []
     }
 
     ans = scene.cast_rays(rays)
@@ -267,6 +302,21 @@ def test_output_shapes(shape):
         ) == expected_shape, 'shape mismatch: expected {} but got {} for {}'.format(
             expected_shape, list(v.shape), k)
 
+    ans = scene.list_intersections(rays)
+    nx = np.sum(scene.count_intersections(rays).numpy()).tolist()
+    for k, v in ans.items():
+        alt_shape = np.copy(shape)
+        if k == 'ray_splits':
+            alt_shape[0] = shape[0] + 1
+        else:
+            alt_shape[0] = nx
+        #use np.append otherwise issues if alt_shape = [0] and last_dim[k] = []
+        expected_shape = np.append(alt_shape, last_dim[k]).tolist()
+        assert list(
+            v.shape
+        ) == expected_shape, 'shape mismatch: expected {} but got {} for {}'.format(
+            expected_shape, list(v.shape), k)
+
 
 def test_sphere_wrong_occupancy():
     # This test checks a specific scenario where the old implementation