From 7c0acac0a50293c52d2adb70967e729f98fa5018 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dmitri=20=C5=A0astin?= <54913410+dmitrishastin@users.noreply.github.com> Date: Fri, 3 Nov 2023 15:41:33 +0000 Subject: [PATCH] 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 | 206 +++++++++++++++++- cpp/open3d/t/geometry/RaycastingScene.h | 27 +++ cpp/pybind/t/geometry/raycasting_scene.cpp | 98 ++++++++- .../test/t/geometry/test_raycasting_scene.py | 51 ++++- 4 files changed, 379 insertions(+), 3 deletions(-) diff --git a/cpp/open3d/t/geometry/RaycastingScene.cpp b/cpp/open3d/t/geometry/RaycastingScene.cpp index 412d6130862..af7dd3f1480 100644 --- a/cpp/open3d/t/geometry/RaycastingScene.cpp +++ b/cpp/open3d/t/geometry/RaycastingScene.cpp @@ -110,6 +110,75 @@ void CountIntersectionsFunc(const RTCFilterFunctionNArguments* args) { } } +struct ListIntersectionsContext { + RTCIntersectContext context; + std::vector>* + 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(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>* + 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 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> + 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& range) { + std::vector 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::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(0, num_rays, BATCH_SIZE), + LoopFn); + }); + } else { + tbb::parallel_for( + tbb::blocked_range(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 +RaycastingScene::ListIntersections(const core::Tensor& rays, + const int nthreads) { + AssertTensorDtypeLastDimDeviceMinNDim(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()); + core::Tensor track_intersections(shape, core::Dtype::FromType()); + auto data = rays.Contiguous(); + impl_->CountIntersections(data.GetDataPtr(), num_rays, + intersections.GetDataPtr(), nthreads); + + // prepare shape with that number of elements + Eigen::Map intersections_vector( + intersections.GetDataPtr(), 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()); + + // generate results structure + std::unordered_map result; + shape.clear(); + shape.push_back(num_rays + 1); + result["ray_splits"] = core::Tensor(shape, core::UInt32); + uint32_t* ptr = result["ray_splits"].GetDataPtr(); + 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(), num_rays, + num_intersections, cumsum, + track_intersections.GetDataPtr(), + result["ray_ids"].GetDataPtr(), + result["geometry_ids"].GetDataPtr(), + result["primitive_ids"].GetDataPtr(), + result["primitive_uvs"].GetDataPtr(), + result["t_hit"].GetDataPtr(), nthreads); + return result; +} + std::unordered_map 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 \ No newline at end of file diff --git a/cpp/open3d/t/geometry/RaycastingScene.h b/cpp/open3d/t/geometry/RaycastingScene.h index b22639148f5..f25d994b0b5 100644 --- a/cpp/open3d/t/geometry/RaycastingScene.h +++ b/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 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 diff --git a/cpp/pybind/t/geometry/raycasting_scene.cpp b/cpp/pybind/t/geometry/raycasting_scene.cpp index bdab66e02e1..e7f6dcecbd5 100644 --- a/cpp/pybind/t/geometry/raycasting_scene.cpp +++ b/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 \ No newline at end of file diff --git a/python/test/t/geometry/test_raycasting_scene.py b/python/test/t/geometry/test_raycasting_scene.py index f89bb0b6644..3ce024a2b29 100644 --- a/python/test/t/geometry/test_raycasting_scene.py +++ b/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,20 @@ 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(): + if k == 'ray_splits': + alt_shape = [np.prod(rays.shape[:-1]) + 1] + else: + alt_shape = [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