Merge branch 'pr/2' into pr/4

hoho/wed.py

@@ -28,40 +28,43 @@ 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=-1, ce=1.0, normalized=True, preregister=True, single_scale=True):
+def update_cv(cv, gt_vertices):
+    if cv < 0:
+        diameter = cdist(gt_vertices, gt_vertices).max()
+        # Cost of adding or deleting a vertex is set to -cv times the diameter of the ground truth wireframe
+        cv = -cv * diameter
+    elif cv == 0:
+        # Cost of adding or deleting a vertex is set to the average distance of the ground truth vertices from their mean
+        cv = np.linalg.norm(np.mean(gt_vertices, axis=0) - gt_vertices, axis=1).mean()
+    return cv
+
+def compute_WED(pd_vertices, pd_edges, gt_vertices, gt_edges, cv_ins=-1/2, cv_del=-1/4, ce=1.0, normalized=True, preregister=True, single_scale=True):
     '''The function computes the Wireframe 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 (the cost in centimeters of missing a vertex, default is -1, which means 1/4 of the diameter of the ground truth mesh)
+    cv_ins: vertex insertion cost: if positive, the cost in centimeters of inserting vertex, if negative, multiplies diameter to compute cost (default is -1/2)
+    cv_del: vertex deletion cost: if positive, the cost in centimeters of deleting a vertex, if negative, multiplies diameter to compute cost (default is -1/2)
     ce: edge cost (multiplier of the edge length for edge deletion and insertion, default is 1.0)
     normalized: if True, the WED is normalized by the total length of the ground truth edges
     preregister: if True, the predicted vertices have their mean and scale matched to the ground truth vertices
     '''
     
-    # Vertex coordinates are in centimeters. When cv and ce are set to 100.0 and 1.0 respectively, 
-    # missing a vertex is equivanlent predicting it 1 meter away from the ground truth vertex.
-    # This is equivalent to setting cv=1 and ce=1 when the vertex coordinates are in meters.
-    # When a negative cv value is set (the default behavior), cv is reset to 1/4 of the diameter of the ground truth wireframe.
-    
     pd_vertices = np.array(pd_vertices)
     gt_vertices = np.array(gt_vertices)
+    pd_edges = np.array(pd_edges)
+    gt_edges = np.array(gt_edges)        
     
-    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
+        
+    cv_del = update_cv(cv_del, gt_vertices)
+    cv_ins = update_cv(cv_ins, gt_vertices)
     
     # Step 0: Prenormalize / preregister
     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 1: Bipartite Matching
     distances = cdist(pd_vertices, gt_vertices, metric='euclidean')
     row_ind, col_ind = linear_sum_assignment(distances)
@@ -70,15 +73,15 @@ def compute_WED(pd_vertices, pd_edges, gt_vertices, gt_edges, cv=-1, ce=1.0, nor
     # Step 2: Vertex Translation
     translation_costs = np.sum(distances[row_ind, col_ind])
     
-    # Additional: Vertex Deletion
+    # Step 3: Vertex Deletion
     unmatched_pd_indices = set(range(len(pd_vertices))) - set(row_ind)
-    deletion_costs = cv * len(unmatched_pd_indices)  
+    deletion_costs = cv_del * len(unmatched_pd_indices)  
     
-    # Step 3: Vertex Insertion
+    # Step 4: Vertex Insertion
     unmatched_gt_indices = set(range(len(gt_vertices))) - set(col_ind)
-    insertion_costs = cv * len(unmatched_gt_indices)  
+    insertion_costs = cv_ins * len(unmatched_gt_indices)  
     
-    # Step 4: Edge Deletion and Insertion
+    # 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]
     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]))

setup.py

@@ -6,7 +6,7 @@ with open('requirements.txt') as f:
     required = f.read().splitlines()
 
 setup(name='hoho',
-	version='0.0.3',
+	version='0.0.4',
 	description='Tools and utilites for the HoHo Dataset and S23DR Competition',
 	url='usm3d.github.io',
 	author='Jack Langerman, Dmytro Mishkin, S23DR Orgainizing Team',