update metric

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

@@ -6,40 +6,96 @@ 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])
+    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)
+    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 compute_WED(pd_vertices, pd_edges, gt_vertices, gt_edges, cv=1.0, ce=1.0, normalized=True, squared=False):
+
+def compute_WED(pd_vertices, pd_edges, gt_vertices, gt_edges, cv=100.0, ce=1.0, normalized=True, prenorm=False, preregister=True, register=True, single_scale=True):
     pd_vertices = np.array(pd_vertices)
     gt_vertices = np.array(gt_vertices)
-    pd_vertices = zeromean_normalize(pd_vertices)
-    gt_vertices = zeromean_normalize(gt_vertices)
+    
+    # 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 1: Bipartite Matching
-    if squared:
-        distances = cdist(pd_vertices, gt_vertices, metric='sqeuclidean')
-    else:
-        distances = cdist(pd_vertices, gt_vertices, metric='euclidean')
 
+    # 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) 
+        
+    
+    # Step 1: Bipartite Matching
+    distances = cdist(pd_vertices, gt_vertices, metric='euclidean')
     row_ind, col_ind = linear_sum_assignment(distances)
-    # Step 2: Vertex Translation
+        
     
-    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])
+    # Step 2: Vertex Translation
+    translation_costs = np.sum(distances[row_ind, col_ind])
     
     # Additional: 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 * len(unmatched_pd_indices)  
     
     # Step 3: 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 * len(unmatched_gt_indices)  
     
     # Step 4: 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]
@@ -61,11 +117,12 @@ 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
-    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)
+    # 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()
         WED = WED / total_length_of_gt_edges
-    print ("Total length", total_length_of_gt_edges)
+        
+    # print ("Total length", total_length_of_gt_edges)
     return WED

requirements.txt

@@ -3,4 +3,6 @@ pillow
 webdataset
 trimesh
 scipy
-datasets
+datasets
+ipywidgets
+matplotlib