(self, lmdb_dir, n_poses, subdivision_stride, pose_resampling_fps, mean_pose, mean_dir_vec,
speaker_model=None, remove_word_timing=False)
| 57 | |
| 58 | class SpeechMotionDataset(Dataset): |
| 59 | def __init__(self, lmdb_dir, n_poses, subdivision_stride, pose_resampling_fps, mean_pose, mean_dir_vec, |
| 60 | speaker_model=None, remove_word_timing=False): |
| 61 | |
| 62 | self.lmdb_dir = lmdb_dir |
| 63 | self.n_poses = n_poses |
| 64 | self.subdivision_stride = subdivision_stride |
| 65 | self.skeleton_resampling_fps = pose_resampling_fps |
| 66 | self.mean_dir_vec = mean_dir_vec |
| 67 | self.remove_word_timing = remove_word_timing |
| 68 | |
| 69 | self.expected_audio_length = int(round(n_poses / pose_resampling_fps * 16000)) |
| 70 | self.expected_spectrogram_length = utils.data_utils.calc_spectrogram_length_from_motion_length( |
| 71 | n_poses, pose_resampling_fps) |
| 72 | |
| 73 | self.lang_model = None |
| 74 | |
| 75 | logging.info("Reading data '{}'...".format(lmdb_dir)) |
| 76 | preloaded_dir = lmdb_dir + '_cache' |
| 77 | if not os.path.exists(preloaded_dir): |
| 78 | logging.info('Creating the dataset cache...') |
| 79 | assert mean_dir_vec is not None |
| 80 | if mean_dir_vec.shape[-1] != 3: |
| 81 | mean_dir_vec = mean_dir_vec.reshape(mean_dir_vec.shape[:-1] + (-1, 3)) |
| 82 | n_poses_extended = int(round(n_poses * 1.25)) # some margin |
| 83 | data_sampler = DataPreprocessor(lmdb_dir, preloaded_dir, n_poses_extended, |
| 84 | subdivision_stride, pose_resampling_fps, mean_pose, mean_dir_vec) |
| 85 | data_sampler.run() |
| 86 | else: |
| 87 | logging.info('Found the cache {}'.format(preloaded_dir)) |
| 88 | |
| 89 | # init lmdb |
| 90 | self.lmdb_env = lmdb.open(preloaded_dir, readonly=True, lock=False) |
| 91 | with self.lmdb_env.begin() as txn: |
| 92 | self.n_samples = txn.stat()['entries'] |
| 93 | |
| 94 | # make a speaker model |
| 95 | if speaker_model is None or speaker_model == 0: |
| 96 | precomputed_model = lmdb_dir + '_speaker_model.pkl' |
| 97 | if not os.path.exists(precomputed_model): |
| 98 | self._make_speaker_model(lmdb_dir, precomputed_model) |
| 99 | else: |
| 100 | with open(precomputed_model, 'rb') as f: |
| 101 | self.speaker_model = pickle.load(f) |
| 102 | else: |
| 103 | self.speaker_model = speaker_model |
| 104 | |
| 105 | def __len__(self): |
| 106 | return self.n_samples |
nothing calls this directly
no test coverage detected