Spaces:
Running
on
L4
Running
on
L4
File size: 3,494 Bytes
0a3525d |
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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 |
"""
Used to transcribe all audio files in one folder into another folder.
e.g.
Directory structure:
--pre_data_root
----SP_1
------01.wav
------02.wav
------......
----SP_2
------01.wav
------02.wav
------......
Use
python tools/whisper_asr.py --audio_dir pre_data_root/SP_1 --save_dir data/SP_1
to transcribe the first speaker.
Use
python tools/whisper_asr.py --audio_dir pre_data_root/SP_2 --save_dir data/SP_2
to transcribe the second speaker.
Note: Be aware of your audio sample rate, which defaults to 44.1kHz.
"""
from pathlib import Path
import click
import librosa
import soundfile as sf
import whisper
from loguru import logger
from merge_asr_files import merge_and_delete_files
from tqdm import tqdm
from fish_speech.utils.file import AUDIO_EXTENSIONS, list_files
@click.command()
@click.option("--model-size", default="large", help="Size of the Whisper model")
@click.option("--audio-dir", required=True, help="Directory containing audio files")
@click.option(
"--save-dir", required=True, help="Directory to save processed audio files"
)
@click.option(
"--sample-rate",
default=None,
type=int,
help="Output sample rate, default to input sample rate",
)
@click.option("--device", default="cuda", help="Device to use")
@click.option("--language", default="ZH", help="Language of the transcription")
def main(model_size, audio_dir, save_dir, sample_rate, device, language):
logger.info("Loading / Downloading OpenAI Whisper model...")
model = whisper.load_model(
name=model_size,
device=device,
download_root=str(Path(".cache/whisper").resolve()),
)
logger.info("Model loaded.")
save_path = Path(save_dir)
save_path.mkdir(parents=True, exist_ok=True)
original_files = []
audio_files = list_files(
path=audio_dir, extensions=AUDIO_EXTENSIONS, recursive=True
)
for file_path in tqdm(audio_files, desc="Processing audio file"):
file_stem = file_path.stem
file_suffix = file_path.suffix
rel_path = Path(file_path).relative_to(audio_dir)
(save_path / rel_path.parent).mkdir(parents=True, exist_ok=True)
if (save_path / rel_path.parent / f"{rel_path.stem}.wav").exists() and (
save_path / rel_path.parent / f"{rel_path.stem}.lab"
).exists():
continue
audio, sr = librosa.load(file_path, sr=sample_rate, mono=False)
transcription = model.transcribe(str(file_path), language=language)
for segment in transcription.get("segments", []):
id, text, start, end = (
segment["id"],
segment["text"],
segment["start"],
segment["end"],
)
extract = audio[..., int(start * sr) : int(end * sr)]
audio_save_path = (
save_path / rel_path.parent / f"{file_stem}-{id}{file_suffix}"
)
sf.write(
audio_save_path,
extract,
samplerate=sr,
)
original_files.append(audio_save_path)
transcript_save_path = save_path / rel_path.parent / f"{file_stem}-{id}.lab"
with open(
transcript_save_path,
"w",
encoding="utf-8",
) as f:
f.write(text)
original_files.append(transcript_save_path)
merge_and_delete_files(save_dir, original_files)
if __name__ == "__main__":
main()
|