Updated chamfer distance function, now returning nn indices as well.

Author: DavidBoja

evaluate_fitting.py

@@ -222,7 +222,7 @@ def evaluate_chamfer(fitting_results_path, scan_path, device, **kwargs):
 
             # compute chamfer
             # dist 1 is 1 x 6890 - chamfer_distance does not return euclidean distance but squared distance
-            dist1, dist2 = chamfer_distance(fit_verts.to(device), scan_vertices.to(device))
+            dist1, dist2, _, _ = chamfer_distance(fit_verts.to(device), scan_vertices.to(device))
             chamfer_standard += (torch.mean(dist1) + torch.mean(dist2)).detach().cpu().item()
             chamfer_bidirectional_average += torch.mean(torch.cat([dist1[0],dist2[0]])).detach().cpu().item()
             chamfer_bm2scan += torch.mean(torch.sqrt(dist1)).detach().cpu().item()
@@ -259,7 +259,7 @@ def evaluate_chamfer(fitting_results_path, scan_path, device, **kwargs):
         fit_verts = torch.from_numpy(fit_verts).unsqueeze(0).float()
 
         # compute chamfer
-        dist1, dist2 = chamfer_distance(fit_verts.to(device), scan_verts.to(device))
+        dist1, dist2, _ , _ = chamfer_distance(fit_verts.to(device), scan_verts.to(device))
         chamfer_standard = (torch.mean(dist1) + torch.mean(dist2)).detach().cpu().item()
         chamfer_bidirectional_average = torch.mean(torch.cat([dist1,dist2])).detach().cpu().item()
         chamfer_bm2scan = torch.mean(torch.sqrt(dist1)).detach().cpu().item()

fit_body_model.py

@@ -124,7 +124,7 @@ def fit_body_model(input_dict: dict, cfg: dict):
                                                    scale)
 
         # compute losses
-        dist1, dist2 = chamfer_distance(body_model_verts.unsqueeze(0), input_vertices)
+        dist1, dist2, _ , _ = chamfer_distance(body_model_verts.unsqueeze(0), input_vertices)
         data_loss = (torch.mean(dist1)) + (torch.mean(dist2))
         data_loss_weighted = data_loss_weight * data_loss
         landmark_loss = summed_L2(body_model_verts[body_model_landmark_inds,:], input_landmarks)

losses.py

@@ -58,7 +58,7 @@ def forward(self,scan_vertices,template_vertices,**kwargs):
                         template to the closest point of the scan closer than
                         partial_data_threshold
         '''
-        _, template2scan_dist = self.chamfer_dist(scan_vertices,template_vertices)
+        _, template2scan_dist, _ , _ = self.chamfer_dist(scan_vertices,template_vertices)
         return template2scan_dist[template2scan_dist < self.partial_data_threshold].sum()