<commit_before>{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filenames = [ './librivox/guidetomen_%02d_rowland_64kb.mp3' % (i,) for i in [1,2,3]]\n",
    "filenames"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###  Build dataset from a series of mp3 files\n",
    "\n",
    "For each file in turn, pull out blocks of spectra each 1024 spectra long (~12sec).\n",
    "First 64 of these will be discarded, so 'step increment' should be (1024-64=960)\n",
    "Ignore tail block.\n",
    "\n",
    "https://medium.com/@chengweizhang2012/an-easy-guide-to-build-new-tensorflow-datasets-and-estimator-with-keras-model-9b0f6b4c1b0d\n",
    "\n",
    "https://github.com/Tony607/Keras_catVSdog_tf_estimator/blob/master/keras_estimator_vgg16-cat_vs_dog-TFRecord.ipynb\n",
    "\n",
    "https://www.tensorflow.org/programmers_guide/datasets\n",
    "\n",
    "http://warmspringwinds.github.io/tensorflow/tf-slim/2016/12/21/tfrecords-guide/\n",
    "\n",
    "https://indico.io/blog/tensorflow-data-inputs-part1-placeholders-protobufs-queues/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "import librosa\n",
    "librosa.__version__  # '0.5.1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_rate= 24000 # input will be standardised to this rate\n",
    "\n",
    "fft_step   = 12.5/1000. # 12.5ms\n",
    "fft_window = 50.0/1000.  # 50ms\n",
    "\n",
    "n_fft = 512*4\n",
    "\n",
    "hop_length = int(fft_step*sample_rate)\n",
    "win_length = int(fft_window*sample_rate)\n",
    "\n",
    "n_mels = 80\n",
    "fmin = 125 # Hz\n",
    "#fmax = ~8000\n",
    "\n",
    "#np.exp(-7.0), np.log(spectra_abs_min)  # \"Audio tests\" suggest a min log of -6 or -8 is fine\n",
    "spectra_abs_min = 0.01 # From Google paper, seems justified\n",
    "\n",
    "win_length, hop_length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# And for the training windowing :\n",
    "steps_total, steps_leadin = 1024, 64\n",
    "\n",
    "# Test the flatten idea\n",
    "#a = np.array([3.4, 55.4, 34.23])\n",
    "a = np.array([[3.4, 55.4,],[34.23, 342.1221]])\n",
    "a.flatten().tolist()\n",
    "\n",
    "#np.angle( 0.-1.j)\n",
    "#np.angle([1.0, 1.0j, 1+1j, -1.0j])        \n",
    "\n",
    "# Test angle re-range idea\n",
    "#_sa = 4.0\n",
    "#np.where(_sa<=np.pi, _sa, _sa-2*np.pi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Based on http://warmspringwinds.github.io/tensorflow/tf-slim/2016/12/21/tfrecords-guide/\n",
    "\n",
    "#def _int64_feature(value):\n",
    "#    return tf.train.Feature(int64_list=tf.train.Int64List( value=[value] ))\n",
    "def _floats_feature(np_arr):\n",
    "    return tf.train.Feature(float_list=tf.train.FloatList( value=np_arr.flatten().tolist() ))\n",
    "\n",
    "def convert_wavs_to_spectra_learnable_records(filename_in):\n",
    "    print(\"convert_wavs_to_spectra_learnable_records(%s)\" % (filename_in,))\n",
    "    filename_base = filename_in.replace('.mp3', '_%s.tfrecords')\n",
    "\n",
    "    samples, _sample_rate = librosa.core.load(filename_in, sr=sample_rate)\n",
    "    samples = samples/np.max(samples)  # Force amplitude of waveform into range ~-1 ... +1.0\n",
    "\n",
    "    spectra_complex = librosa.stft(samples, n_fft=n_fft, \n",
    "                       hop_length=hop_length, \n",
    "                       win_length=win_length, window='hann', )\n",
    "\n",
    "    power_spectra = np.abs(spectra_complex)**2\n",
    "    melspectra = librosa.feature.melspectrogram(S=power_spectra, n_mels=n_mels, fmin=fmin)\n",
    "\n",
    "    with tf.python_io.TFRecordWriter(filename_base % ('train',)) as writer_train, \\\n",
    "         tf.python_io.TFRecordWriter(filename_base % ('valid',)) as writer_valid, \\\n",
    "         tf.python_io.TFRecordWriter(filename_base % ('test',)) as writer_test :\n",
    "                \n",
    "        # Ok, now create a series of Examples with these features\n",
    "        for offset in range(0, melspectra.shape[1]-steps_total, steps_total-steps_leadin):\n",
    "            mel_offset     = melspectra[:, offset:offset+steps_total]   \n",
    "            spectra_offset = spectra_complex[:, offset:offset+steps_total]\n",
    "\n",
    "            # Now do some useful precalculation, instead of relying on TF later\n",
    "            mel_log        = np.log( np.maximum(spectra_abs_min, np.abs(mel_offset) ))\n",
    "            spectra_l_amp  = np.log( np.maximum(spectra_abs_min, np.abs(spectra_offset) ))\n",
    "\n",
    "            #_sa = np.angle(spectra_offset)  # These are [0 .. 2pi) - convert to (-pi .. +pi]\n",
    "            #spectra_angles = np.where(_sa<=np.pi, _sa, _sa-2*np.pi)\n",
    "            \n",
    "            # See : https://github.com/numpy/numpy/issues/7426 (documentation is wrong)\n",
    "            spectra_angles = np.angle(spectra_offset)  # These are (-pi .. +pi]\n",
    "            \n",
    "            spectra_phase0 = spectra_angles[:, steps_leadin-1:steps_leadin]  # Initial angles\n",
    "            spectra_pshift = spectra_angles[:, 1:] - spectra_angles[:, :-1] # Differences\n",
    "            \n",
    "            spectra_target = np.concatenate( (\n",
    "                spectra_l_amp[:, steps_leadin:],\n",
    "                spectra_pshift[:, (steps_leadin-1):], # This shifted along by due to differencing \n",
    "            ), axis=0)  # 1025+1025\n",
    "            #print(mel_offset.shape, spectra_phase0.shape, spectra_pshift.shape, spectra_target.shape)\n",
    "        \n",
    "            example = tf.train.Example(features=tf.train.Features(feature={\n",
    "                # NB: Tensorflow wants the transposed versions to match (batch, T, channels)\n",
    "                #'mel': _floats_feature( mel_offset.T ),\n",
    "                #'spectra_real': _floats_feature( spectra_offset.real.T ),\n",
    "                #'spectra_imag': _floats_feature( spectra_offset.imag.T ),\n",
    "                \n",
    "                # These are the ones required below...\n",
    "                'mel_log': _floats_feature( mel_log.T ),\n",
    "                'spectra_phase0': _floats_feature( spectra_phase0.T ),\n",
    "                'spectra_target': _floats_feature( spectra_target.T ),\n",
    "                #'height': _int64_feature(height),\n",
    "                #'mask_raw': _bytes_feature(annotation_raw)\n",
    "            }))\n",
    "\n",
    "            w = writer_train  # Allocate these between the various train/validation/test files\n",
    "            if np.random.random()>0.8:\n",
    "                w = writer_valid\n",
    "                if np.random.random()>0.5:\n",
    "                    w = writer_test\n",
    "\n",
    "            w.write(example.SerializeToString())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for f in filenames:\n",
    "    convert_wavs_to_spectra_learnable_records(f)\n",
    "print(\"DONE!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Consume the TFrecord data\n",
    "\n",
    "*  https://www.tensorflow.org/programmers_guide/datasets\n",
    "*  https://github.com/Tony607/Keras_catVSdog_tf_estimator/blob/master/keras_estimator_vgg16-cat_vs_dog-TFRecord.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mel_bins, spectra_bins = n_mels, n_fft//2+1 # 80, 1025\n",
    "\n",
    "batch_size, num_epochs = 8, 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mel_to_complex_dataset_from_mp3(filenames, stub='train'):\n",
    "    dataset = tf.data.TFRecordDataset([f.replace('.mp3', '_%s.tfrecords') % stub \n",
    "                                       for f in filenames])\n",
    "\n",
    "    spectra_len = (steps_total-steps_leadin)\n",
    "    features = {\n",
    "      #\"mel\":          tf.FixedLenFeature([mel_bins*steps_total], tf.float32),\n",
    "      #\"spectra_real\": tf.FixedLenFeature([spectra_bins*steps_total], tf.float32),\n",
    "      #\"spectra_imag\": tf.FixedLenFeature([spectra_bins*steps_total], tf.float32),\n",
    "      \"mel_log\":        tf.FixedLenFeature([steps_total*mel_bins], tf.float32),\n",
    "      \"spectra_phase0\": tf.FixedLenFeature([1*spectra_bins], tf.float32),\n",
    "      \"spectra_target\": tf.FixedLenFeature([spectra_len*spectra_bins*2], tf.float32),\n",
    "    }\n",
    "    \n",
    "    def _parse_function_OLD(example_proto):\n",
    "        parsed_features = tf.parse_single_example(example_proto, features)\n",
    "        \n",
    "        mel          = tf.reshape( parsed_features[\"mel\"], (mel_bins, steps_total) ) \n",
    "        \n",
    "        spectra_real = tf.reshape( parsed_features[\"spectra_real\"], (spectra_bins, steps_total) ) \n",
    "        spectra_imag = tf.reshape( parsed_features[\"spectra_imag\"], (spectra_bins, steps_total) ) \n",
    "        \n",
    "        spectra_real = tf.transpose( spectra_real )\n",
    "        spectra_imag = tf.transpose( spectra_imag )\n",
    "        \n",
    "        spectra_complex = tf.complex(spectra_real, spectra_imag)\n",
    "        #spectra_amp   = tf.norm( spectra_complex )\n",
    "        \n",
    "        spectra_amp   = tf.sqrt( tf.square(spectra_real) + tf.square(spectra_imag) )\n",
    "        spectra_log_amp = tf.log(  tf.maximum(0.00001, spectra_amp) )\n",
    "        \n",
    "        spectra_phase = tf.angle(spectra_complex)\n",
    "        #return dict(MelInput=tf.transpose(mel)), (tf.transpose(spectra_log_amp), tf.transpose(spectra_phase))\n",
    "        \n",
    "        spectra_concat = tf.concat( [spectra_log_amp, spectra_phase ], axis=1 )\n",
    "        return dict(MelInput=tf.transpose(mel)), spectra_concat\n",
    "    \n",
    "        #return dict(MelInput=tf.transpose(mel)), (tf.transpose(spectra_real), tf.transpose(spectra_imag))\n",
    "        #spectra = tf.stack( [spectra_real, spectra_imag, ] ) # Should be (spectra_bins, steps_total, 2)\n",
    "        #return dict(MelInput=tf.transpose(mel)), tf.transpose(spectra, perm=[1,0,2])\n",
    "\n",
    "    def _parse_function(example_proto):\n",
    "        parsed_features = tf.parse_single_example(example_proto, features)\n",
    "        \n",
    "        mel_log    = tf.reshape( parsed_features[\"mel_log\"], (steps_total, mel_bins) ) \n",
    "        spectra_phase0 = tf.reshape( parsed_features[\"spectra_phase0\"], (1, spectra_bins) ) \n",
    "        spectra_target = tf.reshape( parsed_features[\"spectra_target\"], (spectra_len, spectra_bins*2) ) \n",
    "        \n",
    "        return dict(MelInput=mel_log, Phase0=spectra_phase0), spectra_target\n",
    "\n",
    "    dataset = dataset.map(_parse_function)\n",
    "    return dataset\n",
    "\n",
    "def input_fn_from(filenames, stub='train', batch_size=1, shuffle=False, repeats=1):\n",
    "    dataset = mel_to_complex_dataset_from_mp3(filenames, stub=stub)\n",
    "    if shuffle:\n",
    "        dataset = dataset.shuffle(buffer_size=100)\n",
    "    dataset = dataset.batch(batch_size).repeat(repeats)\n",
    "    \n",
    "    iterator = dataset.make_one_shot_iterator()    \n",
    "    \n",
    "    batch_features, batch_labels = iterator.get_next()\n",
    "    return batch_features, batch_labels    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Keras model \n",
    "\n",
    "Keras-WaveNet : \n",
    "\n",
    "\n",
    "*  https://github.com/usernaamee/keras-wavenet/blob/master/simple-generative-model-regressor.py \n",
    "*  Beware of GPL3 license!  Code independently...\n",
    "*  But, also, the number of convolutional channels seems pretty arbitrary, so may not be fleshed-out yet\n",
    "\n",
    "TF-WaveNet :\n",
    "*  https://github.com/ibab/tensorflow-wavenet\n",
    "*  MIT license : Feel free to look + adapt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# See : https://github.com/tensorflow/tensorflow/issues/14933\n",
    "#   to understand how broken Google is\n",
    "\n",
    "from tensorflow.python import keras\n",
    "from tensorflow.python.keras import backend as K\n",
    "\n",
    "#tf.reset_default_graph()\n",
    "\n",
    "# Use 'real keras' to get the actual documented functionality for padding='causal'\n",
    "#import keras\n",
    "#from keras import backend as K\n",
    "\n",
    "def wavenet_layer(channels, hidden_channels, kernel_size, dilation_rate, name):\n",
    "    def f(input_):\n",
    "        filter_out = keras.layers.Conv1D(hidden_channels, kernel_size,\n",
    "                                       strides=1, dilation_rate=dilation_rate,\n",
    "                                       padding='valid', use_bias=True, \n",
    "                                       activation='tanh', name='filter_'+name)(input_)\n",
    "        gate_out   = keras.layers.Conv1D(hidden_channels, kernel_size,\n",
    "                                       strides=1, dilation_rate=dilation_rate,\n",
    "                                       padding='valid', use_bias=True, \n",
    "                                       activation='sigmoid', name='gate_'+name)(input_)\n",
    "        mult = keras.layers.Multiply(name='mult_'+name)( [filter_out, gate_out] )\n",
    "        \n",
    "        # Need to pad this result back out to input_ size...\n",
    "        #print(dilation_rate, kernel_size, dilation_rate*(kernel_size-1))\n",
    "        \n",
    "        #def original_shape(input_shape):\n",
    "        #    return (input_shape[0], input_shape[1]+dilation_rate*kernel_size-2, input_shape[2])\n",
    "        # \n",
    "        #mult_padded = keras.layers.Lambda(\n",
    "        #    lambda x: K.temporal_padding(x, padding=(dilation_rate*kernel_size-1,0) ), \n",
    "        #    #output_shape=original_shape,\n",
    "        #    name='mult_padded_'+name)(mult)\n",
    "        \n",
    "        # https://www.tensorflow.org/api_docs/python/tf/keras/layers/ZeroPadding1D\n",
    "        mult_padded = keras.layers.ZeroPadding1D( (dilation_rate*(kernel_size-1), 0) )(mult)\n",
    "\n",
    "        transformed = keras.layers.Conv1D(channels, 1, \n",
    "                                          padding='same', use_bias=True, \n",
    "                                          activation='linear', name='trans_'+name)(mult_padded)\n",
    "        skip_out    = keras.layers.Conv1D(channels, 1, \n",
    "                                          padding='same', use_bias=True, \n",
    "                                          activation='relu', name='skip_'+name)(mult_padded)\n",
    "        \n",
    "        return keras.layers.Add(name='resid_'+name)( [transformed, input_] ), skip_out\n",
    "      \n",
    "    return f\n",
    "\n",
    "#log_amplitude_with_minimum = keras.layers.Lambda( lambda x: K.log( K.maximum(0.00001, x) ))\n",
    "\n",
    "io_channels, hidden_channels = 128,128\n",
    "def model_mel_to_spec( input_shape=(steps_total, mel_bins) ):\n",
    "    #mel = keras.layers.Input(shape=input_shape, name='MelInput')\n",
    "    #mel = tf.keras.layers.Input(shape=input_shape, name='MelInput')\n",
    "    #mel = tf.keras.layers.Input(batch_size=batch_size, shape=input_shape, name='MelInput')\n",
    "    #mel = keras.layers.Input(batch_size=batch_size, shape=input_shape, name='MelInput')\n",
    "    #mel = keras.layers.Input(shape=input_shape, name='MelInput',\n",
    "    #                        _batch_input_shape = (batch_size, steps_total, mel_bins))\n",
    "    \n",
    "    #mel = keras.layers.Input(batch_shape=(batch_size, steps_total, mel_bins), name='MelInput')\n",
    "    #mel._batch_input_shape = (batch_size, steps_total, mel_bins)\n",
    "    #mel = keras.layers.InputLayer(input_shape=input_shape, name='MelInput')\n",
    "    \n",
    "    #mel_floored = K.maximum(0.00001, mel)\n",
    "    #mel_log     = K.log(mel_floored)  # This is (batch, T. channels)\n",
    "    \n",
    "    #mel_log = log_amplitude_with_minimum(mel)\n",
    "    \n",
    "    mel_log = keras.layers.Input(shape=input_shape, name='MelInput')\n",
    "    phase0  = keras.layers.Input(shape=input_shape, name='Phase0')  # Unused\n",
    "\n",
    "    # 'Resize' to make everything 'io_channels' big at the layer interfaces\n",
    "    x = s0 = keras.layers.Conv1D(io_channels, 1, \n",
    "                          padding='same', use_bias=True, \n",
    "                          activation='linear', name='mel_log_expanded')(mel_log)\n",
    "    \n",
    "    x,s1 = wavenet_layer(io_channels, hidden_channels*1, 3, 1, '1')(x)\n",
    "    x,s2 = wavenet_layer(io_channels, hidden_channels*1, 3, 2, '2')(x)\n",
    "    x,s3 = wavenet_layer(io_channels, hidden_channels*1, 3, 4, '3')(x)\n",
    "    x,s4 = wavenet_layer(io_channels, hidden_channels*1, 3, 8, '4')(x)\n",
    "    _,s5 = wavenet_layer(io_channels, hidden_channels*1, 3,16, '5')(x)  # Total footprint is ~64 0.75secs\n",
    "    # x is now irrelevant\n",
    "    \n",
    "    skip_overall = keras.layers.Concatenate( axis=-1 )( [s0,s1,s2,s3,s4,s5] )\n",
    "    \n",
    "    log_amp     = keras.layers.Conv1D(spectra_bins, 1, padding='same', \n",
    "                                  activation='linear', name='log_amp')(skip_overall)\n",
    "    phase_shift = keras.layers.Conv1D(spectra_bins, 1, padding='same', \n",
    "                                  activation='linear', name='phase_shift')(skip_overall)\n",
    "    \n",
    "    #return keras.models.Model(inputs=[mel], outputs=[log_amp, phase])\n",
    "\n",
    "    #amp = K.exp(log_amp)\n",
    "    #amp = keras.layers.Lambda( lambda x: K.exp(x), name='amp')(log_amp)\n",
    "    #return keras.models.Model(inputs=[mel], outputs=[log_amp, phase])\n",
    "\n",
    "    #spec_real = keras.layers.Multiply()( [amp, K.cos(phase)] )\n",
    "    #spec_imag = keras.layers.Multiply()( [amp, K.sin(phase)] )\n",
    "    #return keras.models.Model(inputs=mel, outputs=[spec_real, spec_imag])\n",
    "    \n",
    "    log_amp_valid     = keras.layers.Cropping1D( (steps_leadin,0), name='crop_a' )( log_amp )\n",
    "    phase_shift_valid = keras.layers.Cropping1D( (steps_leadin,0), name='crop_p' )( phase_shift )\n",
    "    \n",
    "    # Concat the amps and phases into one return value\n",
    "    spec_concat = keras.layers.Concatenate( axis=-1, name='spec_concat')( \n",
    "        [log_amp_valid, phase_shift_valid] )\n",
    "    return keras.models.Model(inputs=[mel_log, phase0], outputs=spec_concat)\n",
    "    \n",
    "\n",
    "keras_model = model_mel_to_spec()\n",
    "keras_model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def customLoss(spec_gold, spec_out):\n",
    "    gold_l_amp = keras.layers.Lambda(lambda x : x[:,:,:spectra_bins])(spec_gold)\n",
    "    gold_phase = keras.layers.Lambda(lambda x : x[:,:,spectra_bins:])(spec_gold)\n",
    "    \n",
    "    spec_l_amp = keras.layers.Lambda(lambda x : x[:,:,:spectra_bins])(spec_out)\n",
    "    spec_phase = keras.layers.Lambda(lambda x : x[:,:,spectra_bins:])(spec_out)\n",
    "    \n",
    "    l_amp_loss = keras.losses.mean_squared_error( gold_l_amp, spec_l_amp )\n",
    "    \n",
    "    #phase_diff = K.abs( gold_phase - spec_phase )  # gold are  (-pi ... pi), spec are anywhere\n",
    "    #phase_diff = K.minimum(phase_diff, 2.0*np.pi - phase_diff) # Positive by construction\n",
    "    #phase_loss = K.mean( phase_diff )\n",
    "    #phase_loss = K.mean( -K.cos( gold_phase - spec_phase ) ) # This 'fades out' near equality\n",
    "    phase_loss = keras.losses.mean_squared_error( gold_phase, spec_phase )\n",
    "    \n",
    "    return l_amp_loss + 1.0 * phase_loss\n",
    "\n",
    "keras_model.compile(loss=customLoss, \n",
    "                    optimizer=keras.optimizers.RMSprop(),  # lr=2e-5\n",
    "                    metrics=['mse'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now make it an Estimator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "model_dir = os.path.join(os.getcwd(), 'models', 'mel-to-complex-spectra_04')\n",
    "os.makedirs(model_dir, exist_ok=True)\n",
    "print(\"model_dir: \",model_dir)\n",
    "\n",
    "estimator = tf.keras.estimator.model_to_estimator(keras_model=keras_model, model_dir=model_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Just check that the input name for our model matches what we have in the DataSet reader\n",
    "\n",
    "input_name = keras_model.input_names[0]\n",
    "input_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "train_spec = tf.estimator.TrainSpec(\n",
    "    #input_fn=lambda: imgs_input_fn(path_tfrecords_train, perform_shuffle=True,\n",
    "    #                                repeat_count=5, batch_size=20), \n",
    "    input_fn=lambda: input_fn_from(filenames, stub='train', shuffle=True,\n",
    "                                    repeats=num_epochs, batch_size=batch_size), \n",
    "    )\n",
    "    #max_steps=500)\n",
    "\n",
    "eval_spec = tf.estimator.EvalSpec(\n",
    "    #input_fn=lambda: imgs_input_fn(path_tfrecords_test, perform_shuffle=False, batch_size=1)\n",
    "    input_fn=lambda: input_fn_from(filenames, stub='valid', shuffle=False,\n",
    "                                    repeats=1, batch_size=1), \n",
    "    )\n",
    "\n",
    "import time\n",
    "start_time = time.time()\n",
    "\n",
    "tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)\n",
    "\n",
    "print(\"--- %s seconds ---\" % (time.time() - start_time))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predict_results = estimator.predict(\n",
    "    input_fn=lambda: input_fn_from(filenames, stub='test', shuffle=False,\n",
    "                                    repeats=1, batch_size=1), \n",
    "    )\n",
    "predictions = [p for p in predict_results]\n",
    "len(predictions), predictions[0].keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now check some results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from IPython.display import Audio as audio_playback_widget\n",
    "import soundfile  # For audio_widget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_single_prediction(pred, phase0=None, stub='orig'):\n",
    "    # These are (probably) a numpy array(s)\n",
    "    #print(pred.shape, phase0.shape)\n",
    "    spec_l_amp = pred[:, :spectra_bins]\n",
    "    spec_phase_diff = pred[:, spectra_bins:]\n",
    "    \n",
    "    #print(spec_phase_diff.mean(), spec_phase_diff.std(), )\n",
    "    #spec_l_amp = np.maximum(-6., spec_l_amp)\n",
    "    \n",
    "    if phase0 is None: \n",
    "        phase0 = np.zeros( (1, spectra_bins) )\n",
    "    \n",
    "    #spec_phase_diff += 1.0*np.random.normal( size=spec_phase_diff.shape ) \n",
    "    \n",
    "    # Find cumulative phase amounts\n",
    "    spec_phase = np.cumsum( spec_phase_diff, axis=0 )\n",
    "    \n",
    "    # How sensitive is this to noise?\n",
    "    spec_phase += 1.0*np.random.normal( size=spec_phase.shape ) \n",
    "    \n",
    "    spec_complex = np.exp(spec_l_amp + 1j*(phase0+spec_phase)).T\n",
    "    #spec_complex = np.exp(spec_l_amp + 1j*0.00001*(phase0+spec_phase)).T\n",
    "    \n",
    "    #spectrum = np.log( np.absolute( spec_complex )[:, 0:400] )\n",
    "    spectrum = np.log( np.absolute( spec_complex ) )\n",
    "    \n",
    "    fig, ax = plt.subplots(nrows=1,ncols=1, figsize=(20,4))\n",
    "    cax = ax.matshow(spectrum, interpolation='nearest', aspect='auto', cmap=plt.cm.afmhot, origin='lower')\n",
    "    fig.colorbar(cax)\n",
    "    plt.title(\"Predicted Spectrogram\")\n",
    "    plt.show()\n",
    "    \n",
    "    #plt.figure(figsize=(10, 4))\n",
    "    #librosa.display.specshow(melout[:, 0:400], y_axis='mel', fmin=125, x_axis='time')\n",
    "    #plt.colorbar(format='%+2.0f dB')\n",
    "    #plt.title('Mel spectrogram')\n",
    "    #plt.show()    \n",
    "    samples = librosa.istft(spec_complex, hop_length=hop_length, win_length=win_length)\n",
    "    \n",
    "    f = './tmp/single_%s.wav' % (stub,)\n",
    "    soundfile.write(f, samples/np.max(samples), samplerate=sample_rate)\n",
    "    return audio_playback_widget(f)    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read some test data, and just show it 'raw'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_test_direct = mel_to_complex_dataset_from_mp3(filenames, stub='test')\n",
    "ds_test_iterator = ds_test_direct.make_one_shot_iterator()    \n",
    "ds_test_iterator_next = ds_test_iterator.get_next()\n",
    "\n",
    "with tf.Session() as sess:\n",
    "    first_batch = sess.run(ds_test_iterator_next)\n",
    "\n",
    "spec_orig = first_batch[1] # This is the concat( l_amp and phase )\n",
    "\n",
    "spec_l_amp = spec_orig[:, :spectra_bins]\n",
    "spec_phase_diff = spec_orig[:, spectra_bins:]\n",
    "\n",
    "np.mean(spec_l_amp[:, 100]), np.std(spec_l_amp[:, 100]), np.min(spec_l_amp[:, 100]),"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spec_concat_orig = first_batch[1]\n",
    "show_single_prediction( spec_concat_orig, first_batch[0]['Phase0'] )\n",
    "#show_single_prediction( spec_concat_orig, None )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spec_concat_pred = predictions[0]['spec_concat'].copy()\n",
    "#show_single_prediction( spec_concat_pred, None, 'net' )\n",
    "\n",
    "# Add in original phases (but keep predicted spectra amplitudes)\n",
    "#spec_concat_pred[:, spectra_bins:] = spec_concat_orig[:, spectra_bins:]\n",
    "\n",
    "show_single_prediction( spec_concat_pred, first_batch[0]['Phase0'], 'net' )\n",
    "#show_single_prediction( spec_concat_pred, None, 'net' )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",<commit_msg>Batchnorm still doesn't fix scaling of angles issue...<commit_after>{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filenames = [ './librivox/guidetomen_%02d_rowland_64kb.mp3' % (i,) for i in [1,2,3]]\n",
    "filenames"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###  Build dataset from a series of mp3 files\n",
    "\n",
    "For each file in turn, pull out blocks of spectra each 1024 spectra long (~12sec).\n",
    "First 64 of these will be discarded, so 'step increment' should be (1024-64=960)\n",
    "Ignore tail block.\n",
    "\n",
    "https://medium.com/@chengweizhang2012/an-easy-guide-to-build-new-tensorflow-datasets-and-estimator-with-keras-model-9b0f6b4c1b0d\n",
    "\n",
    "https://github.com/Tony607/Keras_catVSdog_tf_estimator/blob/master/keras_estimator_vgg16-cat_vs_dog-TFRecord.ipynb\n",
    "\n",
    "https://www.tensorflow.org/programmers_guide/datasets\n",
    "\n",
    "http://warmspringwinds.github.io/tensorflow/tf-slim/2016/12/21/tfrecords-guide/\n",
    "\n",
    "https://indico.io/blog/tensorflow-data-inputs-part1-placeholders-protobufs-queues/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "\n",
    "import librosa\n",
    "librosa.__version__  # '0.5.1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_rate= 24000 # input will be standardised to this rate\n",
    "\n",
    "fft_step   = 12.5/1000. # 12.5ms\n",
    "fft_window = 50.0/1000.  # 50ms\n",
    "\n",
    "n_fft = 512*4\n",
    "\n",
    "hop_length = int(fft_step*sample_rate)\n",
    "win_length = int(fft_window*sample_rate)\n",
    "\n",
    "n_mels = 80\n",
    "fmin = 125 # Hz\n",
    "#fmax = ~8000\n",
    "\n",
    "#np.exp(-7.0), np.log(spectra_abs_min)  # \"Audio tests\" suggest a min log of -6 or -8 is fine\n",
    "spectra_abs_min = 0.01 # From Google paper, seems justified\n",
    "\n",
    "win_length, hop_length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# And for the training windowing :\n",
    "steps_total, steps_leadin = 1024, 64\n",
    "\n",
    "# Test the flatten idea\n",
    "#a = np.array([3.4, 55.4, 34.23])\n",
    "a = np.array([[3.4, 55.4,],[34.23, 342.1221]])\n",
    "a.flatten().tolist()\n",
    "\n",
    "#np.angle( 0.-1.j)\n",
    "#np.angle([1.0, 1.0j, 1+1j, -1.0j])        \n",
    "\n",
    "# Test angle re-range idea\n",
    "#_sa = 4.0\n",
    "#np.where(_sa<=np.pi, _sa, _sa-2*np.pi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Based on http://warmspringwinds.github.io/tensorflow/tf-slim/2016/12/21/tfrecords-guide/\n",
    "\n",
    "#def _int64_feature(value):\n",
    "#    return tf.train.Feature(int64_list=tf.train.Int64List( value=[value] ))\n",
    "def _floats_feature(np_arr):\n",
    "    return tf.train.Feature(float_list=tf.train.FloatList( value=np_arr.flatten().tolist() ))\n",
    "\n",
    "def convert_wavs_to_spectra_learnable_records(filename_in):\n",
    "    print(\"convert_wavs_to_spectra_learnable_records(%s)\" % (filename_in,))\n",
    "    filename_base = filename_in.replace('.mp3', '_%s.tfrecords')\n",
    "\n",
    "    samples, _sample_rate = librosa.core.load(filename_in, sr=sample_rate)\n",
    "    samples = samples/np.max(samples)  # Force amplitude of waveform into range ~-1 ... +1.0\n",
    "\n",
    "    spectra_complex = librosa.stft(samples, n_fft=n_fft, \n",
    "                       hop_length=hop_length, \n",
    "                       win_length=win_length, window='hann', )\n",
    "\n",
    "    power_spectra = np.abs(spectra_complex)**2\n",
    "    melspectra = librosa.feature.melspectrogram(S=power_spectra, n_mels=n_mels, fmin=fmin)\n",
    "\n",
    "    with tf.python_io.TFRecordWriter(filename_base % ('train',)) as writer_train, \\\n",
    "         tf.python_io.TFRecordWriter(filename_base % ('valid',)) as writer_valid, \\\n",
    "         tf.python_io.TFRecordWriter(filename_base % ('test',)) as writer_test :\n",
    "                \n",
    "        # Ok, now create a series of Examples with these features\n",
    "        for offset in range(0, melspectra.shape[1]-steps_total, steps_total-steps_leadin):\n",
    "            mel_offset     = melspectra[:, offset:offset+steps_total]   \n",
    "            spectra_offset = spectra_complex[:, offset:offset+steps_total]\n",
    "\n",
    "            # Now do some useful precalculation, instead of relying on TF later\n",
    "            mel_log        = np.log( np.maximum(spectra_abs_min, np.abs(mel_offset) ))\n",
    "            spectra_l_amp  = np.log( np.maximum(spectra_abs_min, np.abs(spectra_offset) ))\n",
    "\n",
    "            #_sa = np.angle(spectra_offset)  # These are [0 .. 2pi) - convert to (-pi .. +pi]\n",
    "            #spectra_angles = np.where(_sa<=np.pi, _sa, _sa-2*np.pi)\n",
    "            \n",
    "            # See : https://github.com/numpy/numpy/issues/7426 (documentation is wrong)\n",
    "            spectra_angles = np.angle(spectra_offset)  # These are (-pi .. +pi]\n",
    "            \n",
    "            spectra_phase0 = spectra_angles[:, steps_leadin-1:steps_leadin]  # Pre-initial angles\n",
    "            spectra_pshift = spectra_angles[:, 1:] - spectra_angles[:, :-1]  # Differences\n",
    "            \n",
    "            spectra_target = np.concatenate( (\n",
    "                spectra_l_amp[:, steps_leadin:],\n",
    "                spectra_pshift[:, (steps_leadin-1):], # This shifted along by due to differencing \n",
    "            ), axis=0)  # 1025+1025\n",
    "            #print(mel_offset.shape, spectra_phase0.shape, spectra_pshift.shape, spectra_target.shape)\n",
    "        \n",
    "            example = tf.train.Example(features=tf.train.Features(feature={\n",
    "                # NB: Tensorflow wants the transposed versions to match (batch, T, channels)\n",
    "                #'mel': _floats_feature( mel_offset.T ),\n",
    "                #'spectra_real': _floats_feature( spectra_offset.real.T ),\n",
    "                #'spectra_imag': _floats_feature( spectra_offset.imag.T ),\n",
    "                \n",
    "                # These are the ones required below...\n",
    "                'mel_log': _floats_feature( mel_log.T ),\n",
    "                'spectra_phase0': _floats_feature( spectra_phase0.T ),\n",
    "                'spectra_target': _floats_feature( spectra_target.T ),\n",
    "                #'height': _int64_feature(height),\n",
    "                #'mask_raw': _bytes_feature(annotation_raw)\n",
    "            }))\n",
    "\n",
    "            w = writer_train  # Allocate these between the various train/validation/test files\n",
    "            if np.random.random()>0.8:\n",
    "                w = writer_valid\n",
    "                if np.random.random()>0.5:\n",
    "                    w = writer_test\n",
    "\n",
    "            w.write(example.SerializeToString())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for f in filenames:\n",
    "    convert_wavs_to_spectra_learnable_records(f)\n",
    "print(\"DONE!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Consume the TFrecord data\n",
    "\n",
    "*  https://www.tensorflow.org/programmers_guide/datasets\n",
    "*  https://github.com/Tony607/Keras_catVSdog_tf_estimator/blob/master/keras_estimator_vgg16-cat_vs_dog-TFRecord.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mel_bins, spectra_bins = n_mels, n_fft//2+1 # 80, 1025\n",
    "\n",
    "batch_size, num_epochs = 8, 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mel_to_complex_dataset_from_mp3(filenames, stub='train'):\n",
    "    dataset = tf.data.TFRecordDataset([f.replace('.mp3', '_%s.tfrecords') % stub \n",
    "                                       for f in filenames])\n",
    "\n",
    "    spectra_len = (steps_total-steps_leadin)\n",
    "    features = {\n",
    "      #\"mel\":          tf.FixedLenFeature([mel_bins*steps_total], tf.float32),\n",
    "      #\"spectra_real\": tf.FixedLenFeature([spectra_bins*steps_total], tf.float32),\n",
    "      #\"spectra_imag\": tf.FixedLenFeature([spectra_bins*steps_total], tf.float32),\n",
    "      \"mel_log\":        tf.FixedLenFeature([steps_total*mel_bins], tf.float32),\n",
    "      \"spectra_phase0\": tf.FixedLenFeature([1*spectra_bins], tf.float32),\n",
    "      \"spectra_target\": tf.FixedLenFeature([spectra_len*spectra_bins*2], tf.float32),\n",
    "    }\n",
    "    \n",
    "    def _parse_function_OLD(example_proto):\n",
    "        parsed_features = tf.parse_single_example(example_proto, features)\n",
    "        \n",
    "        mel          = tf.reshape( parsed_features[\"mel\"], (mel_bins, steps_total) ) \n",
    "        \n",
    "        spectra_real = tf.reshape( parsed_features[\"spectra_real\"], (spectra_bins, steps_total) ) \n",
    "        spectra_imag = tf.reshape( parsed_features[\"spectra_imag\"], (spectra_bins, steps_total) ) \n",
    "        \n",
    "        spectra_real = tf.transpose( spectra_real )\n",
    "        spectra_imag = tf.transpose( spectra_imag )\n",
    "        \n",
    "        spectra_complex = tf.complex(spectra_real, spectra_imag)\n",
    "        #spectra_amp   = tf.norm( spectra_complex )\n",
    "        \n",
    "        spectra_amp   = tf.sqrt( tf.square(spectra_real) + tf.square(spectra_imag) )\n",
    "        spectra_log_amp = tf.log(  tf.maximum(0.00001, spectra_amp) )\n",
    "        \n",
    "        spectra_phase = tf.angle(spectra_complex)\n",
    "        #return dict(MelInput=tf.transpose(mel)), (tf.transpose(spectra_log_amp), tf.transpose(spectra_phase))\n",
    "        \n",
    "        spectra_concat = tf.concat( [spectra_log_amp, spectra_phase ], axis=1 )\n",
    "        return dict(MelInput=tf.transpose(mel)), spectra_concat\n",
    "    \n",
    "        #return dict(MelInput=tf.transpose(mel)), (tf.transpose(spectra_real), tf.transpose(spectra_imag))\n",
    "        #spectra = tf.stack( [spectra_real, spectra_imag, ] ) # Should be (spectra_bins, steps_total, 2)\n",
    "        #return dict(MelInput=tf.transpose(mel)), tf.transpose(spectra, perm=[1,0,2])\n",
    "\n",
    "    def _parse_function(example_proto):\n",
    "        parsed_features = tf.parse_single_example(example_proto, features)\n",
    "        \n",
    "        mel_log    = tf.reshape( parsed_features[\"mel_log\"], (steps_total, mel_bins) ) \n",
    "        spectra_phase0 = tf.reshape( parsed_features[\"spectra_phase0\"], (1, spectra_bins) ) \n",
    "        spectra_target = tf.reshape( parsed_features[\"spectra_target\"], (spectra_len, spectra_bins*2) ) \n",
    "        \n",
    "        return dict(MelInput=mel_log, Phase0=spectra_phase0), spectra_target\n",
    "\n",
    "    dataset = dataset.map(_parse_function)\n",
    "    return dataset\n",
    "\n",
    "def input_fn_from(filenames, stub='train', batch_size=1, shuffle=False, repeats=1):\n",
    "    dataset = mel_to_complex_dataset_from_mp3(filenames, stub=stub)\n",
    "    if shuffle:\n",
    "        dataset = dataset.shuffle(buffer_size=100)\n",
    "    dataset = dataset.batch(batch_size).repeat(repeats)\n",
    "    \n",
    "    iterator = dataset.make_one_shot_iterator()    \n",
    "    \n",
    "    batch_features, batch_labels = iterator.get_next()\n",
    "    return batch_features, batch_labels    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Keras model \n",
    "\n",
    "Keras-WaveNet : \n",
    "\n",
    "\n",
    "*  https://github.com/usernaamee/keras-wavenet/blob/master/simple-generative-model-regressor.py \n",
    "*  Beware of GPL3 license!  Code independently...\n",
    "*  But, also, the number of convolutional channels seems pretty arbitrary, so may not be fleshed-out yet\n",
    "\n",
    "TF-WaveNet :\n",
    "*  https://github.com/ibab/tensorflow-wavenet\n",
    "*  MIT license : Feel free to look + adapt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# See : https://github.com/tensorflow/tensorflow/issues/14933\n",
    "#   to understand how broken Google is\n",
    "\n",
    "from tensorflow.python import keras\n",
    "from tensorflow.python.keras import backend as K\n",
    "\n",
    "#tf.reset_default_graph()\n",
    "\n",
    "# Use 'real keras' to get the actual documented functionality for padding='causal'\n",
    "#import keras\n",
    "#from keras import backend as K\n",
    "\n",
    "def wavenet_layer(channels, hidden_channels, kernel_size, dilation_rate, name):\n",
    "    def f(input_):\n",
    "        filter_out = keras.layers.Conv1D(hidden_channels, kernel_size,\n",
    "                                       strides=1, dilation_rate=dilation_rate,\n",
    "                                       padding='valid', use_bias=True, \n",
    "                                       activation='tanh', name='filter_'+name)(input_)\n",
    "        gate_out   = keras.layers.Conv1D(hidden_channels, kernel_size,\n",
    "                                       strides=1, dilation_rate=dilation_rate,\n",
    "                                       padding='valid', use_bias=True, \n",
    "                                       activation='sigmoid', name='gate_'+name)(input_)\n",
    "        mult = keras.layers.Multiply(name='mult_'+name)( [filter_out, gate_out] )\n",
    "        \n",
    "        # Need to pad this result back out to input_ size...\n",
    "        #print(dilation_rate, kernel_size, dilation_rate*(kernel_size-1))\n",
    "        \n",
    "        #def original_shape(input_shape):\n",
    "        #    return (input_shape[0], input_shape[1]+dilation_rate*kernel_size-2, input_shape[2])\n",
    "        # \n",
    "        #mult_padded = keras.layers.Lambda(\n",
    "        #    lambda x: K.temporal_padding(x, padding=(dilation_rate*kernel_size-1,0) ), \n",
    "        #    #output_shape=original_shape,\n",
    "        #    name='mult_padded_'+name)(mult)\n",
    "        \n",
    "        # https://www.tensorflow.org/api_docs/python/tf/keras/layers/ZeroPadding1D\n",
    "        mult_padded = keras.layers.ZeroPadding1D( (dilation_rate*(kernel_size-1), 0) )(mult)\n",
    "\n",
    "        transformed = keras.layers.Conv1D(channels, 1, \n",
    "                                          padding='same', use_bias=True, \n",
    "                                          activation='linear', name='trans_'+name)(mult_padded)\n",
    "        skip_out    = keras.layers.Conv1D(channels, 1, \n",
    "                                          padding='same', use_bias=True, \n",
    "                                          activation='relu', name='skip_'+name)(mult_padded)\n",
    "        \n",
    "        return keras.layers.Add(name='resid_'+name)( [transformed, input_] ), skip_out\n",
    "      \n",
    "    return f\n",
    "\n",
    "#log_amplitude_with_minimum = keras.layers.Lambda( lambda x: K.log( K.maximum(0.00001, x) ))\n",
    "\n",
    "io_channels, hidden_channels = 128,128\n",
    "def model_mel_to_spec( input_shape=(steps_total, mel_bins) ):\n",
    "    #mel = keras.layers.Input(shape=input_shape, name='MelInput')\n",
    "    #mel = tf.keras.layers.Input(shape=input_shape, name='MelInput')\n",
    "    #mel = tf.keras.layers.Input(batch_size=batch_size, shape=input_shape, name='MelInput')\n",
    "    #mel = keras.layers.Input(batch_size=batch_size, shape=input_shape, name='MelInput')\n",
    "    #mel = keras.layers.Input(shape=input_shape, name='MelInput',\n",
    "    #                        _batch_input_shape = (batch_size, steps_total, mel_bins))\n",
    "    \n",
    "    #mel = keras.layers.Input(batch_shape=(batch_size, steps_total, mel_bins), name='MelInput')\n",
    "    #mel._batch_input_shape = (batch_size, steps_total, mel_bins)\n",
    "    #mel = keras.layers.InputLayer(input_shape=input_shape, name='MelInput')\n",
    "    \n",
    "    #mel_floored = K.maximum(0.00001, mel)\n",
    "    #mel_log     = K.log(mel_floored)  # This is (batch, T. channels)\n",
    "    \n",
    "    #mel_log = log_amplitude_with_minimum(mel)\n",
    "    \n",
    "    mel_log = keras.layers.Input(shape=input_shape, name='MelInput')\n",
    "    phase0  = keras.layers.Input(shape=input_shape, name='Phase0')  # Unused\n",
    "    \n",
    "    x = keras.layers.BatchNormalization()(mel_log)\n",
    "\n",
    "    # 'Resize' to make everything 'io_channels' big at the layer interfaces\n",
    "    x = s0 = keras.layers.Conv1D(io_channels, 1, \n",
    "                          padding='same', use_bias=True, \n",
    "                          activation='linear', name='mel_log_expanded')(x)\n",
    "    \n",
    "    x,s1 = wavenet_layer(io_channels, hidden_channels*1, 3, 1, '1')(x)\n",
    "    x,s2 = wavenet_layer(io_channels, hidden_channels*1, 3, 2, '2')(x)\n",
    "    x,s3 = wavenet_layer(io_channels, hidden_channels*1, 3, 4, '3')(x)\n",
    "    x,s4 = wavenet_layer(io_channels, hidden_channels*1, 3, 8, '4')(x)\n",
    "    _,s5 = wavenet_layer(io_channels, hidden_channels*1, 3,16, '5')(x)  # Total footprint is ~64 0.75secs\n",
    "    # x is now irrelevant\n",
    "    \n",
    "    #skip_overall = keras.layers.Concatenate( axis=-1 )( [s0,s1,s2,s3,s4,s5] )\n",
    "    skip_overall = keras.layers.Concatenate( axis=-1 )( [s0,s0] )\n",
    "    \n",
    "    log_amp     = keras.layers.Conv1D(spectra_bins, 1, padding='same', \n",
    "                                  activation='linear', name='log_amp')(skip_overall)\n",
    "    phase_shift = keras.layers.Conv1D(spectra_bins, 1, padding='same', \n",
    "                                  activation='linear', name='phase_shift')(skip_overall)\n",
    "    \n",
    "    #return keras.models.Model(inputs=[mel], outputs=[log_amp, phase])\n",
    "\n",
    "    #amp = K.exp(log_amp)\n",
    "    #amp = keras.layers.Lambda( lambda x: K.exp(x), name='amp')(log_amp)\n",
    "    #return keras.models.Model(inputs=[mel], outputs=[log_amp, phase])\n",
    "\n",
    "    #spec_real = keras.layers.Multiply()( [amp, K.cos(phase)] )\n",
    "    #spec_imag = keras.layers.Multiply()( [amp, K.sin(phase)] )\n",
    "    #return keras.models.Model(inputs=mel, outputs=[spec_real, spec_imag])\n",
    "    \n",
    "    log_amp_valid     = keras.layers.Cropping1D( (steps_leadin,0), name='crop_a' )( log_amp )\n",
    "    phase_shift_valid = keras.layers.Cropping1D( (steps_leadin,0), name='crop_p' )( phase_shift )\n",
    "    \n",
    "    # Concat the amps and phases into one return value\n",
    "    spec_concat = keras.layers.Concatenate( axis=-1, name='spec_concat')( \n",
    "        [log_amp_valid, phase_shift_valid] )\n",
    "    return keras.models.Model(inputs=[mel_log, phase0], outputs=spec_concat)\n",
    "    \n",
    "\n",
    "keras_model = model_mel_to_spec()\n",
    "keras_model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def customLoss(spec_gold, spec_out):\n",
    "    gold_l_amp = keras.layers.Lambda(lambda x : x[:,:,:spectra_bins])(spec_gold)\n",
    "    gold_phase = keras.layers.Lambda(lambda x : x[:,:,spectra_bins:])(spec_gold)\n",
    "    \n",
    "    spec_l_amp = keras.layers.Lambda(lambda x : x[:,:,:spectra_bins])(spec_out)\n",
    "    spec_phase = keras.layers.Lambda(lambda x : x[:,:,spectra_bins:])(spec_out)\n",
    "    \n",
    "    l_amp_loss = keras.losses.mean_squared_error( gold_l_amp, spec_l_amp )\n",
    "    \n",
    "    #phase_diff = K.abs( gold_phase - spec_phase )  # gold are  (-pi ... pi), spec are anywhere\n",
    "    #phase_diff = K.minimum(phase_diff, 2.0*np.pi - phase_diff) # Positive by construction\n",
    "    #phase_loss = K.mean( phase_diff )\n",
    "    #phase_loss = K.mean( -K.cos( gold_phase - spec_phase ) ) # This 'fades out' near equality\n",
    "    phase_loss = keras.losses.mean_squared_error( gold_phase, spec_phase )\n",
    "    \n",
    "    return l_amp_loss + 1.0 * phase_loss\n",
    "\n",
    "keras_model.compile(loss=customLoss, \n",
    "                    #optimizer=keras.optimizers.RMSprop(),  # lr=2e-5\n",
    "                    optimizer=keras.optimizers.Adam(),  # lr=2e-5\n",
    "                    metrics=['mse'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now make it an Estimator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "model_dir = os.path.join(os.getcwd(), 'models', 'mel-to-complex-spectra_07')\n",
    "os.makedirs(model_dir, exist_ok=True)\n",
    "print(\"model_dir: \",model_dir)\n",
    "\n",
    "estimator = tf.keras.estimator.model_to_estimator(keras_model=keras_model, model_dir=model_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Just check that the input name for our model matches what we have in the DataSet reader\n",
    "\n",
    "input_name = keras_model.input_names[0]\n",
    "input_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "train_spec = tf.estimator.TrainSpec(\n",
    "    #input_fn=lambda: imgs_input_fn(path_tfrecords_train, perform_shuffle=True,\n",
    "    #                                repeat_count=5, batch_size=20), \n",
    "    input_fn=lambda: input_fn_from(filenames, stub='train', shuffle=True,\n",
    "                                    repeats=num_epochs, batch_size=batch_size), \n",
    "    )\n",
    "    #max_steps=500)\n",
    "\n",
    "eval_spec = tf.estimator.EvalSpec(\n",
    "    #input_fn=lambda: imgs_input_fn(path_tfrecords_test, perform_shuffle=False, batch_size=1)\n",
    "    input_fn=lambda: input_fn_from(filenames, stub='valid', shuffle=False,\n",
    "                                    repeats=1, batch_size=1), \n",
    "    )\n",
    "\n",
    "import time\n",
    "start_time = time.time()\n",
    "\n",
    "tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)\n",
    "\n",
    "print(\"--- %s seconds ---\" % (time.time() - start_time))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predict_results = estimator.predict(\n",
    "    input_fn=lambda: input_fn_from(filenames, stub='test', shuffle=False,\n",
    "                                    repeats=1, batch_size=1), \n",
    "    )\n",
    "predictions = [p for p in predict_results]\n",
    "len(predictions), predictions[0].keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Now check some results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from IPython.display import Audio as audio_playback_widget\n",
    "import soundfile  # For audio_widget"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_single_prediction(pred, phase0=None, stub='orig'):\n",
    "    # These are (probably) a numpy array(s)\n",
    "    #print(pred.shape, phase0.shape)\n",
    "    spec_l_amp = pred[:, :spectra_bins]\n",
    "    spec_phase_diff = pred[:, spectra_bins:]\n",
    "    \n",
    "    #print(spec_phase_diff.mean(), spec_phase_diff.std(), )\n",
    "    #spec_l_amp = np.maximum(-6., spec_l_amp)\n",
    "    \n",
    "    if phase0 is None: \n",
    "        phase0 = np.zeros( (1, spectra_bins) )\n",
    "    \n",
    "    #spec_phase_diff += 1.0*np.random.normal( size=spec_phase_diff.shape ) \n",
    "    \n",
    "    # Find cumulative phase amounts\n",
    "    spec_phase = np.cumsum( spec_phase_diff*1.0, axis=0 )\n",
    "    \n",
    "    # How sensitive is this to noise?\n",
    "    #spec_phase += 1.0*np.random.normal( size=spec_phase.shape ) \n",
    "   \n",
    "    #spec_mean = np.mean( spec_phase, axis=0)\n",
    "    #spec_phase = spec_phase - spec_mean\n",
    "    #print(np.max(spec_mean))\n",
    "    \n",
    "    spec_complex = np.exp(spec_l_amp + 1j*(phase0+spec_phase)).T\n",
    "    #spec_complex = np.exp(spec_l_amp + 1j*0.00001*(phase0+spec_phase)).T\n",
    "    \n",
    "    spectrum = np.log( np.absolute( spec_complex ) )\n",
    "    \n",
    "    fig, ax = plt.subplots(nrows=1,ncols=1, figsize=(20,4))\n",
    "    cax = ax.matshow(spectrum, interpolation='nearest', aspect='auto', cmap=plt.cm.afmhot, origin='lower')\n",
    "    fig.colorbar(cax)\n",
    "    plt.title(\"Predicted Spectrogram\")\n",
    "    plt.show()\n",
    "    \n",
    "    #plt.figure(figsize=(10, 4))\n",
    "    #librosa.display.specshow(melout[:, 0:400], y_axis='mel', fmin=125, x_axis='time')\n",
    "    #plt.colorbar(format='%+2.0f dB')\n",
    "    #plt.title('Mel spectrogram')\n",
    "    #plt.show()    \n",
    "    samples = librosa.istft(spec_complex, hop_length=hop_length, win_length=win_length)\n",
    "    \n",
    "    f = './tmp/single_%s.wav' % (stub,)\n",
    "    soundfile.write(f, samples/np.max(samples), samplerate=sample_rate)\n",
    "    return audio_playback_widget(f)    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read some test data, and just show it 'raw'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_test_direct = mel_to_complex_dataset_from_mp3(filenames, stub='test')\n",
    "ds_test_iterator = ds_test_direct.make_one_shot_iterator()    \n",
    "ds_test_iterator_next = ds_test_iterator.get_next()\n",
    "\n",
    "with tf.Session() as sess:\n",
    "    first_batch = sess.run(ds_test_iterator_next)\n",
    "\n",
    "spec_orig = first_batch[1] # This is the concat( l_amp and phase )\n",
    "\n",
    "spec_l_amp = spec_orig[:, :spectra_bins]\n",
    "spec_phase_diff = spec_orig[:, spectra_bins:]\n",
    "\n",
    "np.mean(spec_l_amp[:, 100]), np.std(spec_l_amp[:, 100]), np.min(spec_l_amp[:, 100]),"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spec_concat_true = first_batch[1].copy()\n",
    "show_single_prediction( spec_concat_true, first_batch[0]['Phase0'] )\n",
    "#show_single_prediction( spec_concat_true, None )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spec_concat_pred = predictions[0]['spec_concat'].copy()\n",
    "#show_single_prediction( spec_concat_pred, None, 'net' )\n",
    "\n",
    "# Add in original phases (but keep predicted spectra amplitudes)\n",
    "#spec_concat_pred[:, spectra_bins:] = spec_concat_orig[:, spectra_bins:]\n",
    "\n",
    "show_single_prediction( spec_concat_pred, first_batch[0]['Phase0'], 'net' )\n",
    "#show_single_prediction( spec_concat_pred, None, 'net' )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.set_printoptions(precision=4,suppress=True)\n",
    "(spec_concat_true[50:60, spectra_bins+400:spectra_bins+405], \n",
    " spec_concat_pred[50:60, spectra_bins+400:spectra_bins+405], )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
{'fraction_non_alphanumeric': 0.18548187881568268, 'fraction_numerical': 0.015509757560786252, 'mean_word_length': 2.4532659029978063, 'pattern_counts': {'":': 296, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 22, 'lorem ipsum': 0, 'www.': 6, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 78, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6889366', 'n_tokens_mistral': 21755, 'n_tokens_neox': 20944, 'n_words': 3347}