a2s_project/prepare_slakh_subset_data.py at master · eneugut/a2s_project · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
from new_processing import *
import os
from convert_midi import midi_to_abc

def get_song_length(wav_file_path):
    song = AudioSegment.from_file(os.getcwd()+wav_file_path, format="wav")
    return(len(song)/1000)

def calculate_number_of_subsongs():
    pass

def divide_and_export_as_pkl(wav_file_path, midi_file_path, duration):
    song_length = get_song_length(wav_file_path)
    sections = int(song_length/duration)
    #print("There are "+str(sections)+" sections")

    sound_tensors = []
    abc_files = []
    temp_midi_path = re.sub(r"\.mid",r"_temp.mid", midi_file_path)
    for i in range(sections):
        start = i*duration
        end = (i+1)*duration

        # First, we process the audio
        sound_tensor = subset_audio_segment(wav_file_path,start,end)
        sound_tensors.append(sound_tensor)

        # Second, we process the transcript
        midi_subset = subset_midi_segment(midi_file_path,start,end)
        midi_subset.write(os.getcwd()+temp_midi_path)
        abc_file = midi_to_abc(temp_midi_path)
        abc_files.append(abc_file)
        os.remove(os.getcwd()+temp_midi_path)

    # Now combine and export as pickle
    number = re.findall(r'\d+', wav_file_path) # Find number from file name0
    number = [x for x in number if x != '2100'] # Remove

    type = ""
    if re.search("train", wav_file_path) is not None:
        type = "train"
    if re.search("validation", wav_file_path) is not None:
        type = "validation"
    if re.search("test", wav_file_path) is not None:
        type = "test"

    for i in range(len(sound_tensors)):
        pickle_file_path = r"/slakh-data/processed-slakh/"+type+r"/p"+number[0]+r'-'+f'{i:03d}'+r'.pkl'
        output_pickle((sound_tensors[i],abc_files[i]), pickle_file_path)