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