def prepare_pc(self, pc):
if pc.shape[0] > self.target_pc_size:
pc = utils.regularize_pc_point_count(pc, self.target_pc_size)
pc_mean = np.mean(pc, 0)
pc -= np.expand_dims(pc_mean, 0)
pc = np.tile(pc, (self.num_grasp_samples, 1, 1))
pc = torch.from_numpy(pc).float().to(self.device)
pcs = []
pcs = utils.partition_array_into_subarrays(pc, self.batch_size)
return pcs, pc_mean
def generate_grasps(self, pcs):
grasps_list = []
confidence_list = []
z_list = []
if self.generate_dense_grasps:
latent_samples = self.grasp_sampler.net.module.generate_dense_latents(
self.num_grasps_per_dim)
latent_samples = utils.partition_array_into_subarrays(
latent_samples, self.batch_size)
for latent_sample, pc in zip(latent_samples, pcs):
grasps, confidence, z = self.grasp_sampler.generate_grasps(
pc, latent_sample)
grasps_list.append(grasps)
confidence_list.append(confidence)
z_list.append(z)
else:
for pc in pcs:
grasps, confidence, z = self.grasp_sampler.generate_grasps(pc)
grasps_list.append(grasps)
confidence_list.append(confidence)
z_list.append(z)
return grasps_list, confidence_list, z_list
def generate_grasps(self, pcs):
all_grasps = []
all_confidence = []
all_z = []
if self.generate_dense_grasps:
latent_samples = self.grasp_sampler.net.module.generate_dense_latents(
self.num_grasps_per_dim)
latent_samples = utils.partition_array_into_subarrays(
latent_samples, self.batch_size)
for latent_sample, pc in zip(latent_samples, pcs):
grasps, confidence, z = self.grasp_sampler.generate_grasps(
pc, latent_sample)
all_grasps.append(grasps)
all_confidence.append(confidence)
all_z.append(z)
else:
for pc in pcs:
grasps, confidence, z = self.grasp_sampler.generate_grasps(pc)
all_grasps.append(grasps)
all_confidence.append(confidence)
all_z.append(z)
return all_grasps, all_confidence, all_z