diff --git a/se_extractor.py b/se_extractor.py index b916406..e7d897f 100644 --- a/se_extractor.py +++ b/se_extractor.py @@ -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: