labs_to_rttms.py

#!/usr/bin/env python

# Copyright Brno University of Technology
# Licensed under the Apache License, Version 2.0 (the "License")

from lab_utils import lab_to_hard_labels
from rttm_utils import hard_labels_to_rttm
from types import SimpleNamespace
import argparse
import numpy as np
import os


def parse_arguments() -> SimpleNamespace:
    parser = argparse.ArgumentParser(description='Create RTTM from VAD labels (with speech as speaker).')
    parser.add_argument('--lab-in-dir', type=str, required=True,
                        help='directory with lab files')
    parser.add_argument('--rttm-out-dir', type=str, required=True,
                        help='directory where to save rttm files')
    parser.add_argument('--txt-list-file', type=str, required=True,
                        help='file containing list of files to process')
    parser.add_argument('--lengths', type=str, required=True,
                        help='file containing list of lengths per file')
    parser.add_argument('--overlap', type=bool, required=False, default=False,
                        help='interpret segments as overlap segments (thus create two speech segments)')
    args = parser.parse_args()
    return args


if __name__ == '__main__':
    args = parse_arguments()
    lengths_list = np.loadtxt(args.lengths, dtype=object)
    lengths = {}
    for line in lengths_list:
        name, length = line
        lengths[name] = float(length)

    files_list = np.loadtxt(args.txt_list_file, dtype=object)
    if not os.path.exists(args.rttm_out_dir):
        os.makedirs(args.rttm_out_dir)
    for name in files_list:
        matrix, labels = lab_to_hard_labels(
            os.path.join(args.lab_in_dir, f"{name}.lab"),
            1000,
            args.overlap,
            lengths[name])
        hard_labels_to_rttm(
                    matrix,
                    labels,
                    name,
                    os.path.join(args.rttm_out_dir, f"{name}.rttm"),
                    1000)