|
|
| import math
|
| import numpy as np
|
| import torch
|
| import torch.nn.functional as F
|
| from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
|
|
|
|
|
| def get_sample_indices(original_fps,
|
| total_frames,
|
| target_fps,
|
| num_sample,
|
| fixed_start=None):
|
| required_duration = num_sample / target_fps
|
| required_origin_frames = int(np.ceil(required_duration * original_fps))
|
| if required_duration > total_frames / original_fps:
|
| raise ValueError("required_duration must be less than video length")
|
|
|
| if not fixed_start is None and fixed_start >= 0:
|
| start_frame = fixed_start
|
| else:
|
| max_start = total_frames - required_origin_frames
|
| if max_start < 0:
|
| raise ValueError("video length is too short")
|
| start_frame = np.random.randint(0, max_start + 1)
|
| start_time = start_frame / original_fps
|
|
|
| end_time = start_time + required_duration
|
| time_points = np.linspace(start_time, end_time, num_sample, endpoint=False)
|
|
|
| frame_indices = np.round(np.array(time_points) * original_fps).astype(int)
|
| frame_indices = np.clip(frame_indices, 0, total_frames - 1)
|
| return frame_indices
|
|
|
|
|
| def linear_interpolation(features, input_fps, output_fps, output_len=None):
|
| """
|
| features: shape=[1, T, 512]
|
| input_fps: fps for audio, f_a
|
| output_fps: fps for video, f_m
|
| output_len: video length
|
| """
|
| features = features.transpose(1, 2)
|
| seq_len = features.shape[2] / float(input_fps)
|
| if output_len is None:
|
| output_len = int(seq_len * output_fps)
|
| output_features = F.interpolate(
|
| features, size=output_len, align_corners=True,
|
| mode='linear')
|
| return output_features.transpose(1, 2)
|
|
|
|
|
| class AudioEncoder():
|
|
|
| def __init__(self, device='cpu', model_id="facebook/wav2vec2-base-960h"):
|
|
|
| self.processor = Wav2Vec2Processor.from_pretrained(model_id)
|
| self.model = Wav2Vec2ForCTC.from_pretrained(model_id)
|
|
|
| self.model = self.model.to(device)
|
|
|
| self.video_rate = 30
|
|
|
| def get_audio_embed_bucket(self,
|
| audio_embed,
|
| stride=2,
|
| batch_frames=12,
|
| m=2):
|
| num_layers, audio_frame_num, audio_dim = audio_embed.shape
|
|
|
| if num_layers > 1:
|
| return_all_layers = True
|
| else:
|
| return_all_layers = False
|
|
|
| min_batch_num = int(audio_frame_num / (batch_frames * stride)) + 1
|
|
|
| bucket_num = min_batch_num * batch_frames
|
| batch_idx = [stride * i for i in range(bucket_num)]
|
| batch_audio_eb = []
|
| for bi in batch_idx:
|
| if bi < audio_frame_num:
|
| audio_sample_stride = 2
|
| chosen_idx = list(
|
| range(bi - m * audio_sample_stride,
|
| bi + (m + 1) * audio_sample_stride,
|
| audio_sample_stride))
|
| chosen_idx = [0 if c < 0 else c for c in chosen_idx]
|
| chosen_idx = [
|
| audio_frame_num - 1 if c >= audio_frame_num else c
|
| for c in chosen_idx
|
| ]
|
|
|
| if return_all_layers:
|
| frame_audio_embed = audio_embed[:, chosen_idx].flatten(
|
| start_dim=-2, end_dim=-1)
|
| else:
|
| frame_audio_embed = audio_embed[0][chosen_idx].flatten()
|
| else:
|
| frame_audio_embed = \
|
| torch.zeros([audio_dim * (2 * m + 1)], device=audio_embed.device) if not return_all_layers \
|
| else torch.zeros([num_layers, audio_dim * (2 * m + 1)], device=audio_embed.device)
|
| batch_audio_eb.append(frame_audio_embed)
|
| batch_audio_eb = torch.cat([c.unsqueeze(0) for c in batch_audio_eb],
|
| dim=0)
|
|
|
| return batch_audio_eb, min_batch_num
|
|
|
| def get_audio_embed_bucket_fps(self,
|
| audio_embed,
|
| fps=16,
|
| batch_frames=81,
|
| m=0):
|
| num_layers, audio_frame_num, audio_dim = audio_embed.shape
|
|
|
| if num_layers > 1:
|
| return_all_layers = True
|
| else:
|
| return_all_layers = False
|
|
|
| scale = self.video_rate / fps
|
|
|
| min_batch_num = int(audio_frame_num / (batch_frames * scale)) + 1
|
|
|
| bucket_num = min_batch_num * batch_frames
|
| padd_audio_num = math.ceil(min_batch_num * batch_frames / fps *
|
| self.video_rate) - audio_frame_num
|
| batch_idx = get_sample_indices(
|
| original_fps=self.video_rate,
|
| total_frames=audio_frame_num + padd_audio_num,
|
| target_fps=fps,
|
| num_sample=bucket_num,
|
| fixed_start=0)
|
| batch_audio_eb = []
|
| audio_sample_stride = int(self.video_rate / fps)
|
| for bi in batch_idx:
|
| if bi < audio_frame_num:
|
|
|
| chosen_idx = list(
|
| range(bi - m * audio_sample_stride,
|
| bi + (m + 1) * audio_sample_stride,
|
| audio_sample_stride))
|
| chosen_idx = [0 if c < 0 else c for c in chosen_idx]
|
| chosen_idx = [
|
| audio_frame_num - 1 if c >= audio_frame_num else c
|
| for c in chosen_idx
|
| ]
|
|
|
| if return_all_layers:
|
| frame_audio_embed = audio_embed[:, chosen_idx].flatten(
|
| start_dim=-2, end_dim=-1)
|
| else:
|
| frame_audio_embed = audio_embed[0][chosen_idx].flatten()
|
| else:
|
| frame_audio_embed = \
|
| torch.zeros([audio_dim * (2 * m + 1)], device=audio_embed.device) if not return_all_layers \
|
| else torch.zeros([num_layers, audio_dim * (2 * m + 1)], device=audio_embed.device)
|
| batch_audio_eb.append(frame_audio_embed)
|
| batch_audio_eb = torch.cat([c.unsqueeze(0) for c in batch_audio_eb],
|
| dim=0)
|
|
|
| return batch_audio_eb, min_batch_num
|
|
|
