YAMNet-transfer-learning/datagen_yamnet.py at main · jackgle/YAMNet-transfer-learning · 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
import os
import numpy as np
from random import shuffle
import tensorflow as tf

def get_files_and_labels(train_dir, typ='wav', train_split=0.9):

    classes = sorted(os.listdir(train_dir))
    files_train = list()
    files_val = list()
    labels = dict()
    for cnt, i in enumerate(classes): # loop over classes
        tmp = os.listdir(train_dir+i)
        shuffle(tmp)
        for j in tmp[:round(len(tmp)*train_split)]: # loop over training samples
            if j.split('.')[-1]==typ:
                files_train.append(train_dir + i +'/' + j)
        for j in tmp[round(len(tmp)*train_split):]: # loop over validation samples
            if j.split('.')[-1]==typ:
                files_val.append(train_dir + i +'/' + j)
        labels[i]=cnt

    return files_train, files_val, labels


class DataGenerator(tf.keras.utils.Sequence):

    'Generates YAMNet patches'
    def __init__(self,
                 list_IDs,
                 labels,
                 dim = (96, 64),
                 batch_size=1,
                 n_classes=2,
                 shuffle=True,
                 class_weights=None):
        'Initialization'
        self.batch_size = batch_size
        self.labels = labels
        self.dim = dim
        self.list_IDs = list_IDs
        self.n_classes = n_classes
        self.shuffle = shuffle
        self.class_weights = class_weights
        self.on_epoch_end()


    def __len__(self):
        'Denotes the number of batches per epoch'
        return int(np.floor(len(self.list_IDs) / self.batch_size))

    def __getitem__(self, index):
        'Generate one batch of data'
        # Generate indexes of the batch
        indexes = self.indexes[index*self.batch_size:(index+1)*self.batch_size]

        # Find list of IDs
        list_IDs_temp = [self.list_IDs[k] for k in indexes]

        # Generate data
        if self.class_weights:
            X, y, sample_weights = self.__data_generation(list_IDs_temp)
            return X, y, sample_weights
        else:
            X, y = self.__data_generation(list_IDs_temp)
            return X, y

    def on_epoch_end(self):
        'Updates indexes after each epoch'
        self.indexes = np.arange(len(self.list_IDs))
        if self.shuffle == True:
            np.random.shuffle(self.indexes)

    def __data_generation(self, list_IDs_temp):
        'Generates data containing batch_size samples' # X : (n_samples, *dim, n_channels)
        # Initialization

        X = np.empty((self.batch_size, *self.dim))
        y = np.empty((self.batch_size, self.n_classes))
        sample_weights = np.empty((self.batch_size, ))

        y[:] = 0

        # Generate data
        for i, ID in enumerate(list_IDs_temp):

            class_id = ID.split('/')[-2]
            y[i,self.labels[class_id]] = 1

            sample = np.load(ID)

            # if the waveform for this sample was long enough to contain multiple patches, randomly select one of the patches
            if sample.shape[0] > 1:
                sample = np.squeeze(sample[np.random.choice(range(sample.shape[0]), 1)])

            X[i,] = sample

            if self.class_weights:
                sample_weights[i] = self.class_weights[self.labels[class_id]]


        if self.class_weights is not None:
            return X, y, sample_weights
        else:
            return X, y