File size: 6,173 Bytes
a83935b 7258769 fc034ff a83935b 725639c a0cd1aa 725639c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | import numpy as np
from PIL import Image as PImage
import io
from scipy.spatial.distance import cdist
from scipy.optimize import linear_sum_assignment
from hoho.read_write_colmap import read_cameras_binary, read_images_binary, read_points3D_binary
def convert_entry_to_human_readable(entry):
out = {}
already_good = ['__key__', 'wf_vertices', 'wf_edges', 'edge_semantics', 'mesh_vertices', 'mesh_faces', 'face_semantics', 'K', 'R', 't']
for k, v in entry.items():
if k in already_good:
out[k] = v
continue
if k == 'points3d':
out[k] = read_points3D_binary(fid=io.BytesIO(v))
if k == 'cameras':
out[k] = read_cameras_binary(fid=io.BytesIO(v))
if k == 'images':
out[k] = read_images_binary(fid=io.BytesIO(v))
if k in ['ade20k', 'gestalt']:
out[k] = [PImage.open(io.BytesIO(x)).convert('RGB') for x in v]
if k == 'depthcm':
out[k] = [PImage.open(io.BytesIO(x)) for x in entry['depthcm']]
return out
def to_K(f, cx, cy):
K = np.eye(3)
K[0,0] = K[1,1] = f
K[0,2] = cx
K[1,2] = cy
return K
def quaternion_to_rotation_matrix(qvec):
qw, qx, qy, qz = qvec
R = np.array([
[1 - 2*qy**2 - 2*qz**2, 2*qx*qy - 2*qz*qw, 2*qx*qz + 2*qy*qw],
[2*qx*qy + 2*qz*qw, 1 - 2*qx**2 - 2*qz**2, 2*qy*qz - 2*qx*qw],
[2*qx*qz - 2*qy*qw, 2*qy*qz + 2*qx*qw, 1 - 2*qx**2 - 2*qy**2]
])
return R
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
return cv
def my_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/4)
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)
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
distances = cdist(pd_vertices, gt_vertices, metric='euclidean')
row_ind, col_ind = linear_sum_assignment(distances)
print(row_ind, col_ind)
# Step 2: Vertex Translation
translation_costs = np.sum(distances[row_ind, col_ind])
# Step 3: Vertex Deletion
unmatched_pd_indices = set(range(len(pd_vertices))) - set(row_ind)
deletion_costs = cv_del * len(unmatched_pd_indices)
# Step 4: Vertex Insertion
unmatched_gt_indices = set(range(len(gt_vertices))) - set(col_ind)
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 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]))
# Delete edges not in ground truth
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 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 if len(edge) == 2)
# 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)
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 |