Implementation of 'find_visited_cells' function

[XJay18]

Hi, thanks for the great project!

Recently, I used your implemented TetrahedraTracer to do some experiments. When I tested TetrahedraTracer.find_visited_cells function, I obtained some unexpected results for the output barycentric_coordinates.

To be specific, here is the testing code:

import torch
from tetranerf.utils.extension import TetrahedraTracer, interpolate_values

device = torch.device("cuda:0")

cuda_vertices = torch.tensor([
    [0.,0.,0.], [1.,0.,0.], [0.,1.,0.], [0.,0.,1.],
    [1.,1.,0.], [0.,2.,0.], [0.,1.,1.],
    [-1.,0.,0.], [-1.,1.,0.], [-1.,0.,1.],
], device=device)

cuda_cells = torch.tensor([
    [0,1,2,3],
    [2,4,5,6],
    [7,0,8,9]
], dtype=torch.int32, device=device)

ray_origins = torch.tensor([[0.2,-4.,0.2]], device=device)
ray_directions = torch.tensor([[0.,1.,0.]], device=device)

tetrahedra_field = torch.tensor([
    [0., 1., 2., 3., 4., 5., 6., 7., 8., 9.],
    [-9.,-8.,-7.,-6.,-5.,-4.,-3.,-2.,-1., 0]
], device=device)

tracer = TetrahedraTracer(device)
tracer.load_tetrahedra(cuda_vertices, cuda_cells)

out = tracer.trace_rays(ray_origins, ray_directions, 10)

# Create 5 samples along the ray given the specified distance
distances = torch.tensor([
    [4., 4.4, 5.2, 5.4, 6.],
], device=device)

traced_cells = tracer.find_visited_cells(
    out["num_visited_cells"],
    out["visited_cells"],
    out["barycentric_coordinates"],
    out["hit_distances"],
    distances,
)

print(traced_cells["vertex_indices"])
""" output
tensor([[[0, 1, 2, 3],
         [0, 1, 2, 3],
         [2, 4, 5, 6],
         [2, 4, 5, 6],
         [0, 0, 0, 0]]], device='cuda:0', dtype=torch.int32)
"""
print(traced_cells["vertex_indices"].shape)

barycentric_coords = traced_cells["barycentric_coordinates"]
print(barycentric_coords)
""" output
tensor([[[0.0000, 0.2000, 0.6000, 0.2000],
         [0.4000, 0.2000, 0.2000, 0.2000],
         [0.2000, 0.2000, 0.4000, 0.2000],
         [0.4000, 0.2000, 0.2000, 0.2000],
         [0.0000, 0.0000, 0.0000, 0.0000]]], device='cuda:0')
"""

traced_vert_indices = traced_cells["vertex_indices"][0]
print(cuda_vertices[traced_vert_indices.long()].shape)
""" output
torch.Size([5, 4, 3])
"""

# Barycentric interpolation to get the points at the specified distance
out_verts = cuda_vertices[traced_vert_indices.long()] * barycentric_coords[0].unsqueeze(-1)
out_verts = out_verts.sum(1)
print(out_verts)

""" output
tensor([[0.2000, 0.6000, 0.2000],
        [0.2000, 0.2000, 0.2000],
        [0.2000, 1.4000, 0.2000],
        [0.2000, 1.2000, 0.2000],
        [0.0000, 0.0000, 0.0000]], device='cuda:0')
"""

However, I expect the out_verts to be:

tensor([[0.2000, 0.0000, 0.2000],
        [0.2000, 0.4000, 0.2000],
        [0.2000, 1.2000, 0.2000],
        [0.2000, 1.4000, 0.2000],
        [0.0000, 0.0000, 0.0000]], device='cuda:0')

It seems that the value of traced_cells["barycentric_coordinates"] is unexpected.

After analyzing the source code, I guess the unexpected results may result from the implementation here: src/tetrahedra_tracer.cu Line 224

I guess the code should be:

barycentric_coordinates_out[i * num_samples_per_ray + j] = make_float4(
    (1 - mult) * coords1.x + mult * coords2.x,
    (1 - mult) * coords1.y + mult * coords2.y,
    (1 - mult) * coords1.z + mult * coords2.z,
    (1 - mult) * coords1.w + mult * coords2.w
);

Could you please check the implementation? Sorry for bothering you and thanks for your time in advance!

[jkulhanek]

I believe you are correct. Sorry for introducing the bug in the first place. I will fix it soon.

[jkulhanek]

Should be addressed by: 91d91e51bea7f7330464d4ed04ff5e0d23186254

[XJay18]

Thanks for your response!