mirror of
https://github.com/storytold/artcraft.git
synced 2026-11-04 01:18:34 +00:00
314 lines
13 KiB
Python
314 lines
13 KiB
Python
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import time
|
|
from pathlib import Path
|
|
from typing import List
|
|
|
|
import torch
|
|
from scipy.io import wavfile
|
|
|
|
from nemo.collections import asr as nemo_asr
|
|
from nemo.collections.asr.metrics.wer import WER, word_error_rate
|
|
|
|
parser = argparse.ArgumentParser(description="Cut audio on the segments based on segments")
|
|
parser.add_argument("--output_dir", default='output', type=str, help='Path to output directory')
|
|
parser.add_argument(
|
|
"--alignment",
|
|
type=str,
|
|
required=True,
|
|
help='Path to a data directory with alignments or a single .txt file with timestamps - result of the ctc-segmentation',
|
|
)
|
|
parser.add_argument("--threshold", type=float, default=-5, help='Minimum score value accepted')
|
|
parser.add_argument(
|
|
'--model',
|
|
type=str,
|
|
default='QuartzNet15x5Base-En',
|
|
help='Path to model checkpoint or pre-trained CTC-based ASR model name',
|
|
)
|
|
parser.add_argument('--offset', type=int, default=0, help='Offset in seconds')
|
|
parser.add_argument("--batch_size", type=int, default=64, help='Batch size for inference')
|
|
|
|
|
|
def add_transcript_to_manifest(
|
|
manifest_original: str, manifest_updated: str, asr_model: nemo_asr.models.EncDecCTCModel, batch_size: int
|
|
) -> None:
|
|
"""
|
|
Adds transcripts generated by the asr_model to the manifest_original.
|
|
|
|
Args:
|
|
manifest_original: path to the manifest
|
|
manifest_updated: path to the updated manifest with transcript included
|
|
asr_model: CTC-based ASR model, for example, QuartzNet15x5Base-En
|
|
batch_size: Batch size for asr_model inference
|
|
"""
|
|
transcripts = get_transcript(manifest_original, asr_model, batch_size)
|
|
with open(manifest_original, 'r', encoding='utf8') as f:
|
|
with open(manifest_updated, 'w', encoding='utf8') as f_updated:
|
|
for i, line in enumerate(f):
|
|
info = json.loads(line)
|
|
info['transcript'] = transcripts[i].strip()
|
|
info['WER'] = round(word_error_rate([info['transcript']], [info['text']]) * 100, 2)
|
|
info['CER'] = round(word_error_rate([info['transcript']], [info['text']], use_cer=True) * 100, 2)
|
|
json.dump(info, f_updated, ensure_ascii=False)
|
|
f_updated.write('\n')
|
|
|
|
|
|
def get_transcript(manifest_path: str, asr_model: nemo_asr.models.EncDecCTCModel, batch_size: int) -> List[str]:
|
|
"""
|
|
Returns transcripts for audio segments in the batch
|
|
|
|
Args:
|
|
manifest_path: path to the manifest for inference
|
|
asr_model: CTC-based ASR model, for example, QuartzNet15x5Base-En
|
|
batch_size: batch size
|
|
|
|
Returns: hypotheses: transcripts for the audio segments
|
|
"""
|
|
# batch inference
|
|
try:
|
|
from torch.cuda.amp import autocast
|
|
except ImportError:
|
|
from contextlib import contextmanager
|
|
|
|
@contextmanager
|
|
def autocast(enabled=None):
|
|
yield
|
|
|
|
torch.set_grad_enabled(False)
|
|
asr_model.setup_test_data(
|
|
test_data_config={
|
|
'sample_rate': 16000,
|
|
'manifest_filepath': manifest_path,
|
|
'labels': asr_model.decoder.vocabulary,
|
|
'batch_size': batch_size,
|
|
'normalize_transcripts': False,
|
|
}
|
|
)
|
|
asr_model.eval()
|
|
wer = WER(vocabulary=asr_model.decoder.vocabulary)
|
|
hypotheses = []
|
|
for test_batch in asr_model.test_dataloader():
|
|
if torch.cuda.is_available():
|
|
test_batch = [x.cuda() for x in test_batch]
|
|
with autocast():
|
|
log_probs, encoded_len, greedy_predictions = asr_model(
|
|
input_signal=test_batch[0], input_signal_length=test_batch[1]
|
|
)
|
|
hypotheses += wer.ctc_decoder_predictions_tensor(greedy_predictions)
|
|
del test_batch
|
|
torch.cuda.empty_cache()
|
|
return hypotheses
|
|
|
|
|
|
def process_alignment(alignment_file: str, args):
|
|
""" Cut original audio file into audio segments based on alignment_file
|
|
|
|
Args:
|
|
alignment_file: path to the file with segmented text and corresponding time stamps.
|
|
The first line of the file contains the path to the original audio file
|
|
args: main script args
|
|
"""
|
|
if not os.path.exists(alignment_file):
|
|
raise ValueError(f'{alignment_file} not found')
|
|
|
|
# read the segments, note the first line contains the path to the original audio
|
|
segments = []
|
|
ref_text_processed = []
|
|
ref_text_no_preprocessing = []
|
|
with open(alignment_file, 'r') as f:
|
|
for line in f:
|
|
line = line.split('|')
|
|
# read audio file name from the first line
|
|
if len(line) == 1:
|
|
audio_file = line[0].strip()
|
|
continue
|
|
ref_text_processed.append(line[1].strip())
|
|
ref_text_no_preprocessing.append(line[2].strip())
|
|
line = line[0].split()
|
|
segments.append((float(line[0]) + args.offset / 1000, float(line[1]) + args.offset / 1000, float(line[2])))
|
|
|
|
# cut the audio into segments
|
|
# create manifest in /tmp directory first, then populate transcript values with the batch inference results
|
|
# and save the final manifests at output_dir
|
|
sampling_rate, signal = wavfile.read(audio_file)
|
|
original_duration = len(signal) / sampling_rate
|
|
print(f'Cutting {audio_file} based on {alignment_file}')
|
|
print(f'Original duration: {round(original_duration)}s or ~{round(original_duration / 60)}min')
|
|
|
|
# create directories to store high score .wav audio fragments, low scored once, and deleted
|
|
args.output_dir = os.path.abspath(args.output_dir)
|
|
os.makedirs(args.output_dir, exist_ok=True)
|
|
fragments_dir = os.path.join(args.output_dir, "high_score_clips")
|
|
del_fragments = os.path.join(args.output_dir, 'deleted_clips')
|
|
low_score_segments_dir = os.path.join(args.output_dir, "low_score_clips")
|
|
|
|
os.makedirs(fragments_dir, exist_ok=True)
|
|
os.makedirs(low_score_segments_dir, exist_ok=True)
|
|
os.makedirs(del_fragments, exist_ok=True)
|
|
|
|
base_name = os.path.basename(alignment_file).replace('_segments.txt', '')
|
|
high_score_manifest = f'{base_name}_high_score_manifest.json'
|
|
low_score_manifest = f'{base_name}_low_score_manifest.json'
|
|
del_manifest = f'{base_name}_del_manifest.json'
|
|
manifests_dir = os.path.join(args.output_dir, 'manifests')
|
|
os.makedirs(manifests_dir, exist_ok=True)
|
|
tmp_dir = '/tmp'
|
|
|
|
low_score_dur = 0
|
|
high_score_dur = 0
|
|
with open(os.path.join(tmp_dir, high_score_manifest), 'w', encoding='utf8') as f:
|
|
with open(os.path.join(tmp_dir, low_score_manifest), 'w', encoding='utf8') as low_score_f:
|
|
for i, (st, end, score) in enumerate(segments):
|
|
segment = signal[round(st * sampling_rate) : round(end * sampling_rate)]
|
|
duration = len(segment) / sampling_rate
|
|
if duration > 0:
|
|
text_processed = ref_text_processed[i].strip()
|
|
text_no_preprocessing = ref_text_no_preprocessing[i].strip()
|
|
if score > args.threshold:
|
|
high_score_dur += duration
|
|
audio_filepath = os.path.join(fragments_dir, f'{base_name}_{i:04}.wav')
|
|
file_to_write = f
|
|
else:
|
|
low_score_dur += duration
|
|
audio_filepath = os.path.join(low_score_segments_dir, f'{base_name}_{i:04}.wav')
|
|
file_to_write = low_score_f
|
|
|
|
wavfile.write(audio_filepath, sampling_rate, segment)
|
|
|
|
transcript = 'n/a' # asr_model.transcribe(paths2audio_files=[audio_filepath], batch_size=1)[0]
|
|
info = {
|
|
'audio_filepath': audio_filepath,
|
|
'duration': duration,
|
|
'text': text_processed,
|
|
'text_no_preprocessing': text_no_preprocessing,
|
|
'score': round(score, 2),
|
|
'transcript': transcript.strip(),
|
|
}
|
|
json.dump(info, file_to_write, ensure_ascii=False)
|
|
file_to_write.write('\n')
|
|
|
|
add_transcript_to_manifest(
|
|
os.path.join(tmp_dir, high_score_manifest),
|
|
os.path.join(manifests_dir, high_score_manifest),
|
|
asr_model,
|
|
args.batch_size,
|
|
)
|
|
add_transcript_to_manifest(
|
|
os.path.join(tmp_dir, low_score_manifest),
|
|
os.path.join(manifests_dir, low_score_manifest),
|
|
asr_model,
|
|
args.batch_size,
|
|
)
|
|
print(f'High score files duration: {round(high_score_dur)}s or ~{round(high_score_dur/60)}min at {manifests_dir}')
|
|
print(
|
|
f'Low score files duration: {round(low_score_dur)}s or ~{round(low_score_dur/60)}min saved at {manifests_dir}'
|
|
)
|
|
|
|
# save deleted segments along with manifest
|
|
deleted = []
|
|
del_duration = 0
|
|
begin = 0
|
|
i = 0
|
|
with open(os.path.join(manifests_dir, del_manifest), 'w', encoding='utf8') as f:
|
|
for i, (st, end, _) in enumerate(segments):
|
|
if st - begin > 0.01:
|
|
segment = signal[int(begin * sampling_rate) : int(st * sampling_rate)]
|
|
audio_filepath = os.path.join(del_fragments, f'del_{base_name}_{i:04}.wav')
|
|
wavfile.write(audio_filepath, sampling_rate, segment)
|
|
duration = len(segment) / sampling_rate
|
|
del_duration += duration
|
|
deleted.append((begin, st))
|
|
transcript = 'n/a' # asr_model.transcribe(paths2audio_files=[audio_filepath], batch_size=1)[0]
|
|
info = {'audio_filepath': audio_filepath, 'duration': duration, 'text': transcript}
|
|
json.dump(info, f, ensure_ascii=False)
|
|
f.write('\n')
|
|
begin = end
|
|
|
|
segment = signal[int(begin * sampling_rate) :]
|
|
audio_filepath = os.path.join(del_fragments, f'del_{base_name}_{i+1:04}.wav')
|
|
wavfile.write(audio_filepath, sampling_rate, segment)
|
|
duration = len(segment) / sampling_rate
|
|
del_duration += duration
|
|
deleted.append((begin, original_duration))
|
|
|
|
info = {'audio_filepath': audio_filepath, 'duration': duration, 'text': 'n/a'}
|
|
json.dump(info, f)
|
|
f.write('\n')
|
|
|
|
print(f'Saved DEL files duration: {round(del_duration)}s or ~ {round(del_duration/60)}min at {del_fragments}')
|
|
missing_audio = original_duration - high_score_dur - del_duration - low_score_dur
|
|
if missing_audio > 15:
|
|
raise ValueError(f'{round(missing_audio)}s or ~ {round(missing_audio/60)}min is missing. Check the args')
|
|
|
|
stats = (
|
|
args.output_dir,
|
|
base_name,
|
|
round(original_duration),
|
|
round(high_score_dur),
|
|
round(low_score_dur),
|
|
round(del_duration),
|
|
)
|
|
return stats
|
|
|
|
|
|
if __name__ == '__main__':
|
|
start_time = time.time()
|
|
args = parser.parse_args()
|
|
|
|
os.makedirs(args.output_dir, exist_ok=True)
|
|
|
|
if os.path.exists(args.model):
|
|
asr_model = nemo_asr.models.EncDecCTCModel.restore_from(args.model, strict=False)
|
|
elif args.model in nemo_asr.models.EncDecCTCModel.get_available_model_names():
|
|
asr_model = nemo_asr.models.EncDecCTCModel.from_pretrained(args.model, strict=False)
|
|
else:
|
|
raise ValueError(
|
|
f'Provide path to the pretrained checkpoint or choose from {nemo_asr.models.EncDecCTCModel.list_available_models()}'
|
|
)
|
|
|
|
alignment_files = Path(args.alignment)
|
|
if os.path.isdir(args.alignment):
|
|
alignment_files = alignment_files.glob("*.txt")
|
|
else:
|
|
alignment_files = [Path(alignment_files)]
|
|
|
|
stats_file = os.path.join(args.output_dir, 'stats.tsv')
|
|
with open(stats_file, 'w') as f:
|
|
f.write('Folder\tSegment\tOriginal dur (s)\tHigh quality dur (s)\tLow quality dur (s)\tDeleted dur (s)\n')
|
|
|
|
high_score_dur = 0
|
|
low_score_dur = 0
|
|
del_duration = 0
|
|
|
|
for alignment_file in alignment_files:
|
|
stats = process_alignment(alignment_file, args)
|
|
high_score_dur += stats[-3]
|
|
low_score_dur += stats[-2]
|
|
del_duration += stats[-1]
|
|
stats = '\t'.join([str(t) for t in stats]) + '\n'
|
|
f.write(stats)
|
|
|
|
f.write(f'Total\t\t{round(high_score_dur)}\t{round(low_score_dur)}\t{del_duration}')
|
|
|
|
total_time = time.time() - start_time
|
|
print(f'High score segments duration: {round(high_score_dur)}')
|
|
print(f'Low score segments duration: {round(low_score_dur)}')
|
|
print(f'Deleted segments duration: {round(del_duration)}')
|
|
print(f'Stats saved at {stats_file}')
|
|
print(f'Total execution time: ~{round(total_time / 60)}min')
|