def tps_rpm_bij(x_nd, y_md, fsolve, gsolve, n_iter=20, reg_init=.1, reg_final=.001, rad_init=.1, rad_final=.005, rot_reg=1e-3, outlierprior=1e-1, outlierfrac=2e-1, vis_cost_xy=None, return_corr=False, check_solver=False): """ tps-rpm algorithm mostly as described by chui and rangaran reg_init/reg_final: regularization on curvature rad_init/rad_final: radius for correspondence calculation (meters) plotting: 0 means don't plot. integer n means plot every n iterations """ _, d = x_nd.shape regs = np.around(loglinspace(reg_init, reg_final, n_iter), BEND_COEF_DIGITS) rads = loglinspace(rad_init, rad_final, n_iter) f = ThinPlateSpline(d) scale = (np.max(y_md, axis=0) - np.min(y_md, axis=0)) / ( np.max(x_nd, axis=0) - np.min(x_nd, axis=0)) f.lin_ag = np.diag(scale) # align the mins and max f.trans_g = np.median( y_md, axis=0) - np.median(x_nd, axis=0) * scale # align the medians g = ThinPlateSpline(d) g.lin_ag = np.diag(1. / scale) g.trans_g = -np.diag(1. / scale).dot(f.trans_g) # r_N = None for i in xrange(n_iter): xwarped_nd = f.transform_points(x_nd) ywarped_md = g.transform_points(y_md) fwddist_nm = ssd.cdist(xwarped_nd, y_md, 'euclidean') invdist_nm = ssd.cdist(x_nd, ywarped_md, 'euclidean') r = rads[i] prob_nm = np.exp(-(fwddist_nm + invdist_nm) / (2 * r)) corr_nm, r_N, _ = balance_matrix(prob_nm, 10, outlierprior, outlierfrac) corr_nm += 1e-9 wt_n = corr_nm.sum(axis=1) wt_m = corr_nm.sum(axis=0) xtarg_nd = (corr_nm / wt_n[:, None]).dot(y_md) ytarg_md = (corr_nm / wt_m[None, :]).T.dot(x_nd) fsolve.solve(wt_n, xtarg_nd, regs[i], rot_reg, f) gsolve.solve(wt_m, ytarg_md, regs[i], rot_reg, g) if check_solver: f_test = fit_ThinPlateSpline(x_nd, xtarg_nd, bend_coef=regs[i], wt_n=wt_n, rot_coef=rot_reg) g_test = fit_ThinPlateSpline(y_md, ytarg_md, bend_coef=regs[i], wt_n=wt_m, rot_coef=rot_reg) tol = 1e-4 assert np.allclose(f.trans_g, f_test.trans_g, atol=tol) assert np.allclose(f.lin_ag, f_test.lin_ag, atol=tol) assert np.allclose(f.w_ng, f_test.w_ng, atol=tol) assert np.allclose(g.trans_g, g_test.trans_g, atol=tol) assert np.allclose(g.lin_ag, g_test.lin_ag, atol=tol) assert np.allclose(g.w_ng, g_test.w_ng, atol=tol) f._cost = tps.tps_cost( f.lin_ag, f.trans_g, f.w_ng, f.x_na, xtarg_nd, regs[i], wt_n=wt_n) / wt_n.mean() g._cost = tps.tps_cost( g.lin_ag, g.trans_g, g.w_ng, g.x_na, ytarg_md, regs[i], wt_n=wt_m) / wt_m.mean() if return_corr: return (f, g), corr_nm return f, g
def test_batch_tps_rpm_bij( src_ctx, tgt_ctx, T_init=1e-1, T_final=5e-3, outlierfrac=1e-2, outlierprior=1e-1, outliercutoff=0.5, em_iter=EM_ITER_CHEAP, test_ind=0, ): from lfd.tpsopt.transformations import ThinPlateSpline, set_ThinPlateSpline n_iter = len(src_ctx.bend_coefs) T_vals = loglinspace(T_init, T_final, n_iter) x_nd = src_ctx.pts[test_ind].get()[: src_ctx.dims[test_ind]] y_md = tgt_ctx.pts[0].get()[: tgt_ctx.dims[0]] (n, d) = x_nd.shape (m, _) = y_md.shape f = ThinPlateSpline(d) g = ThinPlateSpline(d) src_ctx.reset_tps_params() tgt_ctx.reset_tps_params() for i, b in enumerate(src_ctx.bend_coefs): T = T_vals[i] for _ in range(em_iter): src_ctx.transform_points() tgt_ctx.transform_points() xwarped_nd = f.transform_points(x_nd) ywarped_md = g.transform_points(y_md) gpu_xw = src_ctx.pts_w[test_ind].get()[:n, :] gpu_yw = tgt_ctx.pts_w[test_ind].get()[:m, :] assert np.allclose(xwarped_nd, gpu_xw, atol=1e-5) assert np.allclose(ywarped_md, gpu_yw, atol=1e-5) xwarped_nd = gpu_xw ywarped_md = gpu_yw src_ctx.get_target_points(tgt_ctx, outlierprior, outlierfrac, outliercutoff, T) fwddist_nm = ssd.cdist(xwarped_nd, y_md, "euclidean") invdist_nm = ssd.cdist(x_nd, ywarped_md, "euclidean") prob_nm = outlierprior * np.ones((n + 1, m + 1), np.float32) prob_nm[:n, :m] = np.exp(-(fwddist_nm + invdist_nm) / float(2 * T)) prob_nm[n, m] = outlierfrac * np.sqrt(n * m) gpu_corr = src_ctx.corr_rm[test_ind].get() gpu_corr = gpu_corr.flatten() gpu_corr = gpu_corr[: (n + 1) * (m + 1)].reshape(n + 1, m + 1).astype(np.float32) assert np.allclose(prob_nm[:n, :m], gpu_corr[:n, :m], atol=1e-5) save_prob_nm = np.array(prob_nm) save_gpu_corr = np.array(gpu_corr) prob_nm[:n, :m] = gpu_corr[:n, :m] r_coefs = np.ones(n + 1, np.float32) c_coefs = np.ones(m + 1, np.float32) a_N = np.ones((n + 1), dtype=np.float32) a_N[n] = m * outlierfrac b_M = np.ones((m + 1), dtype=np.float32) b_M[m] = n * outlierfrac for norm_iter_i in range(DEFAULT_NORM_ITERS): r_coefs = a_N / prob_nm.dot(c_coefs) rn_c_coefs = c_coefs c_coefs = b_M / r_coefs.dot(prob_nm) gpu_r_coefs = src_ctx.r_coefs[test_ind].get()[: n + 1].reshape(n + 1) gpu_c_coefs_cn = src_ctx.c_coefs_cn[test_ind].get()[: m + 1].reshape(m + 1) gpu_c_coefs_rn = src_ctx.c_coefs_rn[test_ind].get()[: m + 1].reshape(m + 1) r_diff = np.abs(r_coefs - gpu_r_coefs) rn_diff = np.abs(rn_c_coefs - gpu_c_coefs_rn) cn_diff = np.abs(c_coefs - gpu_c_coefs_cn) assert np.allclose(r_coefs, gpu_r_coefs, atol=1e-5) assert np.allclose(c_coefs, gpu_c_coefs_cn, atol=1e-5) assert np.allclose(rn_c_coefs, gpu_c_coefs_rn, atol=1e-5) prob_nm = prob_nm[:n, :m] prob_nm *= gpu_r_coefs[:n, None] rn_p_nm = prob_nm * gpu_c_coefs_rn[None, :m] cn_p_nm = prob_nm * gpu_c_coefs_cn[None, :m] wt_n = rn_p_nm.sum(axis=1) gpu_corr_cm = src_ctx.corr_cm[test_ind].get().flatten()[: (n + 1) * (m + 1)] gpu_corr_cm = gpu_corr_cm.reshape(m + 1, n + 1) ## b/c it is column major assert np.allclose(wt_n, gpu_corr_cm[m, :n], atol=1e-4) inlier = wt_n > outliercutoff xtarg_nd = np.empty((n, DATA_DIM), np.float32) xtarg_nd[inlier, :] = rn_p_nm.dot(y_md)[inlier, :] xtarg_nd[~inlier, :] = xwarped_nd[~inlier, :] wt_m = cn_p_nm.sum(axis=0) assert np.allclose(wt_m, gpu_corr[n, :m], atol=1e-4) inlier = wt_m > outliercutoff ytarg_md = np.empty((m, DATA_DIM), np.float32) ytarg_md[inlier, :] = cn_p_nm.T.dot(x_nd)[inlier, :] ytarg_md[~inlier, :] = ywarped_md[~inlier, :] xt_gpu = src_ctx.pts_t[test_ind].get()[:n, :] yt_gpu = tgt_ctx.pts_t[test_ind].get()[:m, :] assert np.allclose(xtarg_nd, xt_gpu, atol=1e-4) assert np.allclose(ytarg_md, yt_gpu, atol=1e-4) src_ctx.update_transform(b) tgt_ctx.update_transform(b) f_p_mat = src_ctx.proj_mats[b][test_ind].get()[: n + d + 1, :n] f_o_mat = src_ctx.offset_mats[b][test_ind].get()[: n + d + 1] b_p_mat = tgt_ctx.proj_mats[b][0].get()[: m + d + 1, :m] b_o_mat = tgt_ctx.offset_mats[b][0].get()[: m + d + 1] f_params = f_p_mat.dot(xtarg_nd) + f_o_mat g_params = b_p_mat.dot(ytarg_md) + b_o_mat gpu_fparams = src_ctx.tps_params[test_ind].get()[: n + d + 1] gpu_gparams = tgt_ctx.tps_params[test_ind].get()[: m + d + 1] assert np.allclose(f_params, gpu_fparams, atol=1e-4) assert np.allclose(g_params, gpu_gparams, atol=1e-4) set_ThinPlateSpline(f, x_nd, gpu_fparams) set_ThinPlateSpline(g, y_md, gpu_gparams) f._cost = tps.tps_cost(f.lin_ag, f.trans_g, f.w_ng, f.x_na, xtarg_nd, 1) g._cost = tps.tps_cost(g.lin_ag, g.trans_g, g.w_ng, g.x_na, ytarg_md, 1) gpu_cost = src_ctx.bidir_tps_cost(tgt_ctx) cpu_cost = f._cost + g._cost assert np.isclose(gpu_cost[test_ind], cpu_cost, atol=1e-4)
def tps_rpm_bij(x_nd, y_md, fsolve, gsolve,n_iter = 20, reg_init = .1, reg_final = .001, rad_init = .1, rad_final = .005, rot_reg = 1e-3, outlierprior=1e-1, outlierfrac=2e-1, vis_cost_xy=None, return_corr=False, check_solver=False): """ tps-rpm algorithm mostly as described by chui and rangaran reg_init/reg_final: regularization on curvature rad_init/rad_final: radius for correspondence calculation (meters) plotting: 0 means don't plot. integer n means plot every n iterations """ _,d=x_nd.shape regs = np.around(loglinspace(reg_init, reg_final, n_iter), BEND_COEF_DIGITS) rads = loglinspace(rad_init, rad_final, n_iter) f = ThinPlateSpline(d) scale = (np.max(y_md,axis=0) - np.min(y_md,axis=0)) / (np.max(x_nd,axis=0) - np.min(x_nd,axis=0)) f.lin_ag = np.diag(scale) # align the mins and max f.trans_g = np.median(y_md,axis=0) - np.median(x_nd,axis=0) * scale # align the medians g = ThinPlateSpline(d) g.lin_ag = np.diag(1./scale) g.trans_g = -np.diag(1./scale).dot(f.trans_g) # r_N = None for i in xrange(n_iter): xwarped_nd = f.transform_points(x_nd) ywarped_md = g.transform_points(y_md) fwddist_nm = ssd.cdist(xwarped_nd, y_md,'euclidean') invdist_nm = ssd.cdist(x_nd, ywarped_md,'euclidean') r = rads[i] prob_nm = np.exp( -(fwddist_nm + invdist_nm) / (2*r) ) corr_nm, r_N, _ = balance_matrix(prob_nm, 10, outlierprior, outlierfrac) corr_nm += 1e-9 wt_n = corr_nm.sum(axis=1) wt_m = corr_nm.sum(axis=0) xtarg_nd = (corr_nm/wt_n[:,None]).dot(y_md) ytarg_md = (corr_nm/wt_m[None,:]).T.dot(x_nd) fsolve.solve(wt_n, xtarg_nd, regs[i], rot_reg, f) gsolve.solve(wt_m, ytarg_md, regs[i], rot_reg, g) if check_solver: f_test = fit_ThinPlateSpline(x_nd, xtarg_nd, bend_coef = regs[i], wt_n=wt_n, rot_coef = rot_reg) g_test = fit_ThinPlateSpline(y_md, ytarg_md, bend_coef = regs[i], wt_n=wt_m, rot_coef = rot_reg) tol = 1e-4 assert np.allclose(f.trans_g, f_test.trans_g, atol=tol) assert np.allclose(f.lin_ag, f_test.lin_ag, atol=tol) assert np.allclose(f.w_ng, f_test.w_ng, atol=tol) assert np.allclose(g.trans_g, g_test.trans_g, atol=tol) assert np.allclose(g.lin_ag, g_test.lin_ag, atol=tol) assert np.allclose(g.w_ng, g_test.w_ng, atol=tol) f._cost = tps.tps_cost(f.lin_ag, f.trans_g, f.w_ng, f.x_na, xtarg_nd, regs[i], wt_n=wt_n)/wt_n.mean() g._cost = tps.tps_cost(g.lin_ag, g.trans_g, g.w_ng, g.x_na, ytarg_md, regs[i], wt_n=wt_m)/wt_m.mean() if return_corr: return (f, g), corr_nm return f,g