(
imshapes,
edges,
edge2pred_i,
edge2pred_j,
edge2conf_i,
edge2conf_j,
im_conf,
min_conf_thr,
device,
has_im_poses=True,
niter_PnP=10,
verbose=True,
save_score_path=None,
)
| 52 | |
| 53 | |
| 54 | def minimum_spanning_tree( |
| 55 | imshapes, |
| 56 | edges, |
| 57 | edge2pred_i, |
| 58 | edge2pred_j, |
| 59 | edge2conf_i, |
| 60 | edge2conf_j, |
| 61 | im_conf, |
| 62 | min_conf_thr, |
| 63 | device, |
| 64 | has_im_poses=True, |
| 65 | niter_PnP=10, |
| 66 | verbose=True, |
| 67 | save_score_path=None, |
| 68 | ): |
| 69 | n_imgs = len(imshapes) |
| 70 | eadge_and_scores = compute_edge_scores(map(i_j_ij, edges), edge2conf_i, edge2conf_j) |
| 71 | sparse_graph = -dict_to_sparse_graph(eadge_and_scores) |
| 72 | msp = sp.csgraph.minimum_spanning_tree(sparse_graph).tocoo() |
| 73 | |
| 74 | # temp variable to store 3d points |
| 75 | pts3d = [None] * len(imshapes) |
| 76 | |
| 77 | todo = sorted(zip(-msp.data, msp.row, msp.col)) # sorted edges |
| 78 | im_poses = [None] * n_imgs |
| 79 | im_focals = [None] * n_imgs |
| 80 | |
| 81 | # init with strongest edge |
| 82 | score, i, j = todo.pop() |
| 83 | if verbose: |
| 84 | print(f" init edge ({i}*,{j}*) {score=}") |
| 85 | i_j = edge_str(i, j) |
| 86 | |
| 87 | pts3d[i] = edge2pred_i[i_j].clone() |
| 88 | pts3d[j] = edge2pred_j[i_j].clone() |
| 89 | done = {i, j} |
| 90 | if has_im_poses: |
| 91 | im_poses[i] = torch.eye(4, device=device) |
| 92 | im_focals[i] = estimate_focal(edge2pred_i[i_j]) |
| 93 | |
| 94 | # set initial pointcloud based on pairwise graph |
| 95 | msp_edges = [(i, j)] |
| 96 | while todo: |
| 97 | # each time, predict the next one |
| 98 | score, i, j = todo.pop() |
| 99 | |
| 100 | if im_focals[i] is None: |
| 101 | im_focals[i] = estimate_focal(edge2pred_i[i_j]) |
| 102 | |
| 103 | if i in done: |
| 104 | if verbose: |
| 105 | print(f" init edge ({i},{j}*) {score=}") |
| 106 | assert j not in done |
| 107 | # align pred[i] with pts3d[i], and then set j accordingly |
| 108 | i_j = edge_str(i, j) |
| 109 | s, R, T = rigid_points_registration( |
| 110 | edge2pred_i[i_j], pts3d[i], conf=edge2conf_i[i_j] |
| 111 | ) |
no test coverage detected