78d419884bc2ae669667bdd8120721fe68db230f
[aubio.git] / python / lib / aubio / slicing.py
1 from aubio import source, sink
2 import os
3
4 max_timestamp = 1e120
5
6 def slice_source_at_stamps(source_file, timestamps, timestamps_end = None,
7         output_dir = None,
8         samplerate = 0,
9         hopsize = 256):
10
11     if timestamps == None or len(timestamps) == 0:
12         raise ValueError ("no timestamps given")
13
14     if timestamps[0] != 0:
15         timestamps = [0] + timestamps
16
17     if timestamps_end != None and len(timestamps_end) != len(timestamps):
18         raise ValueError ("len(timestamps_end) != len(timestamps)")
19     else:
20         timestamps_end = [t - 1 for t in timestamps[1:] ] + [ max_timestamp ]
21
22     source_base_name, source_ext = os.path.splitext(os.path.basename(source_file))
23     if output_dir != None:
24         if not os.path.isdir(output_dir):
25             os.makedirs(output_dir)
26         source_base_name = os.path.join(output_dir, source_base_name)
27
28     def new_sink_name(source_base_name, timestamp, samplerate):
29         timestamp_seconds = timestamp / float(samplerate)
30         #print source_base_name + '_%02.3f' % (timestamp_seconds) + '.wav'
31         return source_base_name + '_%02.3f' % (timestamp_seconds) + '.wav'
32
33     # reopen source file
34     s = source(source_file, samplerate, hopsize)
35     if samplerate == 0: samplerate = s.get_samplerate()
36     total_frames = 0
37     # get next region
38     start_stamp = int(timestamps.pop(0))
39     end_stamp = int(timestamps_end.pop(0))
40
41     # create first sink
42     new_sink_path = new_sink_name(source_base_name, start_stamp, samplerate)
43     #print "new slice", total_frames, "+", remaining, "=", end_stamp
44     g = sink(new_sink_path, samplerate)
45
46     while True:
47         # get hopsize new samples from source
48         vec, read = s()
49         # number of samples until end of region
50         remaining = end_stamp - total_frames
51         # not enough frames remaining, time to split
52         if remaining < read:
53             if remaining != 0:
54                 # write remaining samples from current region
55                 g(vec[0:remaining], remaining)
56             # close this file
57             del g
58             # get the next region
59             start_stamp = int(timestamps.pop(0))
60             end_stamp = int(timestamps_end.pop(0))
61             # create a new file for the new region
62             new_sink_path = new_sink_name(source_base_name, start_stamp, samplerate)
63             #print "new slice", total_frames, "+", remaining, "=", end_stamp
64             g = sink(new_sink_path, samplerate)
65             # write the remaining samples in the new file
66             g(vec[remaining:read], read - remaining)
67         elif read > 0:
68             # write all the samples
69             g(vec[0:read], read)
70         total_frames += read
71         if read < hopsize: break
72
73     # close the last file
74     del g