| 321 | logging.basicConfig(level=logging.ERROR) |
| 322 | |
| 323 | def get_embedding(speech_array, wav2vec_feature_extractor, audio_encoder, sr=16000, device='cpu'): |
| 324 | audio_duration = len(speech_array) / sr |
| 325 | video_length = audio_duration * 25 # Assume the video fps is 25 |
| 326 | |
| 327 | # wav2vec_feature_extractor |
| 328 | audio_feature = np.squeeze( |
| 329 | wav2vec_feature_extractor(speech_array, sampling_rate=sr).input_values |
| 330 | ) |
| 331 | audio_feature = torch.from_numpy(audio_feature).float().to(device=device) |
| 332 | audio_feature = audio_feature.unsqueeze(0) |
| 333 | |
| 334 | # audio encoder |
| 335 | with torch.no_grad(): |
| 336 | embeddings = audio_encoder(audio_feature, seq_len=int(video_length), output_hidden_states=True) |
| 337 | |
| 338 | if len(embeddings) == 0: |
| 339 | print("Fail to extract audio embedding") |
| 340 | return None |
| 341 | |
| 342 | audio_emb = torch.stack(embeddings.hidden_states[1:], dim=1).squeeze(0) |
| 343 | audio_emb = rearrange(audio_emb, "b s d -> s b d") |
| 344 | |
| 345 | audio_emb = audio_emb.cpu().detach() |
| 346 | return audio_emb |
| 347 | |
| 348 | def extract_audio_from_video(filename, sample_rate): |
| 349 | raw_audio_path = filename.split('/')[-1].split('.')[0]+'.wav' |