File size: 4,817 Bytes
6ffdd29 |
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 115 |
# Copyright (c) 2023, 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 os
import soundfile as sf
from utils.constants import BLANK_TOKEN, SPACE_TOKEN
from utils.data_prep import Segment, Word
def make_ctm_files(
utt_obj, output_dir_root, ctm_file_config,
):
"""
Function to save CTM files for all the utterances in the incoming batch.
"""
# don't try to make files if utt_obj.segments_and_tokens is empty, which will happen
# in the case of the ground truth text being empty or the number of tokens being too large vs audio duration
if not utt_obj.segments_and_tokens:
return utt_obj
# get audio file duration if we will need it later
if ctm_file_config.minimum_timestamp_duration > 0:
with sf.SoundFile(utt_obj.audio_filepath) as f:
audio_file_duration = f.frames / f.samplerate
else:
audio_file_duration = None
utt_obj = make_ctm("tokens", utt_obj, output_dir_root, audio_file_duration, ctm_file_config,)
utt_obj = make_ctm("words", utt_obj, output_dir_root, audio_file_duration, ctm_file_config,)
utt_obj = make_ctm("segments", utt_obj, output_dir_root, audio_file_duration, ctm_file_config,)
return utt_obj
def make_ctm(
alignment_level, utt_obj, output_dir_root, audio_file_duration, ctm_file_config,
):
output_dir = os.path.join(output_dir_root, "ctm", alignment_level)
os.makedirs(output_dir, exist_ok=True)
boundary_info_utt = []
for segment_or_token in utt_obj.segments_and_tokens:
if type(segment_or_token) is Segment:
segment = segment_or_token
if alignment_level == "segments":
boundary_info_utt.append(segment)
for word_or_token in segment.words_and_tokens:
if type(word_or_token) is Word:
word = word_or_token
if alignment_level == "words":
boundary_info_utt.append(word)
for token in word.tokens:
if alignment_level == "tokens":
boundary_info_utt.append(token)
else:
token = word_or_token
if alignment_level == "tokens":
boundary_info_utt.append(token)
else:
token = segment_or_token
if alignment_level == "tokens":
boundary_info_utt.append(token)
with open(os.path.join(output_dir, f"{utt_obj.utt_id}.ctm"), "w") as f_ctm:
for boundary_info_ in boundary_info_utt: # loop over every token/word/segment
# skip if t_start = t_end = negative number because we used it as a marker to skip some blank tokens
if not (boundary_info_.t_start < 0 or boundary_info_.t_end < 0):
text = boundary_info_.text
start_time = boundary_info_.t_start
end_time = boundary_info_.t_end
if (
ctm_file_config.minimum_timestamp_duration > 0
and ctm_file_config.minimum_timestamp_duration > end_time - start_time
):
# make the predicted duration of the token/word/segment longer, growing it outwards equal
# amounts from the predicted center of the token/word/segment
token_mid_point = (start_time + end_time) / 2
start_time = max(token_mid_point - ctm_file_config.minimum_timestamp_duration / 2, 0)
end_time = min(
token_mid_point + ctm_file_config.minimum_timestamp_duration / 2, audio_file_duration
)
if not (
text == BLANK_TOKEN and ctm_file_config.remove_blank_tokens
): # don't save blanks if we don't want to
# replace any spaces with <space> so we dont introduce extra space characters to our CTM files
text = text.replace(" ", SPACE_TOKEN)
f_ctm.write(f"{utt_obj.utt_id} 1 {start_time:.2f} {end_time - start_time:.2f} {text}\n")
utt_obj.saved_output_files[f"{alignment_level}_level_ctm_filepath"] = os.path.join(
output_dir, f"{utt_obj.utt_id}.ctm"
)
return utt_obj
|