mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-17 00:23:14 +00:00
Introduce code from Rust monorepo. This is a mess
This commit is contained in:
@@ -0,0 +1,259 @@
|
||||
# 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 glob
|
||||
import os
|
||||
import time
|
||||
from argparse import ArgumentParser
|
||||
from itertools import repeat
|
||||
from multiprocessing import Pool
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
|
||||
"""
|
||||
This script serves two purposes:
|
||||
|
||||
1) gen_overlap_seq:
|
||||
Generate predictions with overlapping input segments by using the frame level prediction from NeMo/examples/asr/vad_infer.py.
|
||||
Then a smoothing filter is applied to decide the label for a frame spanned by multiple segments.
|
||||
|
||||
2)gen_seg_table:
|
||||
Converting frame level prediction to speech/no-speech segment in start and end times format.
|
||||
|
||||
Usage:
|
||||
|
||||
python vad_overlap_posterior.py --gen_overlap_seq --gen_seg_table --frame_folder=<FULL PATH OF YOU STORED FRAME LEVEL PREDICTION> --method='median' --overlap=0.875 --num_workers=20 --threshold=0.8
|
||||
|
||||
"""
|
||||
|
||||
|
||||
def gen_overlap_seq(frame_filepath, per_args):
|
||||
"""
|
||||
Given a frame level prediction, generate predictions with overlapping input segments by using it
|
||||
Args:
|
||||
frame_filepath : frame prediction file to be processed.
|
||||
per_args:
|
||||
method : Median or mean smoothing filter.
|
||||
overlap : Amounts of overlap.
|
||||
seg_len : Length of window for generating the frame.
|
||||
shift_len : Amount of shift of window for generating the frame.
|
||||
out_dir : Output dir of generated prediction.
|
||||
"""
|
||||
|
||||
try:
|
||||
method = per_args['method']
|
||||
overlap = per_args['overlap']
|
||||
seg_len = per_args['seg_len']
|
||||
shift_len = per_args['shift_len']
|
||||
out_dir = per_args['out_dir']
|
||||
|
||||
frame = np.loadtxt(frame_filepath)
|
||||
name = os.path.basename(frame_filepath).split(".frame")[0] + "." + method
|
||||
overlap_filepath = os.path.join(out_dir, name)
|
||||
|
||||
shift = int(shift_len / 0.01) # number of units of shift
|
||||
seg = int((seg_len / 0.01 + 1)) # number of units of each window/segment
|
||||
|
||||
jump_on_target = int(seg * (1 - overlap)) # jump on target generated sequence
|
||||
jump_on_frame = int(jump_on_target / shift) # jump on input frame sequence
|
||||
|
||||
if jump_on_frame < 1:
|
||||
raise ValueError(
|
||||
f"Note we jump over frame sequence to generate overlapping input segments. \n \
|
||||
Your input makes jump_on_fram={jump_on_frame} < 1 which is invalid because it cannot jump and will stuck.\n \
|
||||
Please try different seg_len, shift_len and overlap choices. \n \
|
||||
jump_on_target = int(seg * (1 - overlap)) \n \
|
||||
jump_on_frame = int(jump_on_frame/shift) "
|
||||
)
|
||||
|
||||
target_len = int(len(frame) * shift)
|
||||
|
||||
if method == 'mean':
|
||||
preds = np.zeros(target_len)
|
||||
pred_count = np.zeros(target_len)
|
||||
|
||||
for i, og_pred in enumerate(frame):
|
||||
if i % jump_on_frame != 0:
|
||||
continue
|
||||
start = i * shift
|
||||
end = start + seg
|
||||
preds[start:end] = preds[start:end] + og_pred
|
||||
pred_count[start:end] = pred_count[start:end] + 1
|
||||
|
||||
preds = preds / pred_count
|
||||
last_non_zero_pred = preds[pred_count != 0][-1]
|
||||
preds[pred_count == 0] = last_non_zero_pred
|
||||
|
||||
elif method == 'median':
|
||||
preds = [[] for _ in range(target_len)]
|
||||
for i, og_pred in enumerate(frame):
|
||||
if i % jump_on_frame != 0:
|
||||
continue
|
||||
|
||||
start = i * shift
|
||||
end = start + seg
|
||||
for j in range(start, end):
|
||||
if j <= target_len - 1:
|
||||
preds[j].append(og_pred)
|
||||
|
||||
preds = np.array([np.median(l) for l in preds])
|
||||
nan_idx = np.isnan(preds)
|
||||
last_non_nan_pred = preds[~nan_idx][-1]
|
||||
preds[nan_idx] = last_non_nan_pred
|
||||
|
||||
else:
|
||||
raise ValueError("method should be either mean or median")
|
||||
|
||||
round_final = np.round(preds, 4)
|
||||
np.savetxt(overlap_filepath, round_final, delimiter='\n')
|
||||
print(f"Finished! {overlap_filepath}!")
|
||||
|
||||
except Exception as e:
|
||||
raise (e)
|
||||
|
||||
|
||||
def gen_seg_table(frame_filepath, per_args):
|
||||
|
||||
"""
|
||||
Convert frame level prediction to speech/no-speech segment in start and end times format.
|
||||
And save to csv file in rttm-like format
|
||||
0, 10, speech
|
||||
10,12, no-speech
|
||||
Args:
|
||||
frame_filepath : frame prediction file to be processed.
|
||||
per_args :
|
||||
threshold : threshold for prediction score (from 0 to 1).
|
||||
shift_len : Amount of shift of window for generating the frame.
|
||||
out_dir : Output dir of generated table/csv file.
|
||||
"""
|
||||
threshold = per_args['threshold']
|
||||
shift_len = per_args['shift_len']
|
||||
out_dir = per_args['out_dir']
|
||||
|
||||
print(f"process {frame_filepath}")
|
||||
name = frame_filepath.split("/")[-1].rsplit(".", 1)[0]
|
||||
|
||||
sequence = np.loadtxt(frame_filepath)
|
||||
start = 0
|
||||
end = 0
|
||||
start_list = [0]
|
||||
end_list = []
|
||||
state_list = []
|
||||
|
||||
for i in range(len(sequence) - 1):
|
||||
current_sate = "non-speech" if sequence[i] <= threshold else "speech"
|
||||
next_state = "non-speech" if sequence[i + 1] <= threshold else "speech"
|
||||
if next_state != current_sate:
|
||||
end = i * shift_len + shift_len # shift_len for handling joint
|
||||
state_list.append(current_sate)
|
||||
end_list.append(end)
|
||||
|
||||
start = (i + 1) * shift_len
|
||||
start_list.append(start)
|
||||
|
||||
end_list.append((i + 1) * shift_len + shift_len)
|
||||
state_list.append(current_sate)
|
||||
|
||||
seg_table = pd.DataFrame({'start': start_list, 'end': end_list, 'vad': state_list})
|
||||
|
||||
save_name = name + ".txt"
|
||||
save_path = os.path.join(out_dir, save_name)
|
||||
seg_table.to_csv(save_path, sep='\t', index=False, header=False)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
start = time.time()
|
||||
parser = ArgumentParser()
|
||||
parser.add_argument("--gen_overlap_seq", default=False, action='store_true')
|
||||
parser.add_argument("--gen_seg_table", default=False, action='store_true')
|
||||
parser.add_argument("--frame_folder", type=str, required=True)
|
||||
parser.add_argument(
|
||||
"--method",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Use mean/median for overlapped prediction. Use frame for gen_seg_table of frame prediction",
|
||||
)
|
||||
parser.add_argument("--overlap_out_dir", type=str)
|
||||
parser.add_argument("--table_out_dir", type=str)
|
||||
parser.add_argument("--overlap", type=float, default=0.875, help="Overlap percentatge. Default is 0.875")
|
||||
parser.add_argument("--seg_len", type=float, default=0.63)
|
||||
parser.add_argument("--shift_len", type=float, default=0.01)
|
||||
parser.add_argument("--threshold", type=float, default=0.5)
|
||||
parser.add_argument("--num_workers", type=int, default=4)
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.gen_overlap_seq:
|
||||
p = Pool(processes=args.num_workers)
|
||||
print("Generating predictions with overlapping input segments")
|
||||
frame_filepathlist = glob.glob(args.frame_folder + "/*.frame")
|
||||
|
||||
if not args.overlap_out_dir:
|
||||
overlap_out_dir = "./overlap_smoothing_output" + "_" + args.method + "_" + str(args.overlap)
|
||||
else:
|
||||
overlap_out_dir = args.overlap_out_dir
|
||||
if not os.path.exists(overlap_out_dir):
|
||||
print(f"Creating output dir {overlap_out_dir}")
|
||||
os.mkdir(overlap_out_dir)
|
||||
|
||||
per_args = {
|
||||
"method": args.method,
|
||||
"overlap": args.overlap,
|
||||
"seg_len": args.seg_len,
|
||||
"shift_len": args.shift_len,
|
||||
"out_dir": overlap_out_dir,
|
||||
}
|
||||
|
||||
p.starmap(gen_overlap_seq, zip(frame_filepathlist, repeat(per_args)))
|
||||
p.close()
|
||||
p.join()
|
||||
|
||||
end = time.time()
|
||||
print(f"Generate overlapped prediction takes {end-start} seconds!\n Save to {overlap_out_dir}")
|
||||
|
||||
if args.gen_seg_table:
|
||||
p = Pool(processes=args.num_workers)
|
||||
print("Converting frame level prediction to speech/no-speech segment in start and end times format.")
|
||||
|
||||
if args.gen_overlap_seq:
|
||||
print("Use overlap prediction. Change if you want to use basic frame level prediction")
|
||||
frame_filepath = overlap_out_dir
|
||||
shift_len = 0.01
|
||||
else:
|
||||
print("Use basic frame level prediction")
|
||||
frame_filepath = args.frame_folder
|
||||
shift_len = args.shift_len
|
||||
|
||||
frame_filepathlist = glob.glob(frame_filepath + "/*." + args.method)
|
||||
|
||||
if not args.table_out_dir:
|
||||
table_out_dir = "table_output_" + str(args.threshold)
|
||||
else:
|
||||
table_out_dir = args.table_out_dir
|
||||
if not os.path.exists(table_out_dir):
|
||||
print(f"Creating rttm-like table output dir {table_out_dir}")
|
||||
os.mkdir(table_out_dir)
|
||||
|
||||
per_args = {
|
||||
"threshold": args.threshold,
|
||||
"shift_len": shift_len,
|
||||
"out_dir": table_out_dir,
|
||||
}
|
||||
|
||||
p.starmap(gen_seg_table, zip(frame_filepathlist, repeat(per_args)))
|
||||
p.close()
|
||||
p.join()
|
||||
|
||||
end = time.time()
|
||||
print(f"Generate rttm-like table for {frame_filepath} takes {end-start} seconds!\n Save to {table_out_dir}")
|
||||
@@ -0,0 +1,58 @@
|
||||
# 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 numpy as np
|
||||
|
||||
from nemo.collections.asr.parts.utils.vad_utils import vad_tune_threshold_on_dev
|
||||
from nemo.utils import logging
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--threshold_range", help="range of threshold in list 'START,END,STEP' to be tuned on", required=True
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vad_pred", help="Directory of vad predictions or a file contains the paths of them.", required=True
|
||||
)
|
||||
parser.add_argument(
|
||||
"--groundtruth_RTTM",
|
||||
help="Directory of groundtruch rttm files or a file contains the paths of them",
|
||||
type=str,
|
||||
required=True,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vad_pred_method",
|
||||
help="suffix of prediction file. Should be either in 'frame', 'mean' or 'median'",
|
||||
required=True,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--focus_metric",
|
||||
help="metrics we care most when tuning threshold. Should be either in 'DetER', 'FA', 'MISS' ",
|
||||
type=str,
|
||||
default='DetER',
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
try:
|
||||
start, stop, step = [float(i) for i in args.threshold_range.split(",")]
|
||||
thresholds = np.arange(start, stop, step)
|
||||
except:
|
||||
raise ValueError("Theshold input is invalid! Please enter it as a 'START,STOP,STEP' ")
|
||||
|
||||
best_threhsold = vad_tune_threshold_on_dev(
|
||||
thresholds, args.vad_pred, args.groundtruth_RTTM, args.vad_pred_method, args.focus_metric
|
||||
)
|
||||
logging.info(f"Best threshold selected from {thresholds} is {best_threhsold}!")
|
||||
@@ -0,0 +1,173 @@
|
||||
# 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 json
|
||||
import logging
|
||||
import multiprocessing
|
||||
import os
|
||||
from argparse import ArgumentParser
|
||||
from itertools import repeat
|
||||
|
||||
import librosa
|
||||
import numpy as np
|
||||
|
||||
|
||||
"""
|
||||
This script is designed for inference of frame level Voice Activity Detection (VAD)
|
||||
|
||||
This script serves three goals:
|
||||
(1) Write audio files to manifest
|
||||
(2) Split audio file for avoiding CUDA memory issue
|
||||
(3) Take care of joint of seperate json line for an audio file
|
||||
|
||||
Usage:
|
||||
python write_long_audio_manifest.py --inp_dir=<FULL PATH OF FOLDER OF AUDIO FILES> --manifest_name=<NAME OF OUTPUT JSON FILE> --split_duration=300 --time_length=0.63 --num_worker=10
|
||||
|
||||
"""
|
||||
|
||||
|
||||
def write_manifest(file, args_func):
|
||||
"""
|
||||
Given a list of files, write them to manifest with restrictions.
|
||||
Args:
|
||||
files : file to be processed
|
||||
label : label for audio snippet.
|
||||
split_duration : Max duration of each audio clip (each line in json)
|
||||
time_length : Used for taking care of joint.
|
||||
Length of window for generating the frame.
|
||||
Returns:
|
||||
res : list of generated metadata line of json for file
|
||||
"""
|
||||
|
||||
label = args_func['label']
|
||||
split_duration = args_func['split_duration']
|
||||
time_length = args_func['time_length']
|
||||
|
||||
res = []
|
||||
try:
|
||||
sr = 16000
|
||||
x, _sr = librosa.load(file, sr=sr)
|
||||
duration = librosa.get_duration(x, sr=sr)
|
||||
|
||||
left = duration
|
||||
current_offset = 0
|
||||
status = 'single'
|
||||
|
||||
while left > 0:
|
||||
if left <= split_duration:
|
||||
if status == 'single':
|
||||
write_duration = left
|
||||
current_offset = 0
|
||||
else:
|
||||
status = 'end'
|
||||
write_duration = left + time_length
|
||||
current_offset -= time_length
|
||||
offset_inc = left
|
||||
left = 0
|
||||
else:
|
||||
if status == 'start' or status == 'next':
|
||||
status = 'next'
|
||||
else:
|
||||
status = 'start'
|
||||
|
||||
if status == 'start':
|
||||
write_duration = split_duration
|
||||
offset_inc = split_duration
|
||||
else:
|
||||
write_duration = split_duration + time_length
|
||||
current_offset -= time_length
|
||||
offset_inc = split_duration + time_length
|
||||
|
||||
left -= split_duration
|
||||
|
||||
metadata = {
|
||||
'audio_filepath': file,
|
||||
'duration': write_duration,
|
||||
'label': label,
|
||||
'text': '_',
|
||||
'offset': current_offset,
|
||||
}
|
||||
res.append(metadata)
|
||||
|
||||
current_offset += offset_inc
|
||||
|
||||
except Exception as e:
|
||||
err_file = "error.log"
|
||||
with open(err_file, 'w') as fout:
|
||||
fout.write(file + ":" + str(e))
|
||||
|
||||
return res
|
||||
|
||||
|
||||
def main():
|
||||
parser = ArgumentParser()
|
||||
parser.add_argument("--inp_dir", type=str, required=True, help="(full path) folder of files to be processed")
|
||||
parser.add_argument(
|
||||
"--inp_list", type=str, help="(full path) a file contains NAME of files inside inp_dir to be processed"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--out_dir", type=str, default=".", help="[Optional](full path) location to store generated json file"
|
||||
)
|
||||
parser.add_argument("--manifest_name", type=str, required=True, help="name of generated json file")
|
||||
parser.add_argument("--split_duration", type=int, required=True, help="max duration of each audio clip/line")
|
||||
parser.add_argument(
|
||||
"--time_length", type=float, default=0.63, help="[Optional] time length of segment, default is 0.63s"
|
||||
)
|
||||
parser.add_argument("--num_workers", type=int, default=4, help="[Optional] number of workers for multiprocessing")
|
||||
args = parser.parse_args()
|
||||
|
||||
if not args.inp_list:
|
||||
input_audios = []
|
||||
for root, dirs, files in os.walk(args.inp_dir):
|
||||
for basename in files:
|
||||
if basename.endswith('.wav'):
|
||||
filename = os.path.join(root, basename)
|
||||
input_audios.append(filename)
|
||||
else:
|
||||
name_list = np.loadtxt(args.inp_list, dtype='str')
|
||||
input_audios = [os.path.join(args.inp_dir, name + ".wav") for name in name_list]
|
||||
print(f"Number of wav files to be processed: {len(input_audios)}")
|
||||
|
||||
output_path = os.path.join(args.out_dir, args.manifest_name + '.json')
|
||||
print(f"Save generate manifest to {output_path}!")
|
||||
if not os.path.exists(args.out_dir):
|
||||
logging.info(f'Outdir {args.out_dir} does not exist. Creat directory.')
|
||||
os.mkdir(args.out_dir)
|
||||
if os.path.exists(output_path):
|
||||
print(f"Manifest {output_path} exists! Overwriting")
|
||||
os.remove(output_path)
|
||||
|
||||
p = multiprocessing.Pool(processes=args.num_workers)
|
||||
|
||||
args_func = {
|
||||
'label': 'infer',
|
||||
'split_duration': args.split_duration,
|
||||
'time_length': args.time_length,
|
||||
}
|
||||
|
||||
results = p.starmap(write_manifest, zip(input_audios, repeat(args_func)))
|
||||
|
||||
with open(output_path, 'a') as fout:
|
||||
for res in results:
|
||||
for r in res:
|
||||
json.dump(r, fout)
|
||||
fout.write('\n')
|
||||
fout.flush()
|
||||
|
||||
p.close()
|
||||
|
||||
print(f"Done! Save to {output_path}")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in New Issue
Block a user