def make_xy(pair_index, img, kp, desc, aff, K, R, t, cur_folder): kp_initial = copy.deepcopy(kp) xs = [] xs_initial = [] ys = [] Rs = [] ts = [] img1s = [] img2s = [] cx1s = [] cy1s = [] f1s = [] cx2s = [] cy2s = [] f2s = [] affine = [] # Create a random folder in scratch dump_dir = os.path.join(cur_folder, "dump") if not os.path.exists(dump_dir): os.makedirs(dump_dir) pool_arg = [] for idx in range(pair_index.shape[0]): ii = pair_index[idx, 0] jj = pair_index[idx, 1] print("\rExtracting keypoints {} / {}".format(idx, pair_index.shape[0]), end="") sys.stdout.flush() # Check and extract keypoints if necessary for i in [ii, jj]: i = int(i) dump_file = os.path.join(dump_dir, "kp-aff-desc-{}.h5".format(i)) if not os.path.exists(dump_file): # Correct coordinates using K cx = K[i, 2] cy = K[i, 5] # Correct focals fx = K[i, 0] fy = K[i, 4] kp[i] = (kp[i] - np.array([[cx, cy]])) / np.asarray([[fx, fy]]) # Write descs to harddisk to parallize dump_dict = {} dump_dict["kp_normal"] = kp[i] dump_dict["kp"] = kp_initial[i] dump_dict["desc"] = desc[i] dump_dict["aff"] = aff[i] saveh5(dump_dict, dump_file) else: dump_dict = loadh5(dump_file) kp[i] = dump_dict["kp_normal"] kp_initial[i] = dump_dict["kp"] desc[i] = dump_dict["desc"] aff[i] = dump_dict["aff"] pool_arg += [(dump_dir, idx, int(ii), int(jj))] print("") # Run mp job ratio_CPU = 0.9 number_of_process = int(ratio_CPU * mp.cpu_count()) pool = mp.Pool(processes=number_of_process) manager = mp.Manager() queue = manager.Queue() for idx_arg in xrange(len(pool_arg)): pool_arg[idx_arg] = pool_arg[idx_arg] + (queue, ) # map async pool_res = pool.map_async(dump_data_pair, pool_arg) # monitor loop while True: if pool_res.ready(): break else: size = queue.qsize() print("\rDistMat {} / {}".format(size, len(pool_arg)), end="") sys.stdout.flush() time.sleep(1) pool.close() pool.join() print("") # Pack data idx = 0 total_num = 0 good_num = 0 bad_num = 0 ratio1 = [] ratio2 = [] for idx in range(pair_index.shape[0]): ii = int(pair_index[idx, 0]) jj = int(pair_index[idx, 1]) print("\rWorking on {} / {}".format(idx, pair_index.shape[0]), end="") sys.stdout.flush() K1 = K[ii].reshape(3, 3) K2 = K[jj].reshape(3, 3) # ------------------------------ # Get dR R_i = R[ii].reshape(3, 3) R_j = R[jj].reshape(3, 3) dR = np.dot(R_j, R_i.T) # Get dt t_i = t[ii].reshape(3, 1) t_j = t[jj].reshape(3, 1) dt = t_j - np.dot(dR, t_i) # ------------------------------ # Get keypoints for the first image x1 = kp[ii] y1 = np.concatenate((kp[ii], np.ones((kp[ii].shape[0], 1))), axis=1) # Project the first points into the second image y1p = np.matmul(dR[None], y1[..., None]) + dt[None] # move back to the canonical plane x1p = y1p[:, :2, 0] / y1p[:, 2, 0][..., None] # ------------------------------ # Get keypoints for the second image x2 = kp[jj] # # DEBUG ------------------------------ # # Check if the image projections make sense # draw_val_res( # img[ii], # img[jj], # x1, x1p, np.random.rand(x1.shape[0]) < 0.1, # (img[ii][0].shape[1] - 1.0) * 0.5, # (img[ii][0].shape[0] - 1.0) * 0.5, # parse_geom(geom, geom_type)["K"][ii, 0, 0], # (img[jj][0].shape[1] - 1.0) * 0.5, # (img[jj][0].shape[0] - 1.0) * 0.5, # parse_geom(geom, geom_type)["K"][jj, 0, 0], # "./debug_imgs/", # "debug_img{:04d}.png".format(idx) # ) # ------------------------------ # create x1, y1, x2, y2 as a matrix combo x1mat = np.repeat(x1[:, 0][..., None], len(x2), axis=-1) y1mat = np.repeat(x1[:, 1][..., None], len(x2), axis=1) x1pmat = np.repeat(x1p[:, 0][..., None], len(x2), axis=-1) y1pmat = np.repeat(x1p[:, 1][..., None], len(x2), axis=1) x2mat = np.repeat(x2[:, 0][None], len(x1), axis=0) y2mat = np.repeat(x2[:, 1][None], len(x1), axis=0) # Load precomputed nearest neighbors idx_sort = loadh5( os.path.join(dump_dir, "idx_sort-{}-{}.h5".format(ii, jj)))["idx_sort"] # Move back to tuples idx_sort = (idx_sort[0], idx_sort[1]) x1mat = x1mat[idx_sort] y1mat = y1mat[idx_sort] x1pmat = x1pmat[idx_sort] y1pmat = y1pmat[idx_sort] x2mat = x2mat[idx_sort] y2mat = y2mat[idx_sort] # Turn into x1, x1p, x2 x1 = np.concatenate( [x1mat.reshape(-1, 1), y1mat.reshape(-1, 1)], axis=1) x1p = np.concatenate([x1pmat.reshape(-1, 1), y1pmat.reshape(-1, 1)], axis=1) x2 = np.concatenate( [x2mat.reshape(-1, 1), y2mat.reshape(-1, 1)], axis=1) # make xs in NHWC xs += [ np.concatenate([x1, x2], axis=1).T.reshape(4, 1, -1).transpose( (1, 2, 0)) ] # ------------------------------ # Get the geodesic distance using with x1, x2, dR, dt if config.obj_geod_type == "sampson": geod_d = get_sampsons(x1, x2, dR, dt) elif config.obj_geod_type == "episqr": geod_d = get_episqr(x1, x2, dR, dt) elif config.obj_geod_type == "episym": geod_d = get_episym(x1, x2, dR, dt) # geod_d = get_episym(x1, x2, dR, dt, K1_inv, K2_inv, K1, K2) # Get *rough* reprojection errors. Note that the depth may be noisy. We # ended up not using this... reproj_d = np.sum((x2 - x1p)**2, axis=1) # count inliers and outliers total_num += len(geod_d) good_num += np.sum((geod_d < config.obj_geod_th)) bad_num += np.sum((geod_d >= config.obj_geod_th)) ''' mask = np.zeros(len(geod_d)) for i in range(len(geod_d)): if geod_d[i] < config.obj_geod_th: mask[i] = 1 np.savetxt(os.path.join(dump_dir, "mask-{}-{}.txt".format(ii, jj)), mask) ''' ys += [np.stack([geod_d, reproj_d], axis=1)] # Save R, t for evaluation Rs += [np.array(dR).reshape(3, 3)] # normalize t before saving dtnorm = np.sqrt(np.sum(dt**2)) assert (dtnorm > 1e-5) dt /= dtnorm ts += [np.array(dt).flatten()] # Save img1 and img2 for display img1s += [img[ii]] img2s += [img[jj]] cx = K[ii, 2] cy = K[ii, 5] cx1s += [cx] cy1s += [cy] cx = K[jj, 2] cy = K[jj, 5] cx2s += [cx] cy2s += [cy] fx = K[ii, 0] fy = K[ii, 4] if np.isclose(fx, fy): f = fx else: f = (fx, fy) f1s += [f] fx = K[jj, 0] fy = K[jj, 4] if np.isclose(fx, fy): f = fx else: f = (fx, fy) f2s += [f] # Generate xs_initial and T-transform x1 = kp_initial[ii] x2 = kp_initial[jj] aff1 = aff[ii] aff2 = aff[jj] aff1 = np.repeat(aff1[:, :][..., None], len(aff2), axis=-1).transpose(0, 2, 1) aff2 = np.repeat(aff2[:, :][None], len(aff1), axis=0) x1mat = np.repeat(x1[:, 0][..., None], len(x2), axis=-1) y1mat = np.repeat(x1[:, 1][..., None], len(x2), axis=1) x2mat = np.repeat(x2[:, 0][None], len(x1), axis=0) y2mat = np.repeat(x2[:, 1][None], len(x1), axis=0) idx_sort = loadh5( os.path.join(dump_dir, "idx_sort-{}-{}.h5".format(ii, jj)))["idx_sort"] # Move back to tuples idx_sort = (idx_sort[0], idx_sort[1]) x1mat = x1mat[idx_sort] y1mat = y1mat[idx_sort] x2mat = x2mat[idx_sort] y2mat = y2mat[idx_sort] aff1 = aff1[idx_sort] aff2 = aff2[idx_sort] # Turn into x1, x1p, x2 x1 = np.concatenate( [x1mat.reshape(-1, 1), y1mat.reshape(-1, 1)], axis=1) x2 = np.concatenate( [x2mat.reshape(-1, 1), y2mat.reshape(-1, 1)], axis=1) affine += [ np.concatenate( [aff1.reshape(-1, 9), aff2.reshape(-1, 9)], axis=1) ] xs_initial += [ np.concatenate([x1, x2], axis=1).T.reshape(4, 1, -1).transpose( (1, 2, 0)) ] print("") # Do *not* convert to numpy arrays, as the number of keypoints may differ # now. Simply return it print(".... done") if total_num > 0: print(" Good pairs = {}, Total pairs = {}, Ratio = {}".format( good_num, total_num, float(good_num) / float(total_num))) print(" Bad pairs = {}, Total pairs = {}, Ratio = {}".format( bad_num, total_num, float(bad_num) / float(total_num))) res_dict = {} res_dict["xs"] = xs res_dict["xs_initial"] = xs_initial res_dict["affine"] = affine res_dict["ys"] = ys res_dict["Rs"] = Rs res_dict["ts"] = ts res_dict["img1s"] = img1s res_dict["cx1s"] = cx1s res_dict["cy1s"] = cy1s res_dict["f1s"] = f1s res_dict["img2s"] = img2s res_dict["cx2s"] = cx2s res_dict["cy2s"] = cy2s res_dict["f2s"] = f2s return res_dict
def make_xy(num_sample, pairs, kp, z, desc, img, geom, vis, depth, geom_type, cur_folder): xs = [] ys = [] Rs = [] ts = [] mutuals = [] ratios = [] img1s = [] img2s = [] cx1s = [] cy1s = [] f1s = [] cx2s = [] cy2s = [] f2s = [] # Create a random folder in scratch dump_dir = os.path.join(cur_folder, "dump") if not os.path.exists(dump_dir): os.makedirs(dump_dir) # randomly suffle the pairs and select num_sample amount np.random.seed(1234) cur_pairs = [ pairs[_i] for _i in np.random.permutation(len(pairs))[:num_sample] ] idx = 0 for ii, jj in cur_pairs: idx += 1 print( "\rExtracting keypoints {} / {}".format(idx, len(cur_pairs)), end="") sys.stdout.flush() # Check and extract keypoints if necessary for i in [ii, jj]: dump_file = os.path.join(dump_dir, "kp-z-desc-{}.h5".format(i)) if not os.path.exists(dump_file): if kp[i] is None: cv_kp, cv_desc = sift.detectAndCompute(img[i].transpose( 1, 2, 0), None) cx = (img[i][0].shape[1] - 1.0) * 0.5 cy = (img[i][0].shape[0] - 1.0) * 0.5 # Correct coordinates using K cx += parse_geom(geom, geom_type)["K"][i, 0, 2] cy += parse_geom(geom, geom_type)["K"][i, 1, 2] xy = np.array([_kp.pt for _kp in cv_kp]) # Correct focals fx = parse_geom(geom, geom_type)["K"][i, 0, 0] fy = parse_geom(geom, geom_type)["K"][i, 1, 1] kp[i] = ( xy - np.array([[cx, cy]]) ) / np.asarray([[fx, fy]]) desc[i] = cv_desc if z[i] is None: cx = (img[i][0].shape[1] - 1.0) * 0.5 cy = (img[i][0].shape[0] - 1.0) * 0.5 fx = parse_geom(geom, geom_type)["K"][i, 0, 0] fy = parse_geom(geom, geom_type)["K"][i, 1, 1] xy = kp[i] * np.asarray([[fx, fy]]) + np.array([[cx, cy]]) if len(depth) > 0: z[i] = depth[i][ 0, np.round(xy[:, 1]).astype(int), np.round(xy[:, 0]).astype(int)][..., None] else: z[i] = np.ones((xy.shape[0], 1)) # Write descs to harddisk to parallize dump_dict = {} dump_dict["kp"] = kp[i] dump_dict["z"] = z[i] dump_dict["desc"] = desc[i] saveh5(dump_dict, dump_file) else: dump_dict = loadh5(dump_file) kp[i] = dump_dict["kp"] z[i] = dump_dict["z"] desc[i] = dump_dict["desc"] print("") # Create arguments pool_arg = [] idx = 0 for ii, jj in cur_pairs: idx += 1 pool_arg += [(dump_dir, idx, ii, jj)] # Run mp job ratio_CPU = 0.8 number_of_process = int(ratio_CPU * mp.cpu_count()) pool = mp.Pool(processes=number_of_process) manager = mp.Manager() queue = manager.Queue() # for debugging in dump_data_pair # dump_data_pair(pool_arg[1] + (queue,)) for idx_arg in xrange(len(pool_arg)): pool_arg[idx_arg] = pool_arg[idx_arg] + (queue,) # map async pool_res = pool.map_async(dump_data_pair, pool_arg) # monitor loop while True: if pool_res.ready(): break else: size = queue.qsize() print("\rDistMat {} / {}".format(size, len(pool_arg)), end="") sys.stdout.flush() time.sleep(1) pool.close() pool.join() print("") # Pack data idx = 0 total_num = 0 good_num = 0 bad_num = 0 for ii, jj in cur_pairs: idx += 1 print("\rWorking on {} / {}".format(idx, len(cur_pairs)), end="") sys.stdout.flush() # ------------------------------ # Get dR R_i = parse_geom(geom, geom_type)["R"][ii] R_j = parse_geom(geom, geom_type)["R"][jj] dR = np.dot(R_j, R_i.T) # Get dt t_i = parse_geom(geom, geom_type)["t"][ii].reshape([3, 1]) t_j = parse_geom(geom, geom_type)["t"][jj].reshape([3, 1]) dt = t_j - np.dot(dR, t_i) # ------------------------------ # Get sift points for the first image x1 = kp[ii] y1 = np.concatenate([kp[ii] * z[ii], z[ii]], axis=1) # Project the first points into the second image y1p = np.matmul(dR[None], y1[..., None]) + dt[None] # move back to the canonical plane x1p = y1p[:, :2, 0] / y1p[:, 2, 0][..., None] # ------------------------------ # Get sift points for the second image x2 = kp[jj] # # DEBUG ------------------------------ # # Check if the image projections make sense # draw_val_res( # img[ii], # img[jj], # x1, x1p, np.random.rand(x1.shape[0]) < 0.1, # (img[ii][0].shape[1] - 1.0) * 0.5, # (img[ii][0].shape[0] - 1.0) * 0.5, # parse_geom(geom, geom_type)["K"][ii, 0, 0], # (img[jj][0].shape[1] - 1.0) * 0.5, # (img[jj][0].shape[0] - 1.0) * 0.5, # parse_geom(geom, geom_type)["K"][jj, 0, 0], # "./debug_imgs/", # "debug_img{:04d}.png".format(idx) # ) # ------------------------------ # create x1, y1, x2, y2 as a matrix combo x1mat = np.repeat(x1[:, 0][..., None], len(x2), axis=-1) y1mat = np.repeat(x1[:, 1][..., None], len(x2), axis=1) x1pmat = np.repeat(x1p[:, 0][..., None], len(x2), axis=-1) y1pmat = np.repeat(x1p[:, 1][..., None], len(x2), axis=1) x2mat = np.repeat(x2[:, 0][None], len(x1), axis=0) y2mat = np.repeat(x2[:, 1][None], len(x1), axis=0) # Load precomputed nearest neighbors, ratios and mutual idx_sort = loadh5(os.path.join( dump_dir, "idx_sort-{}-{}.h5".format(ii, jj)))["idx_sort"] mutual = loadh5(os.path.join( dump_dir, "mutual_ratio-{}-{}.h5".format(ii, jj)))["mutual"] ratio = loadh5(os.path.join( dump_dir, "mutual_ratio-{}-{}.h5".format(ii, jj)))["ratio"] mutuals += [mutual] ratios += [ratio] # Move back to tuples idx_sort = (idx_sort[0], idx_sort[1]) x1mat = x1mat[idx_sort] y1mat = y1mat[idx_sort] x1pmat = x1pmat[idx_sort] y1pmat = y1pmat[idx_sort] x2mat = x2mat[idx_sort] y2mat = y2mat[idx_sort] # Turn into x1, x1p, x2 x1 = np.concatenate( [x1mat.reshape(-1, 1), y1mat.reshape(-1, 1)], axis=1) x1p = np.concatenate( [x1pmat.reshape(-1, 1), y1pmat.reshape(-1, 1)], axis=1) x2 = np.concatenate( [x2mat.reshape(-1, 1), y2mat.reshape(-1, 1)], axis=1) # make xs in NHWC xs += [ np.concatenate([x1, x2], axis=1).T.reshape(4, 1, -1).transpose( (1, 2, 0)) ] # ------------------------------ # Get the geodesic distance using with x1, x2, dR, dt if config.obj_geod_type == "sampson": geod_d = get_sampsons(x1, x2, dR, dt) elif config.obj_geod_type == "episqr": geod_d = get_episqr(x1, x2, dR, dt) elif config.obj_geod_type == "episym": geod_d = get_episym(x1, x2, dR, dt) # Get *rough* reprojection errors. Note that the depth may be noisy. We # ended up not using this... reproj_d = np.sum((x2 - x1p)**2, axis=1) # count inliers and outliers total_num += len(geod_d) good_num += np.sum((geod_d < config.obj_geod_th)) bad_num += np.sum((geod_d >= config.obj_geod_th)) ys += [np.stack([geod_d, reproj_d], axis=1)] # Save R, t for evaluation Rs += [np.array(dR).reshape(3, 3)] # normalize t before saving dtnorm = np.sqrt(np.sum(dt**2)) assert (dtnorm > 1e-5) dt /= dtnorm ts += [np.array(dt).flatten()] # Save img1 and img2 for display img1s += [img[ii]] img2s += [img[jj]] cx = (img[ii][0].shape[1] - 1.0) * 0.5 cy = (img[ii][0].shape[0] - 1.0) * 0.5 # Correct coordinates using K cx += parse_geom(geom, geom_type)["K"][ii, 0, 2] cy += parse_geom(geom, geom_type)["K"][ii, 1, 2] fx = parse_geom(geom, geom_type)["K"][ii, 0, 0] fy = parse_geom(geom, geom_type)["K"][ii, 1, 1] if np.isclose(fx, fy): f = fx else: f = (fx, fy) cx1s += [cx] cy1s += [cy] f1s += [f] cx = (img[jj][0].shape[1] - 1.0) * 0.5 cy = (img[jj][0].shape[0] - 1.0) * 0.5 # Correct coordinates using K cx += parse_geom(geom, geom_type)["K"][jj, 0, 2] cy += parse_geom(geom, geom_type)["K"][jj, 1, 2] fx = parse_geom(geom, geom_type)["K"][jj, 0, 0] fy = parse_geom(geom, geom_type)["K"][jj, 1, 1] if np.isclose(fx, fy): f = fx else: f = (fx, fy) cx2s += [cx] cy2s += [cy] f2s += [f] # Do *not* convert to numpy arrays, as the number of keypoints may differ # now. Simply return it print(".... done") if total_num > 0: print(" Good pairs = {}, Total pairs = {}, Ratio = {}".format( good_num, total_num, float(good_num) / float(total_num))) print(" Bad pairs = {}, Total pairs = {}, Ratio = {}".format( bad_num, total_num, float(bad_num) / float(total_num))) res_dict = {} res_dict["xs"] = xs res_dict["ys"] = ys res_dict["Rs"] = Rs res_dict["ts"] = ts res_dict["img1s"] = img1s res_dict["cx1s"] = cx1s res_dict["cy1s"] = cy1s res_dict["f1s"] = f1s res_dict["img2s"] = img2s res_dict["cx2s"] = cx2s res_dict["cy2s"] = cy2s res_dict["f2s"] = f2s res_dict["mutuals"] = mutuals res_dict["ratios"] = ratios res_dict["pairs"] = cur_pairs return res_dict