| 289 | logging.basicConfig(level=logging.ERROR) |
| 290 | |
| 291 | def get_embedding(speech_array, wav2vec_feature_extractor, audio_encoder, sr=16000, device='cpu'): |
| 292 | audio_duration = len(speech_array) / sr |
| 293 | video_length = audio_duration * 25 # Assume the video fps is 25 |
| 294 | |
| 295 | # wav2vec_feature_extractor |
| 296 | audio_feature = np.squeeze( |
| 297 | wav2vec_feature_extractor(speech_array, sampling_rate=sr).input_values |
| 298 | ) |
| 299 | audio_feature = torch.from_numpy(audio_feature).float().to(device=device) |
| 300 | audio_feature = audio_feature.unsqueeze(0) |
| 301 | |
| 302 | # audio encoder |
| 303 | with torch.no_grad(): |
| 304 | embeddings = audio_encoder(audio_feature, seq_len=int(video_length), output_hidden_states=True) |
| 305 | |
| 306 | if len(embeddings) == 0: |
| 307 | print("Fail to extract audio embedding") |
| 308 | return None |
| 309 | |
| 310 | audio_emb = torch.stack(embeddings.hidden_states[1:], dim=1).squeeze(0) |
| 311 | audio_emb = rearrange(audio_emb, "b s d -> s b d") |
| 312 | |
| 313 | audio_emb = audio_emb.cpu().detach() |
| 314 | return audio_emb |
| 315 | |
| 316 | def extract_audio_from_video(filename, sample_rate): |
| 317 | raw_audio_path = filename.split('/')[-1].split('.')[0]+'.wav' |