Dictionary keeping track of the best fit per image in the training set. Ref: https://github.com/nkolot/SPIN/blob/master/train/fits_dict.py
| 28 | |
| 29 | |
| 30 | class FitsDict(): |
| 31 | """Dictionary keeping track of the best fit per image in the training set. |
| 32 | |
| 33 | Ref: https://github.com/nkolot/SPIN/blob/master/train/fits_dict.py |
| 34 | """ |
| 35 | def __init__(self, fits='static') -> None: |
| 36 | assert fits in ['static', 'final'] |
| 37 | self.fits = fits |
| 38 | self.fits_dict = {} |
| 39 | |
| 40 | # array used to flip SMPL pose parameters |
| 41 | self.flipped_parts = torch.tensor(SMPL_POSE_FLIP_PERM, |
| 42 | dtype=torch.int64) |
| 43 | # Load dictionary state |
| 44 | # for ds_name, ds in train_dataset.dataset_dict.items(): |
| 45 | for ds_name in train_datasets: |
| 46 | |
| 47 | # h36m has gt so no static fits |
| 48 | if ds_name == 'h36m' or self.fits == 'static': |
| 49 | dict_file = os.path.join(static_fits_load_dir, |
| 50 | ds_name + '_fits.npy') |
| 51 | content = np.load(dict_file) |
| 52 | self.fits_dict[ds_name] = torch.from_numpy(content) |
| 53 | del content |
| 54 | elif self.fits == 'final': |
| 55 | dict_file = os.path.join('data/final_fits', ds_name + '.npz') |
| 56 | # load like this to save mem |
| 57 | content = np.load(dict_file) |
| 58 | pose = torch.from_numpy(content['pose']) |
| 59 | betas = torch.from_numpy(content['betas']) |
| 60 | del content |
| 61 | params = torch.cat([pose, betas], dim=-1) |
| 62 | self.fits_dict[ds_name] = params |
| 63 | |
| 64 | def save(self): |
| 65 | """Save dictionary state to disk.""" |
| 66 | for ds_name in train_datasets: |
| 67 | dict_file = os.path.join(save_dir, ds_name + '_fits.npy') |
| 68 | np.save(dict_file, self.fits_dict[ds_name].cpu().numpy()) |
| 69 | |
| 70 | def __getitem__(self, x): |
| 71 | """Retrieve dictionary entries.""" |
| 72 | dataset_name, ind, rot, is_flipped = x |
| 73 | batch_size = len(dataset_name) |
| 74 | pose = torch.zeros((batch_size, 72)) |
| 75 | betas = torch.zeros((batch_size, 10)) |
| 76 | for ds, i, n in zip(dataset_name, ind, range(batch_size)): |
| 77 | params = self.fits_dict[ds][i] |
| 78 | pose[n, :] = params[:72] |
| 79 | betas[n, :] = params[72:] |
| 80 | pose = pose.clone() |
| 81 | |
| 82 | # Apply flipping and rotation |
| 83 | pose = self.rotate_pose(self.flip_pose(pose, is_flipped), rot) |
| 84 | |
| 85 | betas = betas.clone() |
| 86 | return pose, betas |
| 87 |