r06944010 / speech-separation-tf2 Goto Github PK

class Libri2Mix():
def __init__(self, config, mode):
    self.mode = mode
    self.config = config
    self.spk = config['training']['n_speaker']
    self.sample_rate = config['dataset']['sample_rate']
    self.path = config['dataset']['path']
    # print(self.path)
    self.batch_size = config['training']['batch_size']
    self.num_stacks = config['model']['num_stacks']
    self.dilations = [2 ** i for i in range(0, config['model']['dilations'] + 1)]

    self.input_length = config['model']['input_length']
    self.num_residual_blocks = len(self.dilations) * self.num_stacks

    if config['training']['path'] == "models/test":
        self.tfr = os.path.join(config['dataset']['path'], 'tfrecord', mode + '_' + 'debug.tfr')
    else:
        self.tfr = os.path.join(config['dataset']['path'], 'tfrecord',
                                str(config['training']['n_speaker']) + 'spk' + '_' + mode + '_' + str(
                                    self.input_length) + '.tfr')

    if not os.path.isfile(self.tfr) or config['dataset']['load_in_mem']:
        print('Find no {}'.format(self.tfr))
        print('[*] Load {} dataset in memory [*]'.format(self.mode))
        self.load_into_memory()

def decode_dataset(self, serialized_example):
    example = tf.parse_single_example(
        serialized_example,
        features={
            "s1": tf.VarLenFeature(tf.float32),
            "s2": tf.VarLenFeature(tf.float32)
        },
    )
    s1 = tf.sparse_tensor_to_dense(example["s1"])
    s2 = tf.sparse_tensor_to_dense(example["s2"])
    audios = tf.stack([s1, s2])
    return audios

def get_iterator(self):
    if os.path.isfile(self.tfr) and not self.config['dataset']['load_in_mem']:
        print("Loading data from \033[93m{} \033[0m".format(self.tfr))
        with tf.name_scope("input"):
            dataset = tf.data.TFRecordDataset(self.tfr).map(self.decode_dataset)
            if self.mode == "tr" or self.mode == "cv":
                dataset = dataset.shuffle(self.batch_size * 100)
            dataset = dataset.batch(self.batch_size, drop_remainder=True)
            dataset = dataset.prefetch(self.batch_size * 10)
            self.iterator = dataset.make_initializable_iterator()
            return self.iterator.get_next()
    else:
        return None

def load_into_memory(self):
    self.file_paths = {}
    self.sequences = {'tr': {'a': [], 'b': []}, 'cv': {'a': [], 'b': []}, 'tt': {'a': [], 'b': []}}
    if self.mode == "tr":
        self.path += "train-100"
    elif self.mode == "cv":
        self.path += "dev"
    for spk in ['a', 'b']:
        if spk == 'a':
            files = os.listdir(self.path + "/s1")
            spk1 = "/s1/"
        else:
            files = os.listdir(self.path + "/s2")
            spk1 = "/s2/"
        sequences = self.load_directory(files, spk1)
        self.sequences[self.mode][spk] = sequences

def load_directory(self, filenames, spk):
    sequences = []
    for filename in tqdm(filenames):
        # print(self.path + spk + filename)
        sequence = util.load_wav(self.path + spk + filename, self.sample_rate)
        length = len(sequence)
        if length < self.input_length:
            continue
        sequences.append(sequence)
    # sequences = np.array(sequences)

    return sequences

def get_next(self):
    n_data = {'tr': 13900, 'cv': 3000, 'tt': 3000}

    indices = np.arange((n_data[self.mode] + self.batch_size - 1) // self.batch_size * self.batch_size)
    indices %= n_data[self.mode]
    np.random.shuffle(indices)

    for i in range(len(indices) // self.batch_size):
        sample_indices = indices[i * self.batch_size:(i + 1) * self.batch_size]
        batch_inputs = []

        for i, sample_i in enumerate(sample_indices):
            speech_a = self.sequences[self.mode]['a'][sample_i]
            speech_b = self.sequences[self.mode]['b'][sample_i]
            offset = np.squeeze(np.random.randint(0, len(speech_a) - self.input_length + 1, 1))
            output_a = speech_a[offset:offset + self.input_length]
            output_b = speech_b[offset:offset + self.input_length]
            batch_inputs.append([output_a, output_b])
        batch_inputs = np.array(batch_inputs, dtype='float32')
        yield batch_inputs, sample_indices