HE/data.py at master · morgan76/HE · 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
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
import random
import numpy as np
import torch
from librosa.util import find_files
import os
from features import get_features
import input_output as io


class CondTripletDset(torch.utils.data.Dataset):


    def __init__(self, experiment_config, split):

        self.ds_path = experiment_config.ds_path
        self.feat_id = experiment_config.feat_id
        self.split = split
        self.experiment_config = experiment_config
        self.tracklist = find_files(os.path.join(self.ds_path, 'audio'), ext=self.experiment_config.dataset.audio_exts)

        if self.split == 'valid':
            d_path = '../msaf_/datasets/BeatlesTUT'
            self.tracklist = find_files(os.path.join(d_path, 'audio'), ext=self.experiment_config.dataset.audio_exts)[:10]

        if self.experiment_config.n_conditions == 1:
            self.triplets = self.buildTriplets(self.tracklist, len(self.tracklist), self.experiment_config.batch_size)
        else:
            self.triplets = self.buildTriplets(self.tracklist, len(self.tracklist), 10)
        print('Number of tracks/triplets for split', split, ':', len(self.tracklist), len(self.triplets))


    def buildTriplets(self, tracklist, n_songs, n_per_song):

        random.seed(None)
        triplets = []

        tracks = np.random.choice(tracklist, size=n_songs, replace=False)

        for track in tracks:
            file_struct = io.FileStruct(track)
            beat_frames = io.read_beats(file_struct.beat_file)
            nb_embeddings = len(beat_frames)
            triplets_temp = []

            if nb_embeddings > 128:

                if self.experiment_config.n_conditions == 1:
                    cs = [0]
                elif self.experiment_config.n_conditions == 4:
                    cs = [0, 1, 2, 3] * 3

                anchor_indexes = np.random.choice(np.arange(0, nb_embeddings, 1), size=n_per_song, replace=False)

                for i in range(len(anchor_indexes)):
                    anchor_index = anchor_indexes[i]
                    for k in range(len(cs)):
                        c = cs[k]
                        if c == 0:
                            positive_index, negative_index = self.sampler(anchor_index, nb_embeddings, c)
                            negative_temp = positive_index
                        else:
                            positive_index, negative_index = self.sampler(anchor_index, nb_embeddings, c)
                            negative_index = negative_temp
                            negative_temp = positive_index

                        assert positive_index != negative_index
                        anchor_index_= beat_frames[anchor_index]
                        positive_index_ = beat_frames[positive_index]
                        negative_index_ = beat_frames[negative_index]

                        triplets_temp.append((track, c, anchor_index_, positive_index_, negative_index_))

                if self.experiment_config.n_conditions == 1:
                    assert len(triplets_temp) == n_per_song
                else:
                    assert len(triplets_temp) == self.experiment_config.batch_size

                random.shuffle(triplets_temp)
                triplets+=triplets_temp

        return triplets


    def sampler(self, anchor_index, nb_embeddings, c, delta_p = 16, delta_n_min = 1, delta_n_max = 96):
        # Randomly samples positive and negative anchors from beat frames
        L = int(nb_embeddings)

        if self.experiment_config.n_conditions == 1:
            delta_p_min = 1
            delta_p = 16
            delta_n_min = 1
            delta_n_max = 128


        elif self.experiment_config.n_conditions == 4:
            if c == 0:
                delta_p_min = 48
                delta_p = 64
                delta_n_min = 64
                delta_n_max = 128
            elif c == 1:
               delta_p_min = 32
               delta_p = 48
               delta_n_min = 40
               delta_n_max = 48
            elif c == 2:
               delta_p_min = 16
               delta_p = 32
               delta_n_min = 32
               delta_n_max = 40
            elif c == 3:
               delta_p_min = 1
               delta_p = 16
               delta_n_min = 24
               delta_n_max = 32


        total_positive = list(np.arange(max(anchor_index-delta_p, 0), max(anchor_index-delta_p_min, 0))) + list(np.arange(min(anchor_index+delta_p_min, L-1), min(anchor_index+delta_p, L-1)))
        positive_index = random.choice(total_positive)


        total_negative = list(np.arange(max(anchor_index-delta_n_max, 0), max(anchor_index-delta_n_min, 0))) + list(np.arange(min(anchor_index+delta_n_min, L-1), min(anchor_index+delta_n_max, L-1)))
        if self.experiment_config.n_conditions == 1:
            total_negative.remove(positive_index)
        negative_index = random.choice(total_negative)

        return positive_index, negative_index


    def __getitem__(self, index):
        track, c, anchor_index_, positive_index_, negative_index_ = self.triplets[index]
        features = get_features(track, self.feat_id, self.experiment_config)

        features = (features-np.min(features))/(np.max(features)-np.min(features))
        features_padded = np.pad(features,
                    pad_width=((0, 0),
                                (int(self.experiment_config.embedding.n_embedding//2),
                                int(self.experiment_config.embedding.n_embedding//2))
                                ),
                    mode='edge')

        anchor_patch = torch.tensor(features_padded[:, anchor_index_:anchor_index_+self.experiment_config.embedding.n_embedding])[None, :, :]
        positive_patch = torch.tensor(features_padded[:, positive_index_:positive_index_+self.experiment_config.embedding.n_embedding])[None, :, :]
        negative_patch = torch.tensor(features_padded[:, negative_index_:negative_index_+self.experiment_config.embedding.n_embedding])[None, :, :]

        return anchor_patch, negative_patch, positive_patch, c


    def __len__(self):
        return len(self.triplets)