Spaces:
Paused
Paused
| import torch | |
| from tqdm import tqdm | |
| from StructDiffusion.diffusion.noise_schedule import extract | |
| from StructDiffusion.diffusion.pose_conversion import get_struct_objs_poses | |
| from StructDiffusion.utils.batch_inference import move_pc_and_create_scene_new | |
| import StructDiffusion.utils.tra3d as tra3d | |
| class Sampler: | |
| def __init__(self, model_class, checkpoint_path, device, debug=False): | |
| self.debug = debug | |
| self.device = device | |
| self.model = model_class.load_from_checkpoint(checkpoint_path) | |
| self.backbone = self.model.model | |
| self.backbone.to(device) | |
| self.backbone.eval() | |
| def sample(self, batch, num_poses): | |
| noise_schedule = self.model.noise_schedule | |
| B = batch["pcs"].shape[0] | |
| x_noisy = torch.randn((B, num_poses, 9), device=self.device) | |
| xs = [] | |
| for t_index in tqdm(reversed(range(0, noise_schedule.timesteps)), | |
| desc='sampling loop time step', total=noise_schedule.timesteps): | |
| t = torch.full((B,), t_index, device=self.device, dtype=torch.long) | |
| # noise schedule | |
| betas_t = extract(noise_schedule.betas, t, x_noisy.shape) | |
| sqrt_one_minus_alphas_cumprod_t = extract(noise_schedule.sqrt_one_minus_alphas_cumprod, t, x_noisy.shape) | |
| sqrt_recip_alphas_t = extract(noise_schedule.sqrt_recip_alphas, t, x_noisy.shape) | |
| # predict noise | |
| pcs = batch["pcs"] | |
| sentence = batch["sentence"] | |
| type_index = batch["type_index"] | |
| position_index = batch["position_index"] | |
| pad_mask = batch["pad_mask"] | |
| # calling the backbone instead of the pytorch-lightning model | |
| with torch.no_grad(): | |
| predicted_noise = self.backbone.forward(t, pcs, sentence, x_noisy, type_index, position_index, pad_mask) | |
| # compute noisy x at t | |
| model_mean = sqrt_recip_alphas_t * (x_noisy - betas_t * predicted_noise / sqrt_one_minus_alphas_cumprod_t) | |
| if t_index == 0: | |
| x_noisy = model_mean | |
| else: | |
| posterior_variance_t = extract(noise_schedule.posterior_variance, t, x_noisy.shape) | |
| noise = torch.randn_like(x_noisy) | |
| x_noisy = model_mean + torch.sqrt(posterior_variance_t) * noise | |
| xs.append(x_noisy) | |
| xs = list(reversed(xs)) | |
| return xs | |
| class SamplerV2: | |
| def __init__(self, diffusion_model_class, diffusion_checkpoint_path, | |
| collision_model_class, collision_checkpoint_path, | |
| device, debug=False): | |
| self.debug = debug | |
| self.device = device | |
| self.diffusion_model = diffusion_model_class.load_from_checkpoint(diffusion_checkpoint_path) | |
| self.diffusion_backbone = self.diffusion_model.model | |
| self.diffusion_backbone.to(device) | |
| self.diffusion_backbone.eval() | |
| self.collision_model = collision_model_class.load_from_checkpoint(collision_checkpoint_path) | |
| self.collision_backbone = self.collision_model.model | |
| self.collision_backbone.to(device) | |
| self.collision_backbone.eval() | |
| def sample(self, batch, num_poses, num_elite, discriminator_batch_size): | |
| noise_schedule = self.diffusion_model.noise_schedule | |
| B = batch["pcs"].shape[0] | |
| x_noisy = torch.randn((B, num_poses, 9), device=self.device) | |
| xs = [] | |
| for t_index in tqdm(reversed(range(0, noise_schedule.timesteps)), | |
| desc='sampling loop time step', total=noise_schedule.timesteps): | |
| t = torch.full((B,), t_index, device=self.device, dtype=torch.long) | |
| # noise schedule | |
| betas_t = extract(noise_schedule.betas, t, x_noisy.shape) | |
| sqrt_one_minus_alphas_cumprod_t = extract(noise_schedule.sqrt_one_minus_alphas_cumprod, t, x_noisy.shape) | |
| sqrt_recip_alphas_t = extract(noise_schedule.sqrt_recip_alphas, t, x_noisy.shape) | |
| # predict noise | |
| pcs = batch["pcs"] | |
| sentence = batch["sentence"] | |
| type_index = batch["type_index"] | |
| position_index = batch["position_index"] | |
| pad_mask = batch["pad_mask"] | |
| # calling the backbone instead of the pytorch-lightning model | |
| with torch.no_grad(): | |
| predicted_noise = self.diffusion_backbone.forward(t, pcs, sentence, x_noisy, type_index, position_index, pad_mask) | |
| # compute noisy x at t | |
| model_mean = sqrt_recip_alphas_t * (x_noisy - betas_t * predicted_noise / sqrt_one_minus_alphas_cumprod_t) | |
| if t_index == 0: | |
| x_noisy = model_mean | |
| else: | |
| posterior_variance_t = extract(noise_schedule.posterior_variance, t, x_noisy.shape) | |
| noise = torch.randn_like(x_noisy) | |
| x_noisy = model_mean + torch.sqrt(posterior_variance_t) * noise | |
| xs.append(x_noisy) | |
| xs = list(reversed(xs)) | |
| visualize = True | |
| struct_pose, pc_poses_in_struct = get_struct_objs_poses(xs[0]) | |
| # struct_pose: B, 1, 4, 4 | |
| # pc_poses_in_struct: B, N, 4, 4 | |
| S = B | |
| B_discriminator = discriminator_batch_size | |
| #################################################### | |
| # only keep one copy | |
| # N, P, 3 | |
| obj_xyzs = batch["pcs"][0][:, :, :3] | |
| print("obj_xyzs shape", obj_xyzs.shape) | |
| # 1, N | |
| # object_pad_mask: padding location has 1 | |
| num_target_objs = num_poses | |
| if self.diffusion_backbone.use_virtual_structure_frame: | |
| num_target_objs -= 1 | |
| object_pad_mask = batch["pad_mask"][0][-num_target_objs:].unsqueeze(0) | |
| target_object_inds = 1 - object_pad_mask | |
| print("target_object_inds shape", target_object_inds.shape) | |
| print("target_object_inds", target_object_inds) | |
| N, P, _ = obj_xyzs.shape | |
| print("S, N, P: {}, {}, {}".format(S, N, P)) | |
| #################################################### | |
| # S, N, ... | |
| struct_pose = struct_pose.repeat(1, N, 1, 1) # S, N, 4, 4 | |
| struct_pose = struct_pose.reshape(S * N, 4, 4) # S x N, 4, 4 | |
| new_obj_xyzs = obj_xyzs.repeat(S, 1, 1, 1) # S, N, P, 3 | |
| current_pc_pose = torch.eye(4).repeat(S, N, 1, 1).to(self.device) # S, N, 4, 4 | |
| current_pc_pose[:, :, :3, 3] = torch.mean(new_obj_xyzs, dim=2) # S, N, 4, 4 | |
| current_pc_pose = current_pc_pose.reshape(S * N, 4, 4) # S x N, 4, 4 | |
| # optimize xyzrpy | |
| obj_params = torch.zeros((S, N, 6)).to(self.device) | |
| obj_params[:, :, :3] = pc_poses_in_struct[:, :, :3, 3] | |
| obj_params[:, :, 3:] = tra3d.matrix_to_euler_angles(pc_poses_in_struct[:, :, :3, :3], "XYZ") # S, N, 6 | |
| # | |
| # new_obj_xyzs_before_cem, goal_pc_pose_before_cem = move_pc(obj_xyzs, obj_params, struct_pose, current_pc_pose, device) | |
| # | |
| # if visualize: | |
| # print("visualizing rearrangements predicted by the generator") | |
| # visualize_batch_pcs(new_obj_xyzs_before_cem, S, N, P, limit_B=5) | |
| #################################################### | |
| # rank | |
| # evaluate in batches | |
| scores = torch.zeros(S).to(self.device) | |
| no_intersection_scores = torch.zeros(S).to(self.device) # the higher the better | |
| num_batches = int(S / B_discriminator) | |
| if S % B_discriminator != 0: | |
| num_batches += 1 | |
| for b in range(num_batches): | |
| if b + 1 == num_batches: | |
| cur_batch_idxs_start = b * B_discriminator | |
| cur_batch_idxs_end = S | |
| else: | |
| cur_batch_idxs_start = b * B_discriminator | |
| cur_batch_idxs_end = (b + 1) * B_discriminator | |
| cur_batch_size = cur_batch_idxs_end - cur_batch_idxs_start | |
| # print("current batch idxs start", cur_batch_idxs_start) | |
| # print("current batch idxs end", cur_batch_idxs_end) | |
| # print("size of the current batch", cur_batch_size) | |
| batch_obj_params = obj_params[cur_batch_idxs_start: cur_batch_idxs_end] | |
| batch_struct_pose = struct_pose[cur_batch_idxs_start * N: cur_batch_idxs_end * N] | |
| batch_current_pc_pose = current_pc_pose[cur_batch_idxs_start * N:cur_batch_idxs_end * N] | |
| new_obj_xyzs, _, subsampled_scene_xyz, _, obj_pair_xyzs = \ | |
| move_pc_and_create_scene_new(obj_xyzs, batch_obj_params, batch_struct_pose, batch_current_pc_pose, | |
| target_object_inds, self.device, | |
| return_scene_pts=False, | |
| return_scene_pts_and_pc_idxs=False, | |
| num_scene_pts=False, | |
| normalize_pc=False, | |
| return_pair_pc=True, | |
| num_pair_pc_pts=self.collision_model.data_cfg.num_scene_pts, | |
| normalize_pair_pc=self.collision_model.data_cfg.normalize_pc) | |
| ####################################### | |
| # predict whether there are pairwise collisions | |
| # if collision_score_weight > 0: | |
| with torch.no_grad(): | |
| _, num_comb, num_pair_pc_pts, _ = obj_pair_xyzs.shape | |
| # obj_pair_xyzs = obj_pair_xyzs.reshape(cur_batch_size * num_comb, num_pair_pc_pts, -1) | |
| collision_logits = self.collision_backbone.forward(obj_pair_xyzs.reshape(cur_batch_size * num_comb, num_pair_pc_pts, -1)) | |
| collision_scores = self.collision_backbone.convert_logits(collision_logits).reshape(cur_batch_size, num_comb) # cur_batch_size, num_comb | |
| # debug | |
| # for bi, this_obj_pair_xyzs in enumerate(obj_pair_xyzs): | |
| # print("batch id", bi) | |
| # for pi, obj_pair_xyz in enumerate(this_obj_pair_xyzs): | |
| # print("pair", pi) | |
| # # obj_pair_xyzs: 2 * P, 5 | |
| # print("collision score", collision_scores[bi, pi]) | |
| # trimesh.PointCloud(obj_pair_xyz[:, :3].cpu()).show() | |
| # 1 - mean() since the collision model predicts 1 if there is a collision | |
| no_intersection_scores[cur_batch_idxs_start:cur_batch_idxs_end] = 1 - torch.mean(collision_scores, dim=1) | |
| if visualize: | |
| print("no intersection scores", no_intersection_scores) | |
| # ####################################### | |
| # if discriminator_score_weight > 0: | |
| # # # debug: | |
| # # print(subsampled_scene_xyz.shape) | |
| # # print(subsampled_scene_xyz[0]) | |
| # # trimesh.PointCloud(subsampled_scene_xyz[0, :, :3].cpu().numpy()).show() | |
| # # | |
| # with torch.no_grad(): | |
| # | |
| # # Important: since this discriminator only uses local structure param, takes sentence from the first and last position | |
| # # local_sentence = sentence[:, [0, 4]] | |
| # # local_sentence_pad_mask = sentence_pad_mask[:, [0, 4]] | |
| # # sentence_disc, sentence_pad_mask_disc, position_index_dic = discriminator_inference.dataset.tensorfy_sentence(raw_sentence_discriminator, raw_sentence_pad_mask_discriminator, raw_position_index_discriminator) | |
| # | |
| # sentence_disc = torch.LongTensor( | |
| # [discriminator_tokenizer.tokenize(*i) for i in raw_sentence_discriminator]) | |
| # sentence_pad_mask_disc = torch.LongTensor(raw_sentence_pad_mask_discriminator) | |
| # position_index_dic = torch.LongTensor(raw_position_index_discriminator) | |
| # | |
| # preds = discriminator_model.forward(subsampled_scene_xyz, | |
| # sentence_disc.unsqueeze(0).repeat(cur_batch_size, 1).to(device), | |
| # sentence_pad_mask_disc.unsqueeze(0).repeat(cur_batch_size, | |
| # 1).to(device), | |
| # position_index_dic.unsqueeze(0).repeat(cur_batch_size, 1).to( | |
| # device)) | |
| # # preds = discriminator_model.forward(subsampled_scene_xyz) | |
| # preds = discriminator_model.convert_logits(preds) | |
| # preds = preds["is_circle"] # cur_batch_size, | |
| # scores[cur_batch_idxs_start:cur_batch_idxs_end] = preds | |
| # if visualize: | |
| # print("discriminator scores", scores) | |
| # scores = scores * discriminator_score_weight + no_intersection_scores * collision_score_weight | |
| scores = no_intersection_scores | |
| sort_idx = torch.argsort(scores).flip(dims=[0])[:num_elite] | |
| elite_obj_params = obj_params[sort_idx] # num_elite, N, 6 | |
| elite_struct_poses = struct_pose.reshape(S, N, 4, 4)[sort_idx] # num_elite, N, 4, 4 | |
| elite_struct_poses = elite_struct_poses.reshape(num_elite * N, 4, 4) # num_elite x N, 4, 4 | |
| elite_scores = scores[sort_idx] | |
| print("elite scores:", elite_scores) | |
| #################################################### | |
| # # visualize best samples | |
| # num_scene_pts = 4096 # if discriminator_num_scene_pts is None else discriminator_num_scene_pts | |
| # batch_current_pc_pose = current_pc_pose[0: num_elite * N] | |
| # best_new_obj_xyzs, best_goal_pc_pose, best_subsampled_scene_xyz, _, _ = \ | |
| # move_pc_and_create_scene_new(obj_xyzs, elite_obj_params, elite_struct_poses, batch_current_pc_pose, | |
| # target_object_inds, self.device, | |
| # return_scene_pts=True, num_scene_pts=num_scene_pts, normalize_pc=True) | |
| # if visualize: | |
| # print("visualizing elite rearrangements ranked by collision model/discriminator") | |
| # visualize_batch_pcs(best_new_obj_xyzs, num_elite, limit_B=num_elite) | |
| # num_elite, N, 6 | |
| elite_obj_params = elite_obj_params.reshape(num_elite * N, -1) | |
| pc_poses_in_struct = torch.eye(4).repeat(num_elite * N, 1, 1).to(self.device) | |
| pc_poses_in_struct[:, :3, :3] = tra3d.euler_angles_to_matrix(elite_obj_params[:, 3:], "XYZ") | |
| pc_poses_in_struct[:, :3, 3] = elite_obj_params[:, :3] | |
| pc_poses_in_struct = pc_poses_in_struct.reshape(num_elite, N, 4, 4) # num_elite, N, 4, 4 | |
| struct_pose = elite_struct_poses.reshape(num_elite, N, 4, 4)[:, 0,].unsqueeze(1) # num_elite, 1, 4, 4 | |
| print(struct_pose.shape) | |
| print(pc_poses_in_struct.shape) | |
| return struct_pose, pc_poses_in_struct |