File size: 2,690 Bytes
30ff368 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
import argparse
import os
import numpy as np
import librosa
import torch
from tqdm import tqdm
from transformers import Wav2Vec2FeatureExtractor, HubertModel
def main(args):
if not torch.cuda.is_available() and args.computed_device == 'cuda':
print('CUDA is not available on this device. Switching to CPU.')
args.computed_device = 'cpu'
device = torch.device(args.computed_device)
model = HubertModel.from_pretrained(args.model_path).to(device)
feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(args.model_path)
model.feature_extractor._freeze_parameters()
model.eval()
os.makedirs(args.audio_feature_saved_path, exist_ok=True)
for wavfile in tqdm(os.listdir(args.audio_dir_path)):
npy_save_path = os.path.join(args.audio_feature_saved_path, os.path.splitext(os.path.basename(wavfile))[0] + '.npy')
if os.path.exists(npy_save_path):
continue
audio, sr = librosa.load(os.path.join(args.audio_dir_path, wavfile), sr=16000)
input_values = feature_extractor(audio, sampling_rate=16000, padding=True, do_normalize=True, return_tensors="pt").input_values
input_values = input_values.to(device)
ws_feats = []
with torch.no_grad():
outputs = model(input_values, output_hidden_states=True)
for i in range(len(outputs.hidden_states)):
ws_feats.append(outputs.hidden_states[i].detach().cpu().numpy())
ws_feat_obj = np.array(ws_feats)
ws_feat_obj = np.squeeze(ws_feat_obj, 1)
if args.padding_to_align_audio:
ws_feat_obj = np.pad(ws_feat_obj, ((0, 0), (0, 1), (0, 0)), 'edge')
np.save(npy_save_path, ws_feat_obj)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Extract audio features using a pre-trained HuBERT model.")
parser.add_argument("--model_path", type=str, default='weights/chinese-hubert-large', help="Path to the pre-trained model weights.")
parser.add_argument("--audio_dir_path", type=str, default='./audio_samples/raw_audios/', help="Directory containing raw audio files.")
parser.add_argument("--audio_feature_saved_path", type=str, default='./audio_samples/audio_features/', help="Directory where extracted audio features will be saved.")
parser.add_argument("--computed_device", type=str, default='cuda', choices=['cuda', 'cpu'], help="Device to compute the audio features on. Use 'cuda' for GPU or 'cpu' for CPU.")
parser.add_argument("--padding_to_align_audio", type=bool, default=True, help="Whether to pad the audio to align features.")
args = parser.parse_args()
main(args) |