python/lib/aubio/cmd.py: add a flush function
[aubio.git] / python / lib / aubio / cmd.py
1 #! /usr/bin/env python
2 # -*- coding: utf-8 -*-
3
4 """aubio command line tool
5
6 This file was written by Paul Brossier <piem@aubio.org> and is released under
7 the GNU/GPL v3.
8
9 Note: this script is mostly about parsing command line arguments. For more
10 readable code examples, check out the `python/demos` folder."""
11
12 import sys
13 import argparse
14 import aubio
15
16 def aubio_parser():
17     epilog = 'use "%(prog)s <command> --help" for more info about each command'
18     parser = argparse.ArgumentParser(epilog=epilog)
19     parser.add_argument('-V', '--version', help="show version",
20             action="store_true", dest="show_version")
21
22     subparsers = parser.add_subparsers(dest='command',
23             description="", metavar="<command>")
24
25     # onset subcommand
26     subparser = subparsers.add_parser('onset',
27             help='get onset times',
28             formatter_class = argparse.ArgumentDefaultsHelpFormatter)
29     parser_add_input(subparser)
30     parser_add_buf_hop_size(subparser)
31     helpstr = "onset novelty function"
32     helpstr += " <default|energy|hfc|complex|phase|specdiff|kl|mkl|specflux>"
33     parser_add_method(subparser, helpstr=helpstr)
34     parser_add_threshold(subparser)
35     parser_add_silence(subparser)
36     parser_add_minioi(subparser)
37     parser_add_time_format(subparser)
38     parser_add_verbose_help(subparser)
39     subparser.set_defaults(process=process_onset)
40
41     # pitch subcommand
42     subparser = subparsers.add_parser('pitch',
43             help='extract fundamental frequency')
44     parser_add_input(subparser)
45     parser_add_buf_hop_size(subparser, buf_size=2048)
46     helpstr = "pitch detection method <default|yinfft|yin|mcomb|fcomb|schmitt>"
47     parser_add_method(subparser, helpstr=helpstr)
48     parser_add_threshold(subparser)
49     parser_add_silence(subparser)
50     parser_add_time_format(subparser)
51     parser_add_verbose_help(subparser)
52     subparser.set_defaults(process=process_pitch)
53
54     # tempo subcommand
55     subparser = subparsers.add_parser('beat',
56             help='get locations of beats')
57     parser_add_input(subparser)
58     parser_add_buf_hop_size(subparser, buf_size=1024, hop_size=512)
59     parser_add_time_format(subparser)
60     parser_add_verbose_help(subparser)
61     subparser.set_defaults(process=process_tempo)
62
63     # notes subcommand
64     subparser = subparsers.add_parser('notes',
65             help='get midi-like notes')
66     parser_add_input(subparser)
67     parser_add_buf_hop_size(subparser)
68     parser_add_time_format(subparser)
69     parser_add_verbose_help(subparser)
70     subparser.set_defaults(process=process_notes)
71
72     # mfcc subcommand
73     subparser = subparsers.add_parser('mfcc',
74             help='extract mel-frequency cepstrum coefficients')
75     parser_add_input(subparser)
76     parser_add_buf_hop_size(subparser)
77     parser_add_time_format(subparser)
78     parser_add_verbose_help(subparser)
79     subparser.set_defaults(process=process_mfcc)
80
81     # melbands subcommand
82     subparser = subparsers.add_parser('melbands',
83             help='extract mel-frequency energies per band')
84     parser_add_input(subparser)
85     parser_add_buf_hop_size(subparser)
86     parser_add_time_format(subparser)
87     parser_add_verbose_help(subparser)
88     subparser.set_defaults(process=process_melbands)
89
90     return parser
91
92 def parser_add_input(parser):
93     parser.add_argument("source_uri", default=None, nargs='?',
94             help="input sound file to analyse", metavar = "<source_uri>")
95     parser.add_argument("-i", "--input", dest = "source_uri2",
96             help="input sound file to analyse", metavar = "<source_uri>")
97     parser.add_argument("-r", "--samplerate",
98             metavar = "<freq>", type=int,
99             action="store", dest="samplerate", default=0,
100             help="samplerate at which the file should be represented")
101
102 def parser_add_verbose_help(parser):
103     parser.add_argument("-v","--verbose",
104             action="count", dest="verbose", default=1,
105             help="make lots of noise [default]")
106     parser.add_argument("-q","--quiet",
107             action="store_const", dest="verbose", const=0,
108             help="be quiet")
109
110 def parser_add_buf_hop_size(parser, buf_size=512, hop_size=256):
111     parser.add_argument("-B","--bufsize",
112             action="store", dest="buf_size", default=buf_size,
113             metavar = "<size>", type=int,
114             help="buffer size [default=%d]" % buf_size)
115     parser.add_argument("-H","--hopsize",
116             metavar = "<size>", type=int,
117             action="store", dest="hop_size", default=hop_size,
118             help="overlap size [default=%d]" % hop_size)
119
120 def parser_add_method(parser, method='default', helpstr='method'):
121     parser.add_argument("-m","--method",
122             metavar = "<method>", type=str,
123             action="store", dest="method", default=method,
124             help="%s [default=%s]" % (helpstr, method))
125
126 def parser_add_threshold(parser, default=None):
127     parser.add_argument("-t","--threshold",
128             metavar = "<threshold>", type=float,
129             action="store", dest="threshold", default=default,
130             help="threshold [default=%s]" % default)
131
132 def parser_add_silence(parser):
133     parser.add_argument("-s", "--silence",
134             metavar = "<value>", type=float,
135             action="store", dest="silence", default=-70,
136             help="silence threshold")
137
138 def parser_add_minioi(parser):
139     parser.add_argument("-M", "--minioi",
140             metavar = "<value>", type=str,
141             action="store", dest="minioi", default="12ms",
142             help="minimum Inter-Onset Interval")
143
144 def parser_add_time_format(parser):
145     helpstr = "select time values output format (samples, ms, seconds)"
146     helpstr += " [default=seconds]"
147     parser.add_argument("-T", "--time-format",
148              metavar='format',
149              dest="time_format",
150              default=None,
151              help=helpstr)
152
153 # some utilities
154
155 def parse_options(args, valid_opts):
156     options = {k :v for k,v in vars(args).items() if k in valid_opts}
157     return options
158
159 def remap_pvoc_options(options):
160     # remap buf_size to win_s, hop_size to hop_s
161     # FIXME: adjust python/ext/py-phasevoc.c to understand buf_size/hop_size
162     options['win_s'] = options['buf_size']
163     del options['buf_size']
164     options['hop_s'] = options['hop_size']
165     del options['hop_size']
166     return options
167
168 def samples2seconds(n_frames, samplerate):
169     return "%f\t" % (n_frames / float(samplerate))
170
171 def samples2milliseconds(n_frames, samplerate):
172     return "%f\t" % (1000. * n_frames / float(samplerate))
173
174 def samples2samples(n_frames, samplerate):
175     return "%d\t" % n_frames
176
177 def timefunc(mode):
178     if mode is None or mode == 'seconds' or mode == 's':
179         return samples2seconds
180     elif mode == 'ms' or mode == 'milliseconds':
181         return samples2milliseconds
182     elif mode == 'samples':
183         return samples2samples
184     else:
185         raise ValueError('invalid time format %s' % mode)
186
187 # definition of processing classes
188
189 class default_process(object):
190     def __init__(self, args):
191         if 'time_format' in args:
192             self.time2string = timefunc(args.time_format)
193         if args.verbose > 2 and hasattr(self, 'options'):
194             name = type(self).__name__.split('_')[1]
195             optstr = ' '.join(['running', name, 'with options', repr(self.options), '\n'])
196             sys.stderr.write(optstr)
197     def flush(self, n_frames, samplerate):
198         pass
199
200 class process_onset(default_process):
201     valid_opts = ['method', 'hop_size', 'buf_size', 'samplerate']
202     def __init__(self, args):
203         self.options = parse_options(args, self.valid_opts)
204         self.onset = aubio.onset(**self.options)
205         if args.threshold is not None:
206             self.onset.set_threshold(args.threshold)
207         if args.minioi:
208             if args.minioi.endswith('ms'):
209                 self.onset.set_minioi_ms(float(args.minioi[:-2]))
210             elif args.minioi.endswith('s'):
211                 self.onset.set_minioi_s(float(args.minioi[:-1]))
212             else:
213                 self.onset.set_minioi(int(args.minioi))
214         if args.silence:
215             self.onset.set_silence(args.silence)
216         super(process_onset, self).__init__(args)
217     def __call__(self, block):
218         return self.onset(block)
219     def repr_res(self, res, frames_read, samplerate):
220         if res[0] != 0:
221             outstr = self.time2string(self.onset.get_last(), samplerate)
222             sys.stdout.write(outstr + '\n')
223
224 class process_pitch(default_process):
225     valid_opts = ['method', 'hop_size', 'buf_size', 'samplerate']
226     def __init__(self, args):
227         self.options = parse_options(args, self.valid_opts)
228         self.pitch = aubio.pitch(**self.options)
229         if args.threshold is not None:
230             self.pitch.set_tolerance(args.threshold)
231         if args.silence is not None:
232             self.pitch.set_silence(args.silence)
233         super(process_pitch, self).__init__(args)
234     def __call__(self, block):
235         return self.pitch(block)
236     def repr_res(self, res, frames_read, samplerate):
237         fmt_out = self.time2string(frames_read, samplerate)
238         sys.stdout.write(fmt_out + "%.6f\n" % res[0])
239
240 class process_tempo(default_process):
241     valid_opts = ['method', 'hop_size', 'buf_size', 'samplerate']
242     def __init__(self, args):
243         self.options = parse_options(args, self.valid_opts)
244         self.tempo = aubio.tempo(**self.options)
245         super(process_tempo, self).__init__(args)
246     def __call__(self, block):
247         return self.tempo(block)
248     def repr_res(self, res, frames_read, samplerate):
249         if res[0] != 0:
250             outstr = self.time2string(self.tempo.get_last(), samplerate)
251             sys.stdout.write(outstr + '\n')
252
253 class process_notes(default_process):
254     valid_opts = ['method', 'hop_size', 'buf_size', 'samplerate']
255     def __init__(self, args):
256         self.options = parse_options(args, self.valid_opts)
257         self.notes = aubio.notes(**self.options)
258         super(process_notes, self).__init__(args)
259     def __call__(self, block):
260         return self.notes(block)
261     def repr_res(self, res, frames_read, samplerate):
262         if res[2] != 0: # note off
263             fmt_out = self.time2string(frames_read, samplerate)
264             sys.stdout.write(fmt_out + '\n')
265         if res[0] != 0: # note on
266             lastmidi = res[0]
267             fmt_out = "%f\t" % lastmidi
268             fmt_out += self.time2string(frames_read, samplerate)
269             sys.stdout.write(fmt_out) # + '\t')
270     def flush(self, frames_read, samplerate):
271         eof = self.time2string(frames_read, samplerate)
272         sys.stdout.write(eof + '\n')
273
274 class process_mfcc(default_process):
275     def __init__(self, args):
276         valid_opts = ['hop_size', 'buf_size']
277         options = parse_options(args, valid_opts)
278         self.options = remap_pvoc_options(options)
279         self.pv = aubio.pvoc(**options)
280
281         valid_opts = ['buf_size', 'n_filters', 'n_coeffs', 'samplerate']
282         options = parse_options(args, valid_opts)
283         self.mfcc = aubio.mfcc(**options)
284         self.options.update(options)
285
286         super(process_mfcc, self).__init__(args)
287
288     def __call__(self, block):
289         fftgrain = self.pv(block)
290         return self.mfcc(fftgrain)
291     def repr_res(self, res, frames_read, samplerate):
292         fmt_out = self.time2string(frames_read, samplerate)
293         fmt_out += ' '.join(["% 9.7f" % f for f in res.tolist()])
294         sys.stdout.write(fmt_out + '\n')
295
296 class process_melbands(default_process):
297     def __init__(self, args):
298         self.args = args
299         valid_opts = ['hop_size', 'buf_size']
300         options = parse_options(args, valid_opts)
301         options = remap_pvoc_options(options)
302         self.pv = aubio.pvoc(**options)
303
304         valid_opts = ['buf_size', 'n_filters']
305         options = {k :v for k,v in vars(args).items() if k in valid_opts}
306         # FIXME
307         options['win_s'] = options['buf_size']
308         del options['buf_size']
309         self.filterbank = aubio.filterbank(**options)
310         self.filterbank.set_mel_coeffs_slaney(args.samplerate)
311
312         super(process_melbands, self).__init__(args)
313     def __call__(self, block):
314         fftgrain = self.pv(block)
315         return self.filterbank(fftgrain)
316     def repr_res(self, res, frames_read, samplerate):
317         fmt_out = self.time2string(frames_read, samplerate)
318         fmt_out += ' '.join(["% 9.7f" % f for f in res.tolist()])
319         sys.stdout.write(fmt_out + '\n')
320
321 def main():
322     parser = aubio_parser()
323     args = parser.parse_args()
324     if args.show_version or ('verbose' in args and args.verbose > 3):
325         sys.stdout.write('aubio version ' + aubio.version + '\n')
326     if args.show_version and args.command is None:
327         sys.exit(0)
328     if args.command is None:
329         sys.stderr.write("Error: a command is required\n")
330         parser.print_help()
331         sys.exit(1)
332     elif not args.source_uri and not args.source_uri2:
333         sys.stderr.write("Error: a source is required\n")
334         parser.print_help()
335         sys.exit(1)
336     elif args.source_uri2 is not None:
337         args.source_uri = args.source_uri2
338     try:
339         # open source_uri
340         with aubio.source(args.source_uri, hop_size=args.hop_size,
341                 samplerate=args.samplerate) as a_source:
342             args.samplerate = a_source.samplerate
343             # create the processor for this subcommand
344             processor = args.process(args)
345             frames_read = 0
346             while True:
347                 # read new block from source
348                 block, read = a_source()
349                 # execute processor on this block
350                 res = processor(block)
351                 # print results for this block
352                 if args.verbose > 0:
353                     processor.repr_res(res, frames_read, a_source.samplerate)
354                 # increment total number of frames read
355                 frames_read += read
356                 # exit loop at end of file
357                 if read < a_source.hop_size: break
358             # flush the processor if needed
359             processor.flush(frames_read, a_source.samplerate)
360             if args.verbose > 1:
361                 fmt_string = "read {:.2f}s"
362                 fmt_string += " ({:d} samples in {:d} blocks of {:d})"
363                 fmt_string += " from {:s} at {:d}Hz\n"
364                 sys.stderr.write(fmt_string.format(
365                         frames_read/float(a_source.samplerate),
366                         frames_read,
367                         frames_read // a_source.hop_size + 1,
368                         a_source.hop_size,
369                         a_source.uri,
370                         a_source.samplerate))
371     except KeyboardInterrupt as e:
372         sys.exit(1)