Hecheng0625's picture
Upload 409 files
c968fc3 verified
# Copyright (c) 2024 Amphion.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from concurrent.futures import ThreadPoolExecutor
import json
import os
import librosa
import numpy as np
import time
import torch
from pydub import AudioSegment
import soundfile as sf
import onnxruntime as ort
import tqdm
import subprocess
import re
from utils.logger import Logger, time_logger
def load_cfg(cfg_path):
"""
Load configuration from a JSON file.
Args:
cfg_path (str): Path to the configuration file.
Returns:
dict: Configuration dictionary.
"""
if not os.path.exists(cfg_path):
raise FileNotFoundError(
f"{cfg_path} not found. Please copy, configure, and rename `config.json.example` to `{cfg_path}`."
)
with open(cfg_path, "r") as f:
try:
cfg = json.load(f)
except json.decoder.JSONDecodeError as e:
raise TypeError(
"Please finish the `// TODO:` in the `config.json` file before running the script. Check README.md for details."
)
return cfg
def write_wav(path, sr, x):
"""Write numpy array to WAV file."""
sf.write(path, x, sr)
def write_mp3(path, sr, x):
"""Convert numpy array to MP3."""
try:
# Ensure x is in the correct format and normalize if necessary
if x.dtype != np.int16:
# Normalize the array to fit in int16 range if it's not already int16
x = np.int16(x / np.max(np.abs(x)) * 32767)
# Create audio segment from numpy array
audio = AudioSegment(
x.tobytes(), frame_rate=sr, sample_width=x.dtype.itemsize, channels=1
)
# Export as MP3 file
audio.export(path, format="mp3")
except Exception as e:
print(e)
print("Error: Failed to write MP3 file.")
def get_audio_files(folder_path):
"""Get all audio files in a folder."""
audio_files = []
for root, _, files in os.walk(folder_path):
if "_processed" in root:
continue
for file in files:
if ".temp" in file:
continue
if file.endswith((".mp3", ".wav", ".flac", ".m4a", ".aac")):
audio_files.append(os.path.join(root, file))
return audio_files
def get_specific_files(folder_path, ext):
"""Get specific files with a given extension in a folder."""
audio_files = []
for root, _, files in os.walk(folder_path):
if "_processed" in root:
continue
for file in files:
if ".temp" in file:
continue
if file.endswith(ext):
audio_files.append(os.path.join(root, file))
return audio_files
def export_to_srt(asr_result, file_path):
"""Export ASR result to SRT file."""
with open(file_path, "w") as f:
def format_time(seconds):
return (
time.strftime("%H:%M:%S", time.gmtime(seconds))
+ f",{int(seconds * 1000 % 1000):03d}"
)
for idx, segment in enumerate(asr_result):
f.write(f"{idx + 1}\n")
f.write(
f"{format_time(segment['start'])} --> {format_time(segment['end'])}\n"
)
f.write(f"{segment['speaker']}: {segment['text']}\n\n")
def detect_gpu():
"""Detect if GPU is available and print related information."""
logger = Logger.get_logger()
if "CUDA_VISIBLE_DEVICES" not in os.environ:
logger.info("ENV: CUDA_VISIBLE_DEVICES not set, use default setting")
else:
gpu_id = os.environ["CUDA_VISIBLE_DEVICES"]
logger.info(f"ENV: CUDA_VISIBLE_DEVICES = {gpu_id}")
if not torch.cuda.is_available():
logger.error("Torch CUDA: No GPU detected. torch.cuda.is_available() = False.")
return False
num_gpus = torch.cuda.device_count()
logger.debug(f"Torch CUDA: Detected {num_gpus} GPUs.")
for i in range(num_gpus):
gpu_name = torch.cuda.get_device_name(i)
logger.debug(f" * GPU {i}: {gpu_name}")
logger.debug("Torch: CUDNN version = " + str(torch.backends.cudnn.version()))
if not torch.backends.cudnn.is_available():
logger.error("Torch: CUDNN is not available.")
return False
logger.debug("Torch: CUDNN is available.")
ort_providers = ort.get_available_providers()
logger.debug(f"ORT: Available providers: {ort_providers}")
if "CUDAExecutionProvider" not in ort_providers:
logger.warning(
"ORT: CUDAExecutionProvider is not available. "
"Please install a compatible version of ONNX Runtime. "
"See https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html"
)
return True
def get_gpu_nums():
"""Get GPU nums by nvidia-smi."""
logger = Logger.get_logger()
try:
result = subprocess.check_output("nvidia-smi -L | wc -l", shell=True)
gpus_count = int(result.decode().strip())
except Exception as e:
logger.error("Error occurred while getting GPU count: " + str(e))
gpus_count = 8 # Default to 8 if GPU count retrieval fails
return gpus_count
def check_env(logger):
"""Check environment variables."""
if "http_proxy" in os.environ:
logger.info(f"ENV: http_proxy = {os.environ['http_proxy']}")
else:
logger.info("ENV: http_proxy not set")
if "https_proxy" in os.environ:
logger.info(f"ENV: https_proxy = {os.environ['https_proxy']}")
else:
logger.info("ENV: https_proxy not set")
if "HF_ENDPOINT" in os.environ:
logger.info(
f"ENV: HF_ENDPOINT = {os.environ['HF_ENDPOINT']}, if downloading slow, try `unset HF_ENDPOINT`"
)
else:
logger.info("ENV: HF_ENDPOINT not set")
hostname = os.popen("hostname").read().strip()
logger.debug(f"HOSTNAME: {hostname}")
environ_path = os.environ["PATH"]
environ_ld_library = os.environ.get("LD_LIBRARY_PATH", "")
logger.debug(f"ENV: PATH = {environ_path}, LD_LIBRARY_PATH = {environ_ld_library}")
@time_logger
def export_to_mp3(audio, asr_result, folder_path, file_name):
"""Export segmented audio to MP3 files."""
sr = audio["sample_rate"]
audio = audio["waveform"]
os.makedirs(folder_path, exist_ok=True)
# Function to process each segment in a separate thread
def process_segment(idx, segment):
start, end = int(segment["start"] * sr), int(segment["end"] * sr)
split_audio = audio[start:end]
split_audio = librosa.to_mono(split_audio)
out_file = f"{file_name}_{idx}.mp3"
out_path = os.path.join(folder_path, out_file)
write_mp3(out_path, sr, split_audio)
# Use ThreadPoolExecutor for concurrent execution
with ThreadPoolExecutor(max_workers=72) as executor:
# Submit each segment processing as a separate thread
futures = [
executor.submit(process_segment, idx, segment)
for idx, segment in enumerate(asr_result)
]
# Wait for all threads to complete
for future in tqdm.tqdm(
futures, total=len(asr_result), desc="Exporting to MP3"
):
future.result()
@time_logger
def export_to_wav(audio, asr_result, folder_path, file_name):
"""Export segmented audio to WAV files."""
sr = audio["sample_rate"]
audio = audio["waveform"]
os.makedirs(folder_path, exist_ok=True)
for idx, segment in enumerate(tqdm.tqdm(asr_result, desc="Exporting to WAV")):
start, end = int(segment["start"] * sr), int(segment["end"] * sr)
split_audio = audio[start:end]
split_audio = librosa.to_mono(split_audio)
out_file = f"{file_name}_{idx}.wav"
out_path = os.path.join(folder_path, out_file)
write_wav(out_path, sr, split_audio)
def get_char_count(text):
"""
Get the number of characters in the text.
Args:
text (str): Input text.
Returns:
int: Number of characters in the text.
"""
# Using regular expression to remove punctuation and spaces
cleaned_text = re.sub(r"[,.!?\"',。!?“”‘’ ]", "", text)
char_count = len(cleaned_text)
return char_count
def calculate_audio_stats(
data, min_duration=3, max_duration=30, min_dnsmos=3, min_char_count=2
):
"""
Reading the proviced json, calculate and return the audio ID and their duration that meet the given filtering criteria.
Args:
data: JSON.
min_duration: Minimum duration of the audio in seconds.
max_duration: Maximum duration of the audio in seconds.
min_dnsmos: Minimum DNSMOS value.
min_char_count: Minimum number of characters.
Returns:
valid_audio_stats: A list containing tuples of audio ID and their duration.
"""
all_audio_stats = []
valid_audio_stats = []
avg_durations = []
# iterate over each entry in the JSON to collect the average duration of the phonemes
for entry in data:
# remove punctuation and spaces
char_count = get_char_count(entry["text"])
duration = entry["end"] - entry["start"]
if char_count > 0:
avg_durations.append(duration / char_count)
# calculate the bounds for the average character duration
if len(avg_durations) > 0:
q1 = np.percentile(avg_durations, 25)
q3 = np.percentile(avg_durations, 75)
iqr = q3 - q1
lower_bound = q1 - 1.5 * iqr
upper_bound = q3 + 1.5 * iqr
else:
# if no valid character data, use default values
lower_bound, upper_bound = 0, np.inf
# iterate over each entry in the JSON to apply all filtering criteria
for idx, entry in enumerate(data):
duration = entry["end"] - entry["start"]
dnsmos = entry["dnsmos"]
# remove punctuation and spaces
char_count = get_char_count(entry["text"])
if char_count > 0:
avg_char_duration = duration / char_count
else:
avg_char_duration = 0
# collect the duration of all audios
all_audio_stats.append((idx, duration))
# apply filtering criteria
if (
(min_duration <= duration <= max_duration) # withing duration range
and (dnsmos >= min_dnsmos)
and (char_count >= min_char_count)
and (
lower_bound <= avg_char_duration <= upper_bound
) # average character duration within bounds
):
valid_audio_stats.append((idx, duration))
return valid_audio_stats, all_audio_stats