TCNTempoDetector

#!/usr/bin/env python
# encoding: utf-8
"""
TCNTempoDetector beat tracking algorithm.

"""

from __future__ import absolute_import, division, print_function

import argparse

import numpy as np

from madmom.audio import SignalProcessor
from madmom.features import ActivationsProcessor
from madmom.features.beats import TCNBeatProcessor
from madmom.features.tempo import TCNTempoHistogramProcessor, TempoEstimationProcessor
from madmom.io import write_events, write_tempo
from madmom.processors import IOProcessor, io_arguments


def main():
    """TCNTempoDetector"""

    # define parser
    p = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter, description='''
    The TCNTempoDetector program detects the tempo in an audio file according
    to the method described in:

    "Multi-Task learning of tempo and beat: learning one to improve the other"
    Sebastian Böck, Matthew Davies and Peter Knees.
    Proc. of the 20th International Society for Music Information Retrieval
    Conference (ISMIR), 2019.

    This program can be run in 'single' file mode to process a single audio
    file and write the detected beats to STDOUT or the given output file.

      $ TCNTempoDetector single INFILE [-o OUTFILE]

    If multiple audio files should be processed, the program can also be run
    in 'batch' mode to save the detected beats to files with the given suffix.

      $ TCNTempoDetector batch [-o OUTPUT_DIR] [-s OUTPUT_SUFFIX] FILES

    If no output directory is given, the program writes the files with the
    detected beats to the same location as the audio files.

    The 'pickle' mode can be used to store the used parameters to be able to
    exactly reproduce experiments.

    ''')
    # version
    p.add_argument('--version', action='version',
                   version='TCNTempoDetector')
    # input/output options
    io_arguments(p, output_suffix='.bpm.txt', online=True)
    ActivationsProcessor.add_arguments(p)
    # signal processing arguments
    SignalProcessor.add_arguments(p, norm=False, gain=0)
    # tempo arguments
    TempoEstimationProcessor.add_arguments(p, hist_smooth=15)

    # parse arguments
    args = p.parse_args()

    # set immutable arguments
    args.tasks = (1, )
    args.interpolate = True
    args.method = None
    args.act_smooth = None

    # print arguments
    if args.verbose:
        print(args)

    # input processor
    if args.load:
        # load the activations from file
        in_processor = ActivationsProcessor(mode='r', **vars(args))
    else:
        # use a TCN to predict beats and tempo
        in_processor = TCNBeatProcessor(**vars(args))

    # output processor
    if args.save:
        # save the TCN activations to file
        out_processor = ActivationsProcessor(mode='w', **vars(args))
    else:
        # extract the tempo histogram from the NN output
        args.histogram_processor = TCNTempoHistogramProcessor(**vars(args))
        # estimate tempo
        tempo_estimator = TempoEstimationProcessor(**vars(args))
        # output handler
        output = write_tempo
        # sequentially process them
        out_processor = [tempo_estimator, output]

    # create an IOProcessor
    processor = IOProcessor(in_processor, out_processor)

    # and call the processing function
    args.func(processor, **vars(args))


if __name__ == '__main__':
    main()