validate edges

hoho/wed.py

@@ -79,7 +79,7 @@ def compute_WED(pd_vertices, pd_edges, gt_vertices, gt_edges, cv_ins=-1/2, cv_de
     insertion_costs = cv_ins * len(unmatched_gt_indices)  
     
     # Step 5: Edge Deletion and Insertion
-    updated_pd_edges = [(col_ind[np.where(row_ind == edge[0])[0][0]], col_ind[np.where(row_ind == edge[1])[0][0]]) for edge in pd_edges if edge[0] in row_ind and edge[1] in row_ind]
+    updated_pd_edges = [(col_ind[np.where(row_ind == edge[0])[0][0]], col_ind[np.where(row_ind == edge[1])[0][0]]) for edge in pd_edges if len(edge)==2 and edge[0] in row_ind and edge[1] in row_ind]
     pd_edges_set = set(map(tuple, [set(edge) for edge in updated_pd_edges]))
     gt_edges_set = set(map(tuple, [set(edge) for edge in gt_edges]))
 
@@ -88,12 +88,12 @@ def compute_WED(pd_vertices, pd_edges, gt_vertices, gt_edges, cv_ins=-1/2, cv_de
     edges_to_delete = pd_edges_set - gt_edges_set
     
     vert_tf = [np.where(col_ind == v)[0][0] if v in col_ind else 0 for v in range(len(gt_vertices))]
-    deletion_edge_costs = ce * sum(np.linalg.norm(pd_vertices[vert_tf[edge[0]]] - pd_vertices[vert_tf[edge[1]]]) for edge in edges_to_delete)
+    deletion_edge_costs = ce * sum(np.linalg.norm(pd_vertices[vert_tf[edge[0]]] - pd_vertices[vert_tf[edge[1]]]) for edge in edges_to_delete if len(edge) == 2)
 
     
     # Insert missing edges from ground truth
     edges_to_insert = gt_edges_set - pd_edges_set
-    insertion_edge_costs = ce * sum(np.linalg.norm(gt_vertices[edge[0]] - gt_vertices[edge[1]]) for edge in edges_to_insert) 
+    insertion_edge_costs = ce * sum(np.linalg.norm(gt_vertices[edge[0]] - gt_vertices[edge[1]]) for edge in edges_to_insert if len(edge) == 2) 
     
     # Step 5: Calculation of WED
     WED = translation_costs + deletion_costs + insertion_costs + deletion_edge_costs + insertion_edge_costs