Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

List intersections #6442

Merged
merged 15 commits into from
Nov 3, 2023
Merged
Show file tree
Hide file tree
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
197 changes: 196 additions & 1 deletion cpp/open3d/t/geometry/RaycastingScene.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -691,6 +837,55 @@ 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 = {intersections_vector.sum(), 1};
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;
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) {
Expand Down Expand Up @@ -961,4 +1156,4 @@ uint32_t RaycastingScene::INVALID_ID() { return RTC_INVALID_GEOMETRY_ID; }

} // namespace geometry
} // namespace t
} // namespace open3d
} // namespace open3d
21 changes: 21 additions & 0 deletions cpp/open3d/t/geometry/RaycastingScene.h
Original file line number Diff line number Diff line change
Expand Up @@ -122,6 +122,27 @@ class RaycastingScene {
/// the closest points within the triangles. The shape is {.., 2}.
/// - \b primitive_normals A tensor with the normals of the
/// closest triangle . The shape is {.., 3}.
std::unordered_map<std::string, core::Tensor> ListIntersections(
const core::Tensor &rays, const int nthreads = 0);

/// \brief Lists the intersections of the rays with the scene
/// \param query_points A tensor with >=2 dims, shape {.., 3} and Dtype
/// Float32 describing the query points. {..} can be any number of
/// dimensions, e.g., to organize the query_point to create a 3D grid the
/// shape can be {depth, height, width, 3}. The last dimension must be 3 and
/// has the format [x, y, z].
/// \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 geometry_ids A tensor with the geometry IDs. The shape is
/// {..}.
/// - \b primitive_ids A tensor with the primitive IDs, which
/// corresponds to the triangle index. The shape is {..}.
/// - \b primitive_uvs A tensor with the barycentric coordinates of
/// the intersection points within the triangles. The shape is
/// {.., 2}.
/// - \b t_hit A tensor with the distance to the hit. The shape is
/// {..}.
std::unordered_map<std::string, core::Tensor> ComputeClosestPoints(
const core::Tensor &query_points, const int nthreads = 0);

Expand Down
85 changes: 84 additions & 1 deletion cpp/pybind/t/geometry/raycasting_scene.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -177,6 +177,89 @@ 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 an offset grid of rays.
p_min = np.min(mesh.vertex['positions'].numpy() - 1, axis=0)
p_max = np.max(mesh.vertex['positions'].numpy() - 1, axis=0)
x = np.arange(p_min[0], p_max[0], .1)
y = np.arange(p_min[1], p_max[1], .1)
xv, yv = np.meshgrid(x, y)
orig = np.vstack([xv.flatten(), yv.flatten(), np.tile(p_min[2], xv.size)]).T
dest = np.copy(orig)
dest[:, 2] = p_max[2] + 1
rays = np.hstack([orig, dest - orig]).astype('float32')

# Compute the ray intersections.
lx = scene.list_intersections(rays)

# Calculate intersection coordinates.
v = mesh.vertex['positions'].numpy()
t = mesh.triangle['indices'].numpy()
tidx = lx['primitive_ids'].numpy()
uv = lx['primitive_uvs'].numpy()
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]

# Visualize the intersections.
entry = o3d.geometry.PointCloud(points = o3d.utility.Vector3dVector(orig))
target = o3d.geometry.PointCloud(points = o3d.utility.Vector3dVector(dest))
correspondence = [(i, i) for i in range(rays.shape[0])]
traj = o3d.geometry.LineSet.create_from_point_cloud_correspondences(entry , target , correspondence)
traj.colors = o3d.utility.Vector3dVector(np.tile([1, 0, 0], [rays.shape[0], 1]))
x = o3d.geometry.PointCloud(points = o3d.utility.Vector3dVector(c))
o3d.visualization.draw([mesh, traj, x])


Args:
rays (open3d.core.Tensor): A tensor with >=2 dims, shape {.., 6}, and Dtype
Float32 describing the rays.
{..} can be any number of dimensions, e.g., to organize rays for
creating an image the shape can be {height, width, 6}.
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_ids
A tensor with ray IDs. The shape is {..}.

geometry_ids
A tensor with the geometry IDs. The shape is {..}.

primitive_ids
A tensor with the primitive IDs, which corresponds to the triangle
index. The shape is {..}.

primitive_uvs
A tensor with the barycentric coordinates of the intersection points within
the triangles. The shape is {.., 2}.

t_hit
A tensor with the distance to the hit. The shape is {..}.

)doc");

raycasting_scene.def("compute_closest_points",
Expand Down Expand Up @@ -350,4 +433,4 @@ The value for invalid IDs
}
} // namespace geometry
} // namespace t
} // namespace open3d
} // namespace open3d
Loading