def __init__(self, demos, actionfile=None,
n_iter=settings.N_ITER, em_iter=settings.EM_ITER,
reg_init=settings.REG[0], reg_final=settings.REG[1],
rad_init=settings.RAD[0], rad_final=settings.RAD[1],
rot_reg=settings.ROT_REG,
outlierprior=settings.OUTLIER_PRIOR, outlierfrac=settings.OURLIER_FRAC,
prior_fn=None,
f_solver_factory=solver.AutoTpsSolverFactory(),
g_solver_factory=solver.AutoTpsSolverFactory(use_cache=False)):
if not lfd.registration._has_cuda:
raise NotImplementedError("CUDA not installed")
super(BatchGpuTpsRpmBijRegistrationFactory, self).__init__(demos=demos,
n_iter=n_iter, em_iter=em_iter,
reg_init=reg_init, reg_final=reg_final,
rad_init=rad_init, rad_final=rad_final,
rot_reg=rot_reg,
outlierprior=outlierprior, outlierfrac=outlierfrac,
prior_fn=prior_fn,
f_solver_factory=f_solver_factory, g_solver_factory=g_solver_factory)
self.actionfile = actionfile
if self.actionfile:
self.bend_coefs = tps.loglinspace(self.reg_init, self.reg_final, self.n_iter)
self.src_ctx = GPUContext(self.bend_coefs)
self.src_ctx.read_h5(actionfile)
self.warn_clip_cloud = True
class LandmarkFeats(MulFeats):
def __init__(self, actionfile):
MulFeats.__init__(self, actionfile)
self.landmark_ctx = None
def set_landmark_file(self, landmarkf):
self.landmark_ctx = GPUContext()
self.landmark_ctx.read_h5(landmarkf)
self.landmark_targ_ctx = TgtContext(self.landmark_ctx)
self.weights = np.zeros(self.src_ctx.N + self.landmark_ctx.N +
MulFeats.N_costs)
def features(self, state, **kwargs):
mul_feats = MulFeats.features(self, state)
self.landmark_targ_ctx.set_cld(state.cloud)
landmark_feats = batch_tps_rpm_bij(self.landmark_ctx,
self.landmark_targ_ctx)
landmark_feats = np.exp(-landmark_feats)
landmark_feats /= np.sum(landmark_feats)
self.costs = np.c_[mul_feats,
np.tile(landmark_feats, (self.src_ctx.N, 1))]
return self.costs
@staticmethod
def get_size(num_actions, num_landmarks=70):
return num_actions + num_landmarks + MulFeats.N_costs
class BatchGpuTpsRpmBijRegistrationFactory(TpsRpmBijRegistrationFactory):
"""
Similar to TpsRpmBijRegistrationFactory but batch_register and batch_cost are computed in batch using the GPU
"""
def __init__(self, demos, actionfile=None,
n_iter=settings.N_ITER, em_iter=settings.EM_ITER,
reg_init=settings.REG[0], reg_final=settings.REG[1],
rad_init=settings.RAD[0], rad_final=settings.RAD[1],
rot_reg=settings.ROT_REG,
outlierprior=settings.OUTLIER_PRIOR, outlierfrac=settings.OURLIER_FRAC,
prior_fn=None,
f_solver_factory=solver.AutoTpsSolverFactory(),
g_solver_factory=solver.AutoTpsSolverFactory(use_cache=False)):
if not lfd.registration._has_cuda:
raise NotImplementedError("CUDA not installed")
super(BatchGpuTpsRpmBijRegistrationFactory, self).__init__(demos=demos,
n_iter=n_iter, em_iter=em_iter,
reg_init=reg_init, reg_final=reg_final,
rad_init=rad_init, rad_final=rad_final,
rot_reg=rot_reg,
outlierprior=outlierprior, outlierfrac=outlierfrac,
prior_fn=prior_fn,
f_solver_factory=f_solver_factory, g_solver_factory=g_solver_factory)
self.actionfile = actionfile
if self.actionfile:
self.bend_coefs = tps.loglinspace(self.reg_init, self.reg_final, self.n_iter)
self.src_ctx = GPUContext(self.bend_coefs)
self.src_ctx.read_h5(actionfile)
self.warn_clip_cloud = True
def _clip_cloud(self, cloud):
if len(cloud) > settings.MAX_CLD_SIZE:
cloud = cloud[np.random.choice(range(len(cloud)), size=settings.MAX_CLD_SIZE, replace=False)]
if self.warn_clip_cloud:
import warnings
warnings.warn("The cloud has more points than the maximum for GPU and it is being clipped")
self.warn_clip_cloud = False
return cloud
def batch_register(self, test_scene_state):
raise NotImplementedError
def batch_cost(self, test_scene_state):
if not(self.actionfile):
raise ValueError('No actionfile provided for gpu context')
tgt_ctx = TgtContext(self.src_ctx)
cloud = test_scene_state.cloud
cloud = self._clip_cloud(cloud)
tgt_ctx.set_cld(cloud)
cost_array = batch_tps_rpm_bij(self.src_ctx, tgt_ctx,
T_init=self.rad_init, T_final=self.rad_final,
outlierfrac=self.outlierfrac, outlierprior=self.outlierprior,
outliercutoff=settings.OUTLIER_CUTOFF,
em_iter=self.em_iter,
component_cost=True)
costs = dict(zip(self.src_ctx.seg_names, cost_array))
return costs
class LandmarkFeats(MulFeats):
def __init__(self, actionfile):
MulFeats.__init__(self, actionfile)
self.landmark_ctx = None
def set_landmark_file(self, landmarkf):
self.landmark_ctx = GPUContext()
self.landmark_ctx.read_h5(landmarkf)
self.landmark_targ_ctx = TgtContext(self.landmark_ctx)
self.weights = np.zeros(self.src_ctx.N + self.landmark_ctx.N + MulFeats.N_costs)
def features(self, state, **kwargs):
mul_feats = MulFeats.features(self, state)
self.landmark_targ_ctx.set_cld(state.cloud)
landmark_feats = batch_tps_rpm_bij(self.landmark_ctx, self.landmark_targ_ctx)
landmark_feats = np.exp(-landmark_feats)
landmark_feats /= np.sum(landmark_feats)
self.costs = np.c_[mul_feats, np.tile(landmark_feats, (self.src_ctx.N, 1))]
return self.costs
@staticmethod
def get_size(num_actions, num_landmarks=70):
return num_actions + num_landmarks + MulFeats.N_costs
def set_landmark_file(self, landmarkf):
self.landmark_ctx = GPUContext()
self.landmark_ctx.read_h5(landmarkf)
self.landmark_targ_ctx = TgtContext(self.landmark_ctx)
self.weights = np.zeros(self.src_ctx.N + self.landmark_ctx.N +
MulFeats.N_costs)
def set_landmark_file(self, landmarkf):
self.landmark_ctx = GPUContext()
self.landmark_ctx.read_h5(landmarkf)
self.landmark_targ_ctx = TgtContext(self.landmark_ctx)
self.weights = np.zeros(self.src_ctx.N + self.landmark_ctx.N + MulFeats.N_costs)
