update se_extractor

This commit is contained in:
wl-zhao
2024-01-06 00:21:39 +08:00
+15 -10
View File
@@ -1,6 +1,9 @@
import os
import glob
import torch
import hashlib
import librosa
import base64
from glob import glob
import numpy as np
from pydub import AudioSegment
@@ -13,7 +16,7 @@ from whisper_timestamped.transcribe import get_audio_tensor, get_vad_segments
model_size = "medium"
# Run on GPU with FP16
model = None
def split_audio_whisper(audio_path, target_dir, audio_name):
def split_audio_whisper(audio_path, audio_name, target_dir='processed'):
global model
if model is None:
model = WhisperModel(model_size, device="cuda", compute_type="float16")
@@ -71,7 +74,7 @@ def split_audio_whisper(audio_path, target_dir, audio_name):
return wavs_folder
def split_audio_vad(audio_path, target_dir, audio_name, split_seconds=10.0):
def split_audio_vad(audio_path, audio_name, target_dir, split_seconds=10.0):
SAMPLE_RATE = 16000
audio_vad = get_audio_tensor(audio_path)
segments = get_vad_segments(
@@ -112,19 +115,21 @@ def split_audio_vad(audio_path, target_dir, audio_name, split_seconds=10.0):
count += 1
return wavs_folder
def hash_numpy_array(array):
def hash_numpy_array(audio_path):
array, _ = librosa.load(audio_path, sr=None, mono=True)
# Convert the array to bytes
array_bytes = array.tobytes()
# Calculate the hash of the array bytes
hash_object = hashlib.sha256(array_bytes)
hash_value = hash_object.digest()
# Convert the hash value to base64
base64_value = base64.b64encode(hash_value)
return base64_value.decode('utf-8')[:16].replace('/', '&')
return base64_value.decode('utf-8')[:16].replace('/', '_^')
def get_se(audio_path, vc_model, target_dir='processed', vad=True):
device = vc_model.device
audio_hash = hash_numpy_array(librosa.load(audio_path, mono=True)[0])
audio_name = os.path.basename(audio_path).rsplit('.', 1)[0] + '_' + audio_hash
audio_name = f"{os.path.basename(audio_path).rsplit('.', 1)[0]}_{hash_numpy_array(audio_path)}"
se_path = os.path.join(target_dir, audio_name, 'se.pth')
if os.path.isfile(se_path):
@@ -133,9 +138,9 @@ def get_se(audio_path, vc_model, target_dir='processed', vad=True):
if os.path.isdir(audio_path):
wavs_folder = audio_path
elif vad:
wavs_folder = split_audio_vad(audio_path, target_dir, audio_name)
wavs_folder = split_audio_vad(audio_path, target_dir=target_dir, audio_name=audio_name)
else:
wavs_folder = split_audio_whisper(audio_path, target_dir, audio_name)
wavs_folder = split_audio_whisper(audio_path, target_dir=target_dir, audio_name=audio_name)
audio_segs = glob(f'{wavs_folder}/*.wav')
if len(audio_segs) == 0: