| import torch |
| import numpy as np |
| from transformers import Wav2Vec2FeatureExtractor, AutoModel |
|
|
| class FeatureExtractorMERT: |
| def __init__(self, model_name="m-a-p/MERT-v1-95M", device = "None", sr=24000): |
| self.model_name = model_name |
| self.sr = sr |
| if device == "None": |
| use_cuda = torch.cuda.is_available() |
| device = torch.device("cuda" if use_cuda else "cpu") |
| else: |
| self.device = device |
| |
| self.model = AutoModel.from_pretrained(self.model_name, trust_remote_code=True).to(self.device) |
| self.processor = Wav2Vec2FeatureExtractor.from_pretrained(self.model_name, trust_remote_code=True) |
|
|
| def extract_features_from_segment(self, segment, sample_rate, save_path): |
| input_audio = segment.float() |
| model_inputs = self.processor(input_audio, sampling_rate=sample_rate, return_tensors="pt") |
| model_inputs = model_inputs.to(self.device) |
|
|
| with torch.no_grad(): |
| model_outputs = self.model(**model_inputs, output_hidden_states=True) |
|
|
| |
| all_layer_hidden_states = torch.stack(model_outputs.hidden_states).squeeze()[1:, :, :].unsqueeze(0) |
| all_layer_hidden_states = all_layer_hidden_states.mean(dim=2) |
| features = all_layer_hidden_states.cpu().detach().numpy() |
|
|
| |
| np.save(save_path, features) |
