{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "help(tf.lite.TFLiteConverter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "help(tf.lite.TargetSpec)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_dir = \"/path/to/original/google_bird_classification/google_perch\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "converter = tf.lite.TFLiteConverter.from_saved_model(model_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tflite_model = converter.convert()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The inclusion of the tensorflow ops makes the conversion work, but increases the size of the model substantially. It would be better to have an alternative, perhaps obtainable using the the google perch repo: https://github.com/google-research/perch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "converter.target_spec.supported_ops = [\n",
    "  tf.lite.OpsSet.TFLITE_BUILTINS, # enable TensorFlow Lite ops.\n",
    "  tf.lite.OpsSet.SELECT_TF_OPS # enable TensorFlow ops.\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tflite_model = converter.convert()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(model_dir + \"/converted_model.tflite\", \"wb\") as modelfile: \n",
    "    modelfile.write(tflite_model)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Try out our new model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.signal import frame as tf_split_signal_into_chunks\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tflite_runtime.interpreter as tflite\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_dir = \"/path/to/model/google_bird_classification/google_perch\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "interpreter = tflite.Interpreter(model_path=model_dir + \"/model.tflite\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audiopath = \"/path/to/sample/soundfile/soundscape.wav\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np  \n",
    "import librosa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess(path: str, sample_rate: int, sample_secs: float, overlap: float) -> list:\n",
    "    chunks = []\n",
    "    data, actual_sampling_rate = librosa.load(\n",
    "        path, sr=sample_rate, mono=True, res_type=\"kaiser_fast\"\n",
    "    )\n",
    "\n",
    "    duration = librosa.get_duration(y=data, sr=sample_rate)\n",
    "    # raise when sampling rate is unequal.\n",
    "    if actual_sampling_rate != sample_rate:\n",
    "        raise RuntimeError(\n",
    "            \"Sampling rate is not the desired one. Desired sampling rate: {sample_rate}, actual sampling rate: {actual_sampling_rate}\"\n",
    "        )\n",
    "\n",
    "    frame_length = int(sample_secs * sample_rate)\n",
    "    step_length = int(sample_secs - overlap) * sample_rate\n",
    "\n",
    "    chunks = tf_split_signal_into_chunks(\n",
    "        data, frame_length, step_length, pad_end=True\n",
    "    ).numpy()\n",
    "\n",
    "    print(\n",
    "        \"process audio data google: complete, read \",\n",
    "        str(len(chunks)),\n",
    "        \"chunks.\",\n",
    "        flush=True\n",
    "    )\n",
    "\n",
    "    return chunks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = preprocess(audiopath, 32000, 5.0, 0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_signature = interpreter.get_signature_runner()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_details = my_signature.get_input_details()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_details = my_signature.get_output_details()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_details\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_details"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = my_signature(inputs = tf.constant([1], shape = (1, 160000), dtype=tf.float32 ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output[\"output_0\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output[\"output_1\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(data[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = my_signature(inputs = data[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "interpreter = tflite.Interpreter(model_path=model_dir + \"/model.tflite\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "interpreter.allocate_tensors()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get input and output tensors.\n",
    "input_details = interpreter.get_input_details()\n",
    "output_details = interpreter.get_output_details()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_shape = input_details[0]['shape']\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_data = np.array(np.random.random_sample(input_shape), dtype=np.float32)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "interpreter.set_tensor(input_details[0]['index'], input_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "interpreter.invoke()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_data = interpreter.get_tensor(output_details[0]['index'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample = np.array([data[0]], dtype=\"float32\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "interpreter.set_tensor(input_details[0]['index'], sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "interpreter.invoke()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logits = interpreter.get_tensor(output_details[1]['index'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_data = tf.nn.softmax(logits).numpy()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# try to run the same thing with the original model and then compare the output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_model_from_file_pb(path: str, _):\n",
    "\n",
    "    if \".\" in Path(path).name or \".pb\" in Path(path).name:\n",
    "        # tensorflow assumes a model file to be named \"saved_model.pb\" and the path given to be a directory\n",
    "        path = Path(path).parent\n",
    "\n",
    "    if Path(path).exists() is False:\n",
    "        raise FileNotFoundError(\"The desired model file does not exist\")\n",
    "\n",
    "    try:\n",
    "        model = tf.saved_model.load(path)\n",
    "        return model\n",
    "    except Exception as e:\n",
    "        raise RuntimeError from e\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "other_model = load_model_from_file_pb(model_dir, None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logits, _ = other_model.infer_tf(\n",
    "    sample\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = tf.nn.softmax(logits).numpy()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.max(np.abs(results - output_data))"
   ]
  }
 ],
 "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",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}