{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"# !pip install --quiet tensorflow_io\n# !pip install --quiet tensorflow_addons\n!pip install --quiet tensorflow_probability","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"import tensorflow as tf\n# import tensorflow_io as tfio\n# import tensorflow_addons as tfa\nimport tensorflow_probability as tfp\nfrom kaggle_datasets import KaggleDatasets\nimport pandas as pd\nimport numpy as np\nfrom sklearn import model_selection\nimport os\nimport glob\nimport tqdm\nimport math\nimport matplotlib.pyplot as plt\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### TPU stuff"},{"metadata":{"trusted":true},"cell_type":"code","source":"tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\ntf.config.experimental_connect_to_cluster(tpu)\ntf.tpu.experimental.initialize_tpu_system(tpu)\nprint(\"All devices: \", tf.config.list_logical_devices('TPU'))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"strategy = tf.distribute.experimental.TPUStrategy(tpu)\nAUTOTUNE = tf.data.experimental.AUTOTUNE\nGCS_DS_PATH = KaggleDatasets().get_gcs_path('rfcx-species-audio-detection')\nTRAIN_TFREC = GCS_DS_PATH + \"/tfrecords/train\"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 1. Helper functions for TF dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"def _int32(x):\n    return tf.cast(x, 'int32')\n\ndef random_truncated_normal(mean, stddev, min_val, max_val):\n    i = tf.constant( 1.0)\n    x = tf.constant(-1.0)\n    cond = lambda i, x: x<min_val or x>max_val\n    body = lambda i, x: (tf.add(i, 1.0), tf.random.normal([], mean, stddev*i))\n    _, value = tf.while_loop(cond, body, [i, x])\n    return value\n\ndef obtain_onehot_label(labels, t_starts, t_ends, window_start, window_end):\n    onehot_label = tf.zeros(24, dtype='float32')\n    for i in range(len(labels)):\n        species = tf.cast(labels[i], 'int32')\n        t_start = t_starts[i]\n        t_end = t_ends[i]\n        if tf.logical_and(tf.greater(window_end, t_start),\n                          tf.less(window_start, t_end)):\n            onehot_label += tf.one_hot(species, 24)\n    return tf.clip_by_value(onehot_label, 0, 1)\n    \ndef slice_record(features, window_size=5.0):\n    \n    items = tf.strings.split(features['label_info'], sep=',').to_tensor()\n\n    spid = tf.strings.to_number(tf.gather(items, 0, axis=1))\n    tmin = tf.strings.to_number(tf.gather(items, 2, axis=1))\n    tmax = tf.strings.to_number(tf.gather(items, 4, axis=1))\n    \n    idx = tf.random.uniform(\n        shape=(), minval=0, maxval=tf.shape(spid)[0], dtype='int32')\n    mean = (tmax[idx]*features['rate'] + tmin[idx]*features['rate']) / 2\n    \n    window_size = window_size*features['rate']\n    duration = tmax[idx]*features['rate'] - tmin[idx]*features['rate']\n    \n    stddev = tf.math.maximum(duration, window_size) / 2.0 / 3.0\n\n    value = random_truncated_normal(\n        mean, stddev, min_val=window_size/2, max_val=60*features['rate']-window_size)\n    \n    idx_center = _int32(value)\n    idx_start = idx_center - (_int32(window_size) // 2)\n    idx_end  = idx_start + _int32(window_size)\n\n    features['label_info'] = obtain_onehot_label(\n        spid, _int32(tmin*features['rate']), _int32(tmax*features['rate']), idx_start, idx_end)\n    features['audio_wav'] = features['audio_wav'][idx_start: idx_end]\n    \n    return {\n        'signal': features['audio_wav'],\n        'rate': features['rate'],\n        'label_info': features['label_info']\n    }\n\ndef parse_function(example_proto, mode='training'):\n    \n    feature_description = {\n        'recording_id': tf.io.FixedLenFeature([], tf.string, default_value=''),\n        'audio_wav': tf.io.FixedLenFeature([], tf.string, default_value=''),\n        'label_info': tf.io.FixedLenFeature([], tf.string, default_value=''),\n    }\n    example = tf.io.parse_single_example(example_proto, feature_description)\n    wav, rate = tf.audio.decode_wav(example['audio_wav']) # mono\n    example['audio_wav'] = tf.squeeze(tf.cast(wav, 'float32'), axis=-1)\n    example['rate'] = tf.cast(rate, 'float32')\n\n    label_info = tf.strings.split(example['label_info'], sep='\"')[1]\n    label_info = tf.strings.split(label_info, sep=';')\n    example['label_info'] = label_info\n    example = slice_record(example)\n    return example['signal'] # return only signal for the sake of demonstration\n   ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 2. (LogMel)SpecLayer and random augmentation"},{"metadata":{"trusted":true},"cell_type":"code","source":"def random_white_noise(x, p, snr_min=15.0, snr_max=30.0):\n    def _random_snr(x, size, minval, maxval):\n        snr = tf.random.uniform((size, 1), minval, maxval)\n        amplitude = tf.math.reduce_max(tf.math.abs(x))\n        return amplitude / (10**(snr/20))\n    shape = tf.shape(x)\n    snr = _random_snr(x, shape[0], snr_min, snr_max)\n    mul = tf.cast(tfp.distributions.Bernoulli(probs=p).sample((shape[0], 1)), 'float32')\n    noise = tf.random.normal(shape=shape)\n    ampl_noise = tf.math.reduce_max(tf.math.abs(noise))\n    noise = (noise - tf.reduce_mean(noise)) * 1 / ampl_noise * snr * mul\n    return x + noise\n\n\ndef random_volume_control(x, p, db_limit=10, mode='uniform'):\n    shape = tf.cast(tf.shape(x), 'float32')\n    db = tf.random.uniform((tf.cast(shape[0], 'int32'), 1), -db_limit, +db_limit)\n    r = tf.range(shape[1], dtype='float32')\n    r = tf.reshape(tf.tile(r, [shape[0]]), (shape[0], -1))\n    if mode == 'uniform':\n        noise_type = tf.ones(tf.cast(shape, 'int32'), dtype='float32')\n    elif mode == 'fade':\n        noise_type = r[::-1]/(shape[1] - 1)\n    elif mode == 'cosine':\n        noise_type = tf.math.cos(r/shape[1] * math.pi * 2)\n    elif mode == 'sine':\n        noise_type = tf.math.sin(r/shape[1] * math.pi * 2)\n    else:\n        raise ValueError(\n            \"mode has to be either 'uniform', 'fade', 'cosine' or 'sine'\")\n    mul = tf.cast(tfp.distributions.Bernoulli(probs=p).sample((tf.cast(shape[0], 'int32'), 1)), 'float32')\n    noise = 10 ** (db * noise_type / 20 * mul)\n    return x * noise\n\ndef random_time_mask(input, p, max_mask=50):\n    batch_size = tf.shape(input)[0]\n    time_max = tf.cast(tf.repeat(tf.shape(input)[1], [batch_size]), 'float32')\n    time_min = tf.cast(tf.repeat(tf.constant(0), [batch_size]), 'float32')\n    mask_max = tf.cast(tf.repeat(tf.constant(max_mask), [batch_size]), 'float32')\n\n    t = tfp.distributions.Uniform(low=time_min, high=mask_max).sample()\n    t0 = tfp.distributions.Uniform(low=time_min, high=time_max-t).sample()\n\n    mul = tf.cast(tfp.distributions.Bernoulli(probs=p).sample((batch_size)), 'float32')\n    t = t * mul\n    indices = tf.reshape(tf.range(time_max[0]), (1, -1, 1))\n    condition = tf.math.logical_and(\n        tf.math.greater_equal(indices, t0), tf.math.less(indices, t0 + t)\n    )\n    condition = tf.transpose(condition, (2, 1, 0))\n    zero = tf.constant(0, dtype=input.dtype)\n    return tf.where(condition, zero, input)\n\ndef random_freq_mask(input, p, max_mask=16):\n    batch_size = tf.shape(input)[0]\n    freq_max = tf.cast(tf.repeat(tf.shape(input)[2], [batch_size]), 'float32')\n    freq_min = tf.cast(tf.repeat(tf.constant(0), [batch_size]), 'float32')\n    mask_max = tf.cast(tf.repeat(tf.constant(max_mask), [batch_size]), 'float32')\n\n    f = tfp.distributions.Uniform(low=freq_min, high=mask_max).sample()\n    f0 = tfp.distributions.Uniform(low=freq_min, high=freq_max-f).sample()\n\n    mul = tf.cast(tfp.distributions.Bernoulli(probs=p).sample((batch_size)), 'float32')\n    f = f * mul\n    indices = tf.reshape(tf.range(freq_max[0]), (1, -1, 1))\n    condition = tf.math.logical_and(\n        tf.math.greater_equal(indices, f0), tf.math.less(indices, f0 + f)\n    )\n    condition = tf.transpose(condition, (2, 0, 1))\n    zero = tf.constant(0, dtype=input.dtype)\n    return tf.where(condition, zero, input)\n\ndef random_brightness(input, p, max_delta=0.2):\n    batch_size = tf.shape(input)[0]\n    delta = tfp.distributions.Uniform(low=-max_delta, high=max_delta).sample((batch_size, 1, 1))\n    mul = tf.cast(tfp.distributions.Bernoulli(probs=p).sample((batch_size, 1, 1)), 'float32')\n    delta = delta * mul\n    return tf.math.add(input, delta)\n\ndef random_gaussian_noise(input, p, scale=0.2):\n    batch_size = tf.shape(input)[0]\n    shape = tf.shape(input)\n    noise = tfp.distributions.Normal(loc=0, scale=scale).sample(shape)\n    mul = tf.cast(tfp.distributions.Bernoulli(probs=p).sample((batch_size, 1, 1)), 'float32')\n    noise = noise * mul\n    return input + noise\n\n\nclass SpecLayer(tf.keras.layers.Layer):\n\n    def __init__(self,\n                 fft_length=2048,\n                 frame_length=2048,\n                 frame_step=512,\n                 mel_power=2.0,\n                 num_mel_bins=384,\n                 sample_rate=48_000,\n                 lower_edge_hertz=40,\n                 upper_edge_hertz=20_000,\n                 name='spec_layer',\n                 dtype='float32',\n                 **kwargs):\n\n        super(SpecLayer, self).__init__(name=name, dtype=dtype, **kwargs)\n  \n        self.fft_length = fft_length\n        self.frame_length = frame_length\n        self.frame_step = frame_step\n        self.mel_power = mel_power\n        self.mel_matrix = tf.signal.linear_to_mel_weight_matrix(\n            num_mel_bins=num_mel_bins,\n            num_spectrogram_bins=fft_length//2+1,\n            sample_rate=sample_rate,\n            lower_edge_hertz=lower_edge_hertz,\n            upper_edge_hertz=upper_edge_hertz)\n\n\n    def build(self, input_shape):\n        self.non_trainable_weights.append(self.mel_matrix)\n        super(SpecLayer, self).build(input_shape)\n\n    def call(self, inputs, training=False, tta=False):\n\n        x = inputs\n        \n        if training:\n            x = random_volume_control(x, p=0.5, mode='sine')\n            x = random_white_noise(x, p=0.3)\n\n        x = tf.math.abs(\n            tf.signal.stft(\n                x,\n                frame_length=self.frame_length,\n                frame_step=self.frame_step,\n                fft_length=self.fft_length,\n                window_fn=tf.signal.hann_window,\n                pad_end=True\n            )\n        )\n        x = tf.matmul(tf.math.pow(x, self.mel_power), self.mel_matrix)\n        # log10\n        x = tf.math.log(x + 1e-6) / tf.math.log(10.0)\n        x = self.standardize(x)\n\n        if training:\n            x = random_time_mask(x, p=0.5)\n            x = random_freq_mask(x, p=0.5)\n            x = random_brightness(x, p=0.3)\n            x = random_gaussian_noise(x, p=0.2)\n\n        return self.preprocess(x)\n\n\n    @staticmethod\n    def standardize(x):\n        x -= tf.math.reduce_mean(x, [1,2], True)\n        x /= tf.math.reduce_std( x, [1,2], True) + 1e-6\n        return x\n\n    @staticmethod\n    def preprocess(x):\n        x -= tf.math.reduce_min(x, [1,2], True)\n        x /= tf.math.reduce_max(x, [1,2], True)\n        x *= tf.constant(255.)\n        x = tf.transpose(x, (0, 2, 1))\n        return x","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 3. Creating dataset\n\ntfrecords --> waveform signal"},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset = tf.data.TFRecordDataset(\n    filenames=tf.io.gfile.glob(TRAIN_TFREC + '/*.tfrec'),\n    num_parallel_reads=AUTOTUNE)\ndataset = dataset.map(parse_function, AUTOTUNE)\ndataset = dataset.batch(32*strategy.num_replicas_in_sync)\ndataset = dataset.prefetch(AUTOTUNE)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 4. Create (and fit) model\n\nwaveform signal --> log-mel spectrogram<br>\niterating over 4.727 training examples in ~20 sec\n"},{"metadata":{"trusted":true},"cell_type":"code","source":"with strategy.scope():\n    model = tf.keras.Sequential([SpecLayer()])\n    model.compile()\n    model.fit(dataset, epochs=5)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 5. Some output"},{"metadata":{"trusted":true},"cell_type":"code","source":"fig, axes = plt.subplots(2,1, figsize=(10, 10))\n\nwave = next(iter(dataset))\naxes[0].plot(wave.numpy()[0])\nlogmelspec = SpecLayer()(wave)\naxes[1].imshow(logmelspec.numpy()[0].astype(np.uint8))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}