Cleaned-up, and added diameter-based cv cost

hoho/wed.py

@@ -3,12 +3,6 @@ from scipy.optimize import linear_sum_assignment
 import numpy as np 
 
 
-def zeromean_normalize(vertices):
-    vertices = np.array(vertices)
-    vertices = vertices - vertices.mean(axis=0)
-    vertices = vertices / (1e-6 + np.linalg.norm(vertices, axis=1)[:, None]) # project all verts to sphere (not what we meant)
-    return vertices
-
 def preregister_mean_std(verts_to_transform, target_verts, single_scale=True):
     mu_target = target_verts.mean(axis=0)
     mu_in = verts_to_transform.mean(axis=0)
@@ -34,52 +28,38 @@ def preregister_mean_std(verts_to_transform, target_verts, single_scale=True):
     return transformed_verts
 
 
-def compute_WED(pd_vertices, pd_edges, gt_vertices, gt_edges, cv=1000.0, ce=1.0, normalized=True, prenorm=False, preregister=True, register=False, single_scale=True):
+def compute_WED(pd_vertices, pd_edges, gt_vertices, gt_edges, cv=-1, ce=1.0, normalized=True, preregister=True, single_scale=True):
+    '''The function computes the Weighted Edge Distance (WED) between two graphs.
+    pd_vertices: list of predicted vertices
+    pd_edges: list of predicted edges
+    gt_vertices: list of ground truth vertices
+    gt_edges: list of ground truth edges
+    cv: vertex cost
+    ce: edge cost
+    normalized: if True, the WED is normalized by the total length of the ground truth edges
+    preregister: if True, the predicted vertices are pre-registered to the ground truth vertices
+    '''
+    
+    # vertex coordinates are in centimeters, so  cv and ce are set to 100.0 and 1.0 respectively. 
+    # This means the missing a vertex is equivanlent predicting it 1 meters off,
+    # and that is the same as cv and ce equal to 1.0, if GT is in meters
+    
     pd_vertices = np.array(pd_vertices)
     gt_vertices = np.array(gt_vertices)
     
+    diameter = cdist(gt_vertices, gt_vertices).max()
+    
+    if cv < 0:
+        cv = diameter / 4.0 
+        # Cost of addining or deleting a vertex is set to 1/4 of the diameter of the ground truth mesh
+    
     # Step 0: Prenormalize / preregister
-    if prenorm:
-        pd_vertices = zeromean_normalize(pd_vertices)
-        gt_vertices = zeromean_normalize(gt_vertices)
-        
     if preregister:
         pd_vertices = preregister_mean_std(pd_vertices, gt_vertices, single_scale=single_scale)
         
 
     pd_edges = np.array(pd_edges)
-    gt_edges = np.array(gt_edges)
-    
-
-    # Step 0.5: Register
-    if register:
-        # find the optimal rotation, translation, and scale
-        from scipy.spatial.transform import Rotation as R
-        from scipy.optimize import minimize
-        
-        def transform(x, pd_vertices):
-            # x is a 7-element vector, first 3 elements are the rotation vector, next 3 elements are the translation vector, finally scale
-            rotation = R.from_rotvec(x[:3])
-            translation = x[3:6]
-            scale = x[6]
-            return scale * rotation.apply(pd_vertices) + translation
-        
-        def cost_function(x, pd_vertices, gt_vertices):
-            pd_vertices_transformed = transform(x, pd_vertices)
-            distances = cdist(pd_vertices_transformed, gt_vertices, metric='euclidean')
-            row_ind, col_ind = linear_sum_assignment(distances)
-            translation_costs = np.sum(distances[row_ind, col_ind])
-            
-            return translation_costs
-        
-        x0 = np.array([0, 0, 0, 0, 0, 0, 1])
-        # minimize subject to scale > 1e-6
-        # res = minimize(cost_function, x0, args=(pd_vertices, gt_vertices), constraints={'type': 'ineq', 'fun': lambda x: x[6] - 1e-6})
-        res = minimize(cost_function, x0, args=(pd_vertices, gt_vertices), bounds=[(-np.pi, np.pi), (-np.pi, np.pi), (-np.pi, np.pi), (-500, 500), (-500, 500), (-500, 500), (0.1, 3)])
-        # print("scale:", res.x)
-
-        pd_vertices = transform(res.x, pd_vertices) 
-        
+    gt_edges = np.array(gt_edges)        
     
     # Step 1: Bipartite Matching
     distances = cdist(pd_vertices, gt_vertices, metric='euclidean')
@@ -106,7 +86,6 @@ def compute_WED(pd_vertices, pd_edges, gt_vertices, gt_edges, cv=1000.0, ce=1.0,
     # Delete edges not in ground truth
     edges_to_delete = pd_edges_set - gt_edges_set
     
-    #deletion_edge_costs = ce * sum(np.linalg.norm(pd_vertices[edge[0]] - pd_vertices[edge[1]]) for edge in edges_to_delete)
     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)
 
@@ -117,8 +96,7 @@ def compute_WED(pd_vertices, pd_edges, gt_vertices, gt_edges, cv=1000.0, ce=1.0,
     
     # Step 5: Calculation of WED
     WED = translation_costs + deletion_costs + insertion_costs + deletion_edge_costs + insertion_edge_costs
-    # print("translation_costs, deletion_costs, insertion_costs, deletion_edge_costs, insertion_edge_costs")
-    # print(translation_costs, deletion_costs, insertion_costs, deletion_edge_costs, insertion_edge_costs)
+
     
     if normalized:
         total_length_of_gt_edges = np.linalg.norm((gt_vertices[gt_edges[:, 0]] - gt_vertices[gt_edges[:, 1]]), axis=1).sum()