import matplotlib.pyplot as plt
import torch
from kaolin.io import obj
from kaolin.ops.mesh import index_vertices_by_faces
from kaolin.metrics.trianglemesh import point_to_mesh_distance

mesh = obj.import_mesh("/path/to/model.obj", with_normals=True, triangulate=True)

resolution = 64
# hardcoded to match the model bounds
x = torch.linspace(-10, 10, resolution)
y = torch.linspace(-10, 10, resolution)
z = torch.linspace(-10, 10, resolution)
grid_points = torch.stack(torch.meshgrid(x, y, z, indexing="ij"), dim=-1).view(-1, 3)

vertices = mesh.vertices.cuda()
faces = mesh.faces.cuda()
grid_points = grid_points.cuda()

face_vertices = index_vertices_by_faces(vertices[None], faces)
distance, _index, _dist_type = point_to_mesh_distance(grid_points[None], face_vertices)

plt.imshow(distance.reshape(resolution, resolution, resolution).cpu()[:, :, 31])
plt.savefig("df-kaolin.png")