Method sample

ADHMR/lib/runners/scorenet.py:296–507

(self, state, construct_dpo=False)

    

    def sample(self, state, construct_dpo=False):
        args, config = self.args, self.config
        
        model = state['model']
        model_cond = state['model_cond']
        model_score = state['model_score']
        model_score_cond = state['model_score_cond']
        dataset = state['dataset']
        dataloader = state['dataloader']
        epoch = state['epoch']
        dataset_type = state['dataset_type']
            
        rank = torch.distributed.get_rank()
        
        multi_n = self.multihypo_n
        with torch.no_grad():
            pred = {}
            dpo_pair_list = []
            for i in range(multi_n+2):
                pred[i] = {}
            for i, (inps, labels, img_ids, bboxes, img_path) in enumerate(tqdm(dataloader, ncols=100)):
                n = inps.size(0)
                output = {}
                for k, _ in labels.items():
                    labels[k] = labels[k].to(self.device)
                input = inps.float().to(self.device)
                save_trans_inv = labels['trans_inv'].cpu().numpy().astype(np.float32)
                scale = torch.tensor([self.image_size[0],self.image_size[1],dataset.bbox_3d_shape[2]]).float().to(self.device) / self.config.diffusion.scale
                labels['trans_inv']= labels['trans_inv'].unsqueeze(1).repeat(1,multi_n+2,1,1).view(-1,2,3)
                labels['intrinsic_param']= labels['intrinsic_param'].unsqueeze(1).repeat(1,multi_n+2,1,1).view(-1,3,3)
                labels['joint_root']= labels['joint_root'].unsqueeze(1).repeat(1,multi_n+2,1).view(-1,3)

                state['input'] = input.clone()
                state['scale'] = scale.clone()
                output, score_input = self.gen_mesh(state, multi_n)                     # get denoised twist and joints

                if construct_dpo:
                    # construct with GT as winner
                    save_trans_inv = save_trans_inv.tolist()
                    save_pred_joints = output['pred_joints'].cpu().numpy().astype(np.float32).tolist()
                    save_pred_twist = output['pred_twist'].cpu().numpy().astype(np.float32).tolist()
                    save_score_joint = score_input['joint'].cpu().numpy().astype(np.float32).tolist()
                    save_score_twist = score_input['twist'].cpu().numpy().astype(np.float32).tolist()
                    save_joints_uvd_29 = labels['joints_uvd_29'].cpu().numpy().astype(np.float32).tolist()
                    save_w_twist = labels['target_twist'].cpu().numpy().astype(np.float32).tolist()
                    save_trans = labels['trans'].cpu().numpy().astype(np.float32).tolist()
                    # import ipdb; ipdb.set_trace()
                    for idx in range(n):
                        for idx_hypo in range(multi_n):
                            dpo_pair = {
                                'img_path': img_path[idx],
                                # 'db_idx': labels['db_idx'][idx].item(),
                                'img_idx': img_ids[idx].item(),                     # absolute img_ids 
                                'l_joints': save_pred_joints[idx][idx_hypo],
                                'l_twist': save_pred_twist[idx][idx_hypo],
                                'w_joints': save_joints_uvd_29[idx],
                                'w_twist': save_w_twist[idx],
                                'trans': save_trans[idx],