def point_mesh_face_distance(meshes: Meshes, pcls: Pointclouds): """ Computes the distance between a pointcloud and a mesh within a batch. Given a pair `(mesh, pcl)` in the batch, we define the distance to be the sum of two distances, namely `point_face(mesh, pcl) + face_point(mesh, pcl)` `point_face(mesh, pcl)`: Computes the squared distance of each point p in pcl to the closest triangular face in mesh and averages across all points in pcl `face_point(mesh, pcl)`: Computes the squared distance of each triangular face in mesh to the closest point in pcl and averages across all faces in mesh. The above distance functions are applied for all `(mesh, pcl)` pairs in the batch and then averaged across the batch. Args: meshes: A Meshes data structure containing N meshes pcls: A Pointclouds data structure containing N pointclouds Returns: loss: The `point_face(mesh, pcl) + face_point(mesh, pcl)` distance between all `(mesh, pcl)` in a batch averaged across the batch. """ if len(meshes) != len(pcls): raise ValueError("meshes and pointclouds must be equal sized batches") N = len(meshes) # packed representation for pointclouds points = pcls.points_packed() # (P, 3) points_first_idx = pcls.cloud_to_packed_first_idx() max_points = pcls.num_points_per_cloud().max().item() # packed representation for faces verts_packed = meshes.verts_packed() faces_packed = meshes.faces_packed() tris = verts_packed[faces_packed] # (T, 3, 3) tris_first_idx = meshes.mesh_to_faces_packed_first_idx() max_tris = meshes.num_faces_per_mesh().max().item() # point to face distance: shape (P,) point_to_face = point_face_distance(points, points_first_idx, tris, tris_first_idx, max_points) # weight each example by the inverse of number of points in the example point_to_cloud_idx = pcls.packed_to_cloud_idx() # (sum(P_i),) num_points_per_cloud = pcls.num_points_per_cloud() # (N,) weights_p = num_points_per_cloud.gather(0, point_to_cloud_idx) weights_p = 1.0 / weights_p.float() point_to_face = point_to_face * weights_p point_dist = point_to_face.sum() / N # face to point distance: shape (T,) face_to_point = face_point_distance(points, points_first_idx, tris, tris_first_idx, max_tris) # weight each example by the inverse of number of faces in the example tri_to_mesh_idx = meshes.faces_packed_to_mesh_idx() # (sum(T_n),) num_tris_per_mesh = meshes.num_faces_per_mesh() # (N, ) weights_t = num_tris_per_mesh.gather(0, tri_to_mesh_idx) weights_t = 1.0 / weights_t.float() face_to_point = face_to_point * weights_t face_dist = face_to_point.sum() / N return point_dist + face_dist
def point_mesh_edge_distance(meshes: Meshes, pcls: Pointclouds): """ Computes the distance between a pointcloud and a mesh within a batch. Given a pair `(mesh, pcl)` in the batch, we define the distance to be the sum of two distances, namely `point_edge(mesh, pcl) + edge_point(mesh, pcl)` `point_edge(mesh, pcl)`: Computes the squared distance of each point p in pcl to the closest edge segment in mesh and averages across all points in pcl `edge_point(mesh, pcl)`: Computes the squared distance of each edge segment in mesh to the closest point in pcl and averages across all edges in mesh. The above distance functions are applied for all `(mesh, pcl)` pairs in the batch and then averaged across the batch. Args: meshes: A Meshes data structure containing N meshes pcls: A Pointclouds data structure containing N pointclouds Returns: loss: The `point_edge(mesh, pcl) + edge_point(mesh, pcl)` distance between all `(mesh, pcl)` in a batch averaged across the batch. """ if len(meshes) != len(pcls): raise ValueError("meshes and pointclouds must be equal sized batches") N = len(meshes) # packed representation for pointclouds points = pcls.points_packed() # (P, 3) points_first_idx = pcls.cloud_to_packed_first_idx() max_points = pcls.num_points_per_cloud().max().item() # packed representation for edges verts_packed = meshes.verts_packed() edges_packed = meshes.edges_packed() segms = verts_packed[edges_packed] # (S, 2, 3) segms_first_idx = meshes.mesh_to_edges_packed_first_idx() max_segms = meshes.num_edges_per_mesh().max().item() # point to edge distance: shape (P,) point_to_edge = point_edge_distance(points, points_first_idx, segms, segms_first_idx, max_points) # weight each example by the inverse of number of points in the example point_to_cloud_idx = pcls.packed_to_cloud_idx() # (sum(P_i), ) num_points_per_cloud = pcls.num_points_per_cloud() # (N,) weights_p = num_points_per_cloud.gather(0, point_to_cloud_idx) weights_p = 1.0 / weights_p.float() point_to_edge = point_to_edge * weights_p point_dist = point_to_edge.sum() / N # edge to edge distance: shape (S,) edge_to_point = edge_point_distance(points, points_first_idx, segms, segms_first_idx, max_segms) # weight each example by the inverse of number of edges in the example segm_to_mesh_idx = meshes.edges_packed_to_mesh_idx() # (sum(S_n),) num_segms_per_mesh = meshes.num_edges_per_mesh() # (N,) weights_s = num_segms_per_mesh.gather(0, segm_to_mesh_idx) weights_s = 1.0 / weights_s.float() edge_to_point = edge_to_point * weights_s edge_dist = edge_to_point.sum() / N return point_dist + edge_dist