RASMUS/Finnish-ASR-Canary-v2

1# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14 15import argparse16import multiprocessing17import shutil18from collections import OrderedDict19from pathlib import Path20from pprint import pprint21from typing import Dict22 23import matplotlib.pyplot as plt24import numpy as np25import seaborn as sns26import sox27from scipy.stats import expon28from tqdm import tqdm29 30from nemo.collections.asr.parts.utils.vad_utils import (31    get_nonspeech_segments,32    load_speech_overlap_segments_from_rttm,33    plot_sample_from_rttm,34)35 36"""37This script analyzes multi-speaker speech dataset and generates statistics.38The input directory </path/to/rttm_and_wav_directory> is required to contain the following files:39    - rttm files (*.rttm)40    - wav files (*.wav)41 42Usage:43    python <NEMO_ROOT>/scripts/speaker_tasks/multispeaker_data_analysis.py \44        </path/to/rttm_and_wav_directory> \45        --session_dur 20 \46        --silence_mean 0.2 \47        --silence_var 100 \48        --overlap_mean 0.15 \49        --overlap_var 50 \50        --num_workers 8 \51        --num_samples 10 \52        --output_dir <path/to/output_directory>53"""54 55 56def process_sample(sess_dict: Dict) -> Dict:57    """58    Process each synthetic sample59 60    Args:61        sess_dict (dict): dictionary containing the following keys62            rttm_file (str): path to the rttm file63            session_dur (float): duration of the session (specified by argument)64            precise (bool): whether to measure the precise duration of the session using sox65 66    Returns:67        results (dict): dictionary containing the following keys68            session_dur (float): duration of the session69            silence_len_list (list): list of silence durations of each silence occurrence70            silence_dur (float): total silence duration in a session71            silence_ratio (float): ratio of silence duration to session duration72            overlap_len_list (list): list of overlap durations of each overlap occurrence73            overlap_dur (float): total overlap duration74            overlap_ratio (float): ratio of overlap duration to speech (non-silence) duration75    """76 77    rttm_file = sess_dict["rttm_file"]78    session_dur = sess_dict["session_dur"]79    precise = sess_dict["precise"]80    if precise or session_dur is None:81        wav_file = rttm_file.parent / Path(rttm_file.stem + ".wav")82        session_dur = sox.file_info.duration(str(wav_file))83 84    speech_seg, overlap_seg = load_speech_overlap_segments_from_rttm(rttm_file)85    speech_dur = sum([sess_dict[1] - sess_dict[0] for sess_dict in speech_seg])86 87    silence_seg = get_nonspeech_segments(speech_seg, session_dur)88    silence_len_list = [sess_dict[1] - sess_dict[0] for sess_dict in silence_seg]89    silence_dur = max(0, session_dur - speech_dur)90    silence_ratio = silence_dur / session_dur91 92    overlap_len_list = [sess_dict[1] - sess_dict[0] for sess_dict in overlap_seg]93    overlap_dur = sum(overlap_len_list) if len(overlap_len_list) else 094    overlap_ratio = overlap_dur / speech_dur95 96    results = {97        "session_dur": session_dur,98        "silence_len_list": silence_len_list,99        "silence_dur": silence_dur,100        "silence_ratio": silence_ratio,101        "overlap_len_list": overlap_len_list,102        "overlap_dur": overlap_dur,103        "overlap_ratio": overlap_ratio,104    }105 106    return results107 108 109def run_multispeaker_data_analysis(110    input_dir,111    session_dur=None,112    silence_mean=None,113    silence_var=None,114    overlap_mean=None,115    overlap_var=None,116    precise=False,117    save_path=None,118    num_workers=1,119) -> Dict:120    rttm_list = list(Path(input_dir).glob("*.rttm"))121    """122    Analyze the multispeaker data and plot the distribution of silence and overlap durations.123 124    Args:125        input_dir (str): path to the directory containing the rttm files126        session_dur (float): duration of the session (specified by argument)127        silence_mean (float): mean of the silence duration distribution128        silence_var (float): variance of the silence duration distribution129        overlap_mean (float): mean of the overlap duration distribution130        overlap_var (float): variance of the overlap duration distribution131        precise (bool): whether to measure the precise duration of the session using sox132        save_path (str): path to save the plots133 134    Returns:135        stats (dict): dictionary containing the statistics of the analyzed data136    """137 138    print(f"Found {len(rttm_list)} files to be processed")139    if len(rttm_list) == 0:140        raise ValueError(f"No rttm files found in {input_dir}")141 142    silence_duration = 0.0143    total_duration = 0.0144    overlap_duration = 0.0145 146    silence_ratio_all = []147    overlap_ratio_all = []148    silence_length_all = []149    overlap_length_all = []150 151    queue = []152    for rttm_file in tqdm(rttm_list):153        queue.append(154            {"rttm_file": rttm_file, "session_dur": session_dur, "precise": precise,}155        )156 157    if num_workers <= 1:158        results = [process_sample(sess_dict) for sess_dict in tqdm(queue)]159    else:160        with multiprocessing.Pool(processes=num_workers) as p:161            results = list(tqdm(p.imap(process_sample, queue), total=len(queue), desc='Processing', leave=True,))162 163    for item in results:164        total_duration += item["session_dur"]165        silence_duration += item["silence_dur"]166        overlap_duration += item["overlap_dur"]167 168        silence_length_all += item["silence_len_list"]169        overlap_length_all += item["overlap_len_list"]170 171        silence_ratio_all.append(item["silence_ratio"])172        overlap_ratio_all.append(item["overlap_ratio"])173 174    actual_silence_mean = silence_duration / total_duration175    actual_silence_var = np.var(silence_ratio_all)176    actual_overlap_mean = overlap_duration / (total_duration - silence_duration)177    actual_overlap_var = np.var(overlap_ratio_all)178 179    stats = OrderedDict()180    stats["total duration (hours)"] = f"{total_duration / 3600:.2f}"181    stats["number of sessions"] = len(rttm_list)182    stats["average session duration (seconds)"] = f"{total_duration / len(rttm_list):.2f}"183    stats["actual silence ratio mean/var"] = f"{actual_silence_mean:.4f}/{actual_silence_var:.4f}"184    stats["actual overlap ratio mean/var"] = f"{actual_overlap_mean:.4f}/{actual_overlap_var:.4f}"185    stats["expected silence ratio mean/var"] = f"{silence_mean}/{silence_var}"186    stats["expected overlap ratio mean/var"] = f"{overlap_mean}/{overlap_var}"187    stats["save_path"] = save_path188 189    print("-----------------------------------------------")190    print("                    Results                    ")191    print("-----------------------------------------------")192    for k, v in stats.items():193        print(k, ": ", v)194    print("-----------------------------------------------")195 196    fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(14, 14))197    fig.suptitle(198        f"Average session={total_duration/len(rttm_list):.2f} seconds, num sessions={len(rttm_list)}, total={total_duration/3600:.2f} hours"199    )200    sns.histplot(silence_ratio_all, ax=ax1)201    ax1.set_xlabel("Silence ratio in a session")202    ax1.set_title(203        f"Target silence mean={silence_mean}, var={silence_var}. \nActual silence ratio={actual_silence_mean:.4f}, var={actual_silence_var:.4f}"204    )205 206    _, scale = expon.fit(silence_length_all, floc=0)207    sns.histplot(silence_length_all, ax=ax2)208    ax2.set_xlabel("Per-silence length in seconds")209    ax2.set_title(f"Per-silence length histogram, \nfitted exponential distribution with mean={scale:.4f}")210 211    sns.histplot(overlap_ratio_all, ax=ax3)212    ax3.set_title(213        f"Target overlap mean={overlap_mean}, var={overlap_var}. \nActual ratio={actual_overlap_mean:.4f}, var={actual_overlap_var:.4f}"214    )215    ax3.set_xlabel("Overlap ratio in a session")216    _, scale2 = expon.fit(overlap_length_all, floc=0)217    sns.histplot(overlap_length_all, ax=ax4)218    ax4.set_title(f"Per overlap length histogram, \nfitted exponential distribution with mean={scale2:.4f}")219    ax4.set_xlabel("Duration in seconds")220 221    if save_path:222        fig.savefig(save_path)223        print(f"Figure saved at: {save_path}")224 225    return stats226 227 228def visualize_multispeaker_data(input_dir: str, output_dir: str, num_samples: int = 10) -> None:229    """230    Visualize a set of randomly sampled data in the input directory231 232    Args:233        input_dir (str): Path to the input directory234        output_dir (str): Path to the output directory235        num_samples (int): Number of samples to visualize236    """237    rttm_list = list(Path(input_dir).glob("*.rttm"))238    idx_list = np.random.permutation(len(rttm_list))[:num_samples]239    print(f"Visualizing {num_samples} random samples")240    for idx in idx_list:241        rttm_file = rttm_list[idx]242        audio_file = rttm_file.parent / Path(rttm_file.stem + ".wav")243        output_file = Path(output_dir) / Path(rttm_file.stem + ".png")244        plot_sample_from_rttm(audio_file=audio_file, rttm_file=rttm_file, save_path=str(output_file), show=False)245    print(f"Sample plots saved at: {output_dir}")246 247 248if __name__ == "__main__":249    parser = argparse.ArgumentParser()250    parser.add_argument("input_dir", default="", help="Input directory")251    parser.add_argument("-sd", "--session_dur", default=None, type=float, help="Duration per session in seconds")252    parser.add_argument("-sm", "--silence_mean", default=None, type=float, help="Expected silence ratio mean")253    parser.add_argument("-sv", "--silence_var", default=None, type=float, help="Expected silence ratio variance")254    parser.add_argument("-om", "--overlap_mean", default=None, type=float, help="Expected overlap ratio mean")255    parser.add_argument("-ov", "--overlap_var", default=None, type=float, help="Expected overlap ratio variance")256    parser.add_argument("-w", "--num_workers", default=1, type=int, help="Number of CPU workers to use")257    parser.add_argument("-s", "--num_samples", default=10, type=int, help="Number of random samples to plot")258    parser.add_argument("-o", "--output_dir", default="analysis/", type=str, help="Directory for saving output figure")259    parser.add_argument(260        "--precise", action="store_true", help="Set to get precise duration, with significant time cost"261    )262    args = parser.parse_args()263 264    print("Running with params:")265    pprint(vars(args))266 267    output_dir = Path(args.output_dir)268    if output_dir.exists():269        print(f"Removing existing output directory: {args.output_dir}")270        shutil.rmtree(str(output_dir))271    output_dir.mkdir(parents=True)272 273    run_multispeaker_data_analysis(274        input_dir=args.input_dir,275        session_dur=args.session_dur,276        silence_mean=args.silence_mean,277        silence_var=args.silence_var,278        overlap_mean=args.overlap_mean,279        overlap_var=args.overlap_var,280        precise=args.precise,281        save_path=str(Path(args.output_dir, "statistics.png")),282        num_workers=args.num_workers,283    )284 285    visualize_multispeaker_data(input_dir=args.input_dir, output_dir=args.output_dir, num_samples=args.num_samples)286 287    print("The multispeaker data analysis has been completed.")288    print(f"Please check the output directory: \n{args.output_dir}")289