separate-cv

hoho/vis.py

@@ -133,7 +133,8 @@ def create_image_grid(images, target_length=312, num_per_row=2):
     return grid_img
 
 
-import matplotlib
+import matplotlib.pyplot as plt
+
 def visualize_depth(depth, min_depth=None, max_depth=None, cmap='rainbow'):
     depth = np.array(depth)
     
@@ -148,7 +149,7 @@ def visualize_depth(depth, min_depth=None, max_depth=None, cmap='rainbow'):
     depth = np.clip(depth, 0, 1)
     
     # Use the matplotlib colormap to convert the depth to an RGB image
-    cmap = matplotlib.cm.get_cmap(cmap)
+    cmap = plt.get_cmap(cmap)
     depth_image = (cmap(depth) * 255).astype(np.uint8)
     
     # Convert the depth image to a PIL image

hoho/wed.py

@@ -2,43 +2,97 @@ from scipy.spatial.distance import cdist
 from scipy.optimize import linear_sum_assignment
 import numpy as np 
 
-def compute_WED(pd_vertices, pd_edges, gt_vertices, gt_edges, cv=1.0, ce=1.0, normalized=True, squared=False):
+
+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)
+    std_target = np.std(target_verts, axis=0)
+    std_in = np.std(verts_to_transform, axis=0)
+    
+    if np.any(std_in == 0):
+        std_in[std_in == 0] = 1
+    if np.any(std_target == 0):
+        std_target[std_target == 0] = 1
+    if np.any(np.isnan(std_in)):
+        std_in[np.isnan(std_in)] = 1
+    if np.any(np.isnan(std_target)):
+        std_target[np.isnan(std_target)] = 1
+        
+    if single_scale:
+        std_target = np.linalg.norm(std_target)
+        std_in = np.linalg.norm(std_in)
+    
+    transformed_verts = (verts_to_transform - mu_in) / std_in
+    transformed_verts = transformed_verts * std_target + mu_target
+    
+    return transformed_verts
+
+
+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_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
+    '''
+    
     pd_vertices = np.array(pd_vertices)
     gt_vertices = np.array(gt_vertices)
     pd_edges = np.array(pd_edges)
-    gt_edges = np.array(gt_edges)
+    gt_edges = np.array(gt_edges)        
+    
+        
+    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)
     
-    # Step 1: Bipartite Matching
-    if squared:
-        distances = cdist(pd_vertices, gt_vertices, metric='sqeuclidean')
-    else:
-        distances = cdist(pd_vertices, gt_vertices, metric='euclidean')
 
+    # Step 1: Bipartite Matching
+    distances = cdist(pd_vertices, gt_vertices, metric='euclidean')
     row_ind, col_ind = linear_sum_assignment(distances)
+        
     
     # Step 2: Vertex Translation
+    translation_costs = np.sum(distances[row_ind, col_ind])
     
-    if squared:
-        translation_costs = cv * np.sqrt(np.sum(distances[row_ind, col_ind]))
-    else:
-        translation_costs = cv * 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)  # Assuming a fixed cost for vertex deletion
+    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)  # Assuming a fixed cost for vertex insertion
+    insertion_costs = cv_ins * len(unmatched_gt_indices)  
     
-    # Step 4: Edge Deletion and Insertion
-    updated_pd_edges = [(row_ind[np.where(col_ind == edge[0])[0][0]], row_ind[np.where(col_ind == edge[1])[0][0]]) for edge in pd_edges if edge[0] in col_ind and edge[1] in col_ind]
-    pd_edges_set = set(map(tuple, updated_pd_edges))
-    gt_edges_set = set(map(tuple, gt_edges))
+    # 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]))
+
     
     # 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)
+
     
     # Insert missing edges from ground truth
     edges_to_insert = gt_edges_set - pd_edges_set
@@ -46,9 +100,11 @@ def compute_WED(pd_vertices, pd_edges, gt_vertices, gt_edges, cv=1.0, ce=1.0, no
     
     # Step 5: Calculation of WED
     WED = 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()
         WED = WED / total_length_of_gt_edges
-       
+        
+    # print ("Total length", total_length_of_gt_edges)
     return WED

requirements.txt

@@ -1,8 +1,10 @@
+datasets
+ipywidgets
+matplotlib
 numpy
 pillow
-webdataset
-trimesh
-scipy
-datasets
+plotly
 pycolmap
-plotly
+scipy
+trimesh
+webdataset

setup.py

@@ -6,7 +6,7 @@ with open('requirements.txt') as f:
     required = f.read().splitlines()
 
 setup(name='hoho',
-	version='0.0.2',
+	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',