Spaces:
Running
on
Zero
Running
on
Zero
import argparse | |
import os | |
import traceback | |
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com" | |
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" | |
import torch | |
from faster_whisper import WhisperModel | |
from tqdm import tqdm | |
from tools.asr.config import check_fw_local_models | |
language_code_list = [ | |
"af", "am", "ar", "as", "az", | |
"ba", "be", "bg", "bn", "bo", | |
"br", "bs", "ca", "cs", "cy", | |
"da", "de", "el", "en", "es", | |
"et", "eu", "fa", "fi", "fo", | |
"fr", "gl", "gu", "ha", "haw", | |
"he", "hi", "hr", "ht", "hu", | |
"hy", "id", "is", "it", "ja", | |
"jw", "ka", "kk", "km", "kn", | |
"ko", "la", "lb", "ln", "lo", | |
"lt", "lv", "mg", "mi", "mk", | |
"ml", "mn", "mr", "ms", "mt", | |
"my", "ne", "nl", "nn", "no", | |
"oc", "pa", "pl", "ps", "pt", | |
"ro", "ru", "sa", "sd", "si", | |
"sk", "sl", "sn", "so", "sq", | |
"sr", "su", "sv", "sw", "ta", | |
"te", "tg", "th", "tk", "tl", | |
"tr", "tt", "uk", "ur", "uz", | |
"vi", "yi", "yo", "zh", "yue", | |
"auto"] | |
def execute_asr(input_folder, output_folder, model_size, language, precision): | |
if '-local' in model_size: | |
model_size = model_size[:-6] | |
model_path = f'tools/asr/models/faster-whisper-{model_size}' | |
else: | |
model_path = model_size | |
if language == 'auto': | |
language = None #不设置语种由模型自动输出概率最高的语种 | |
print("loading faster whisper model:",model_size,model_path) | |
device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
try: | |
model = WhisperModel(model_path, device=device, compute_type=precision) | |
except: | |
return print(traceback.format_exc()) | |
input_file_names = os.listdir(input_folder) | |
input_file_names.sort() | |
output = [] | |
output_file_name = os.path.basename(input_folder) | |
for file_name in tqdm(input_file_names): | |
try: | |
file_path = os.path.join(input_folder, file_name) | |
segments, info = model.transcribe( | |
audio = file_path, | |
beam_size = 5, | |
vad_filter = True, | |
vad_parameters = dict(min_silence_duration_ms=700), | |
language = language) | |
text = '' | |
if info.language == "zh": | |
print("检测为中文文本, 转 FunASR 处理") | |
if("only_asr"not in globals()): | |
from tools.asr.funasr_asr import \ | |
only_asr # #如果用英文就不需要导入下载模型 | |
text = only_asr(file_path) | |
if text == '': | |
for segment in segments: | |
text += segment.text | |
output.append(f"{file_path}|{output_file_name}|{info.language.upper()}|{text}") | |
except: | |
print(traceback.format_exc()) | |
output_folder = output_folder or "output/asr_opt" | |
os.makedirs(output_folder, exist_ok=True) | |
output_file_path = os.path.abspath(f'{output_folder}/{output_file_name}.list') | |
with open(output_file_path, "w", encoding="utf-8") as f: | |
f.write("\n".join(output)) | |
print(f"ASR 任务完成->标注文件路径: {output_file_path}\n") | |
return output_file_path | |
if __name__ == '__main__': | |
parser = argparse.ArgumentParser() | |
parser.add_argument("-i", "--input_folder", type=str, required=True, | |
help="Path to the folder containing WAV files.") | |
parser.add_argument("-o", "--output_folder", type=str, required=True, | |
help="Output folder to store transcriptions.") | |
parser.add_argument("-s", "--model_size", type=str, default='large-v3', | |
choices=check_fw_local_models(), | |
help="Model Size of Faster Whisper") | |
parser.add_argument("-l", "--language", type=str, default='ja', | |
choices=language_code_list, | |
help="Language of the audio files.") | |
parser.add_argument("-p", "--precision", type=str, default='float16', choices=['float16','float32','int8'], | |
help="fp16, int8 or fp32") | |
cmd = parser.parse_args() | |
output_file_path = execute_asr( | |
input_folder = cmd.input_folder, | |
output_folder = cmd.output_folder, | |
model_size = cmd.model_size, | |
language = cmd.language, | |
precision = cmd.precision, | |
) | |