diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..d8ee3b7 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,4 @@ +# Exclude files that are not needed in the serving image from the Docker build +# context. + +/python/data_processing/audio_utils.py \ No newline at end of file diff --git a/notebooks/README.md b/notebooks/README.md new file mode 100644 index 0000000..40bab9f --- /dev/null +++ b/notebooks/README.md @@ -0,0 +1,23 @@ +# HeAR Notebooks + +* Quick start with Hugging Face - + Example of encoding a health-related audio clip into an + embedding vector by running the model locally from Hugging Face. There are + different versions of this notebook for the following model formats: + + * [TensorFlow](quick_start_with_hugging_face.ipynb) + * [PyTorch](quick_start_with_hugging_face_pytorch.ipynb) + +* [Quick start with Vertex Model Garden](quick_start_with_model_garden.ipynb) - + Example of serving the model on + [Vertex AI](https://cloud.google.com/vertex-ai/docs/predictions/overview) + and using Vertex AI APIs to encode health-related audio to + embeddings in online or batch workflows. + +* [Train a data efficient classifier](train_data_efficient_classifier.ipynb) - + Example of using the generated embeddings to train a custom classifier with + less data and compute. + +* [Health event detector demo](hear_event_detector_demo.ipynb) - Example of + using HeAR with the health event detector models to first identify and then + generate embeddings for health-related sounds like coughing and breathing. \ No newline at end of file diff --git a/notebooks/hear_event_detector_demo.ipynb b/notebooks/hear_event_detector_demo.ipynb new file mode 100644 index 0000000..03d2140 --- /dev/null +++ b/notebooks/hear_event_detector_demo.ipynb @@ -0,0 +1,638 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "Mt64XBHqJiYo" + }, + "source": [ + "~~~markdown\n", + "Copyright 2025 Google LLC\n", + "\n", + "Licensed under the Apache License, Version 2.0 (the \"License\");\n", + "you may not use this file except in compliance with the License.\n", + "You may obtain a copy of the License at\n", + "\n", + " https://www.apache.org/licenses/LICENSE-2.0\n", + "\n", + "Unless required by applicable law or agreed to in writing, software\n", + "distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "See the License for the specific language governing permissions and\n", + "limitations under the License.\n", + "~~~\n", + "\n", + "# HeAR Event Detector Demo\n", + "\u003ctable\u003e\u003ctbody\u003e\u003ctr\u003e\n", + " \u003ctd style=\"text-align: center\"\u003e\n", + " \u003ca href=\"https://colab.research.google.com/github/google-health/hear/blob/master/notebooks/hear_event_detector_demo.ipynb\"\u003e\n", + " \u003cimg alt=\"Google Colab logo\" src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" width=\"32px\"\u003e\u003cbr\u003e Run in Google Colab\n", + " \u003c/a\u003e\n", + " \u003c/td\u003e \n", + " \u003ctd style=\"text-align: center\"\u003e\n", + " \u003ca href=\"https://github.com/google-health/hear/blob/master/notebooks/hear_event_detector_demo.ipynb\"\u003e\n", + " \u003cimg alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"\u003e\u003cbr\u003e View on GitHub\n", + " \u003c/a\u003e\n", + " \u003c/td\u003e\n", + " \u003ctd style=\"text-align: center\"\u003e\n", + " \u003ca href=\"https://huggingface.co/google/hear\"\u003e\n", + " \u003cimg alt=\"Hugging Face logo\" src=\"https://huggingface.co/front/assets/huggingface_logo-noborder.svg\" width=\"32px\"\u003e\u003cbr\u003e View on Hugging Face\n", + " \u003c/a\u003e\n", + " \u003c/td\u003e\n", + "\u003c/tr\u003e\u003c/tbody\u003e\u003c/table\u003e\n", + "\n", + "\n", + "This Colab notebook demonstrates using the HeAR (Health Acoustic Representations) model along with the included Health Event Detectors directly from Hugging Face, to identify audio clips with relevent health sounds such as coughing, breathing or sneezing, then create and utilize embeddings from this subset of health-related audio clips.\n", + "\n", + "This notebook is similar to `train_data_efficient_classifier.ipynb` and also uses the small [Wikimedia Commons](https://commons.wikimedia.org/wiki/Commons:Welcome) dataset of relevant health sounds. In this example the audio files are reduced to a smaller subset of clips using the event detector to identify clips containing interesting health sounds to embed with HeAR.\n", + "\n", + "\n", + "\n", + "#### This notebook demonstrates:\n", + "\n", + "1. Loading all supported Hugging Face Models (HeAR, Event Detector and Frontend).\n", + "\n", + "2. Detecting 2-second clips within the [Wikimedia Commons](https://commons.wikimedia.org/wiki/Commons:Welcome) dataset with high probability of containing one or more of the supported event detection labels, then generating HeAR embedddings for these clips.\n", + "\n", + "3. Finding the most similar audio files to a given query audio file based on the Cosine Similarity between the respective HeAR embeddings of the audio files.\n", + "\n", + "4. Optimizing the event detectors and frontend for low latency on-device usage using TFLite to support large scale feature event detection or feature generation" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "_uZFXCSuqr1V" + }, + "source": [ + "# Authenticate with HuggingFace, skip if you have a HF_TOKEN secret" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "l-5Tj0uqS3dI" + }, + "outputs": [], + "source": [ + "from huggingface_hub.utils import HfFolder\n", + "\n", + "if HfFolder.get_token() is None:\n", + " from huggingface_hub import notebook_login\n", + " notebook_login()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "mubK5MmPQcUS" + }, + "source": [ + "# Clone HuggingFace repository snapshot\n", + "\n", + "This will store the HeAR and event detector models in local cache so they can be loaded later.\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "WjfuNYeZQJdc" + }, + "outputs": [], + "source": [ + "import numpy as np\n", + "import tensorflow as tf\n", + "\n", + "print(\"TensorFlow version:\", tf.__version__)\n", + "print(\"Keras version:\", tf.keras.__version__)\n", + "\n", + "from huggingface_hub import snapshot_download\n", + "hugging_face_repo = \"google/hear\"\n", + "local_snapshot_path = snapshot_download(repo_id=hugging_face_repo)\n", + "print(f\"Saved {hugging_face_repo} to {local_snapshot_path}\\n\")\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "mO-Z5BOtj3D1" + }, + "source": [ + "# Download Audio Data" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "NBZnJwzzj3D1" + }, + "source": [ + " Wiki Commons\n", + "https://commons.wikimedia.org/wiki/Category:Coughing_audio\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "tJ55XsJsj3D2" + }, + "outputs": [], + "source": [ + "# @title Download Public Domain Cough Examples to Notebook\n", + "import os\n", + "import subprocess\n", + "from urllib.parse import urlparse\n", + "\n", + "# More examples: https://commons.wikimedia.org/wiki/Category:Coughing_audio\n", + "wiki_cough_file_urls = [\n", + " 'https://upload.wikimedia.org/wikipedia/commons/c/cc/Man_coughing.ogg',\n", + " 'https://upload.wikimedia.org/wikipedia/commons/6/6a/Cough_1.ogg',\n", + " 'https://upload.wikimedia.org/wikipedia/commons/d/d9/Cough_2.ogg',\n", + " 'https://upload.wikimedia.org/wikipedia/commons/b/be/Woman_coughing_three_times.wav',\n", + " 'https://upload.wikimedia.org/wikipedia/commons/d/d0/Sneezing.ogg',\n", + " 'https://upload.wikimedia.org/wikipedia/commons/e/ef/Laughter_and_clearing_voice.ogg',\n", + " 'https://upload.wikimedia.org/wikipedia/commons/c/c6/Laughter.ogg',\n", + " 'https://upload.wikimedia.org/wikipedia/commons/1/1c/Knocking_on_wood_or_door.ogg',\n", + "]\n", + "\n", + "# Download the files.\n", + "files_map = {} # file name to file path map\n", + "for url in wiki_cough_file_urls:\n", + " filename = os.path.basename(urlparse(url).path)\n", + " print(f'Downloading {filename}...')\n", + " res = subprocess.run(['wget', '-nv', '-O', filename, url], capture_output=True, text=True)\n", + " if res.returncode != 0:\n", + " print(f\" Download failed. Return code: {res.returncode}\\nError: {res.stderr}\")\n", + " files_map[filename] = url\n", + "print(f'\\nLocal Files:\\n{os.listdir():}\\n')" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "K2Eyu4VIsbUB" + }, + "source": [ + "# Load Models and Run Inference\n", + "\n", + "### HeAR Model\n", + "\n", + "The HeAR model uses a powerful [ViT](https://huggingface.co/docs/transformers/en/model_doc/vit) backbone and generates rich 512 length embeddings from a 2 second single channel 16kHz audio clip. See the [HuggingFace Model Card](https://huggingface.co/google/hear) for more details.\n", + "\n", + "### Event Detector Models\n", + "\n", + "The event detector models use the efficient [MobileNet-V3](https://huggingface.co/docs/timm/en/models/mobilenet-v3) backbone paired with our custom TensorFlow spectrogram frontend.\n", + "\n", + "As with HeAR, these models expect a 2 second single channel 16kHz audio clip and output 8 detection probability scores for the following labels:\n", + "\n", + "```\n", + "['Cough', 'Snore', 'Baby Cough', 'Breathe', 'Sneeze','Throat Clear', 'Laugh', 'Speech']\n", + "```\n", + "\n", + "The event detector has two size variants which can be used interchangeably depending on the use-case.\n", + "\n", + "* `event_detector_small/` is based on [MobileNetV3Small](https://www.tensorflow.org/api_docs/python/tf/keras/applications/MobileNetV3Small) with approximately **1M** parameters (3.60 MB)\n", + "\n", + "* `event_detector_large/` is based on [MobileNetV3Large](https://www.tensorflow.org/api_docs/python/tf/keras/applications/MobileNetV3Large) with approximately **3M** parameters (11.46 MB)\n", + "\n", + "#### Spectrogram Frontend\n", + "\n", + "Our event detectors are fused with a custom, on-device optimized spectrogram frontend which efficiently converts 2 seconds of 16kHz audio into [PCEN](https://research.google/pubs/trainable-frontend-for-robust-and-far-field-keyword-spotting) scaled [Mel-spectrogram](https://huggingface.co/learn/audio-course/en/chapter1/audio_data#mel-spectrogram) features with 200 time steps and 48 Mel-frequency bins.\n", + "\n", + "* `spectrogram_frontend/` is based on the PCEN implementation from [LEAF](https://research.google/blog/leaf-a-learnable-frontend-for-audio-classification/) and has only **5k** non-trainable parameters (18.56 KB) which are frozen and not configurable.\n", + "\n", + "We provide a standalone version of this frontend so that features can be pre-computed for new event detector training applications. See **Extract Batch Frontend Spectrogram Features** section for usage examples.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "cQVAxkBHZ_JO" + }, + "outputs": [], + "source": [ + "# @title Load HeAR, Event Detector and Frontend Models\n", + "from huggingface_hub import from_pretrained_keras\n", + "\n", + "# Constants for all included models, each input should be 32,000 samples.\n", + "SAMPLE_RATE = 16000\n", + "CLIP_DURATION = 2\n", + "\n", + "# Select event detector variant\n", + "EVENT_DETECTOR = \"event_detector_small\" # @param [\"event_detector_large\", \"event_detector_small\"]\n", + "# Included event detectors are trained to doutput detection probabilities in this order.\n", + "LABEL_LIST = ['Cough', 'Snore', 'Baby Cough', 'Breathe', 'Sneeze', 'Throat Clear', 'Laugh', 'Speech']\n", + "\n", + "# HeAR Embedding Model\n", + "print(f\"\\nLoading HeAR model\")\n", + "hear_model = from_pretrained_keras(local_snapshot_path)\n", + "hear_infer = hear_model.signatures[\"serving_default\"]\n", + "\n", + "# Event detector models and frontend are nested in the \"event_detector/\" folder.\n", + "# Detector Frontend Model for efficiently computing spectrogram feature\n", + "frontend_path = os.path.join(\"event_detector/\", \"spectrogram_frontend\")\n", + "print(f\"\\nLoading frontend model from: {frontend_path}\")\n", + "frontend_model = from_pretrained_keras(\n", + " os.path.join(local_snapshot_path, frontend_path)\n", + ")\n", + "\n", + "# Detector Model based on size variant selection, frontend included\n", + "event_detector_path = os.path.join(\"event_detector/\", EVENT_DETECTOR)\n", + "print(f\"\\nLoading detector model from: {event_detector_path}\")\n", + "event_detector = from_pretrained_keras(\n", + " os.path.join(local_snapshot_path, event_detector_path)\n", + ")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "33G4zKJHjGGc" + }, + "outputs": [], + "source": [ + "# @title Plot Helpers\n", + "import librosa\n", + "import matplotlib.pyplot as plt\n", + "import librosa.display\n", + "from IPython.display import Audio\n", + "import matplotlib.cm as cm\n", + "import warnings\n", + "\n", + "# Suppress the specific warning\n", + "warnings.filterwarnings(\"ignore\", category=UserWarning, module=\"soundfile\")\n", + "warnings.filterwarnings(\"ignore\", module=\"librosa\")\n", + "\n", + "def plot_waveform(sound, sr, title, figsize=(12, 4), color='blue', alpha=0.7):\n", + " \"\"\"Plots the waveform of the audio using librosa.display.\"\"\"\n", + " plt.figure(figsize=figsize)\n", + " librosa.display.waveshow(sound, sr=sr, color=color, alpha=alpha)\n", + " plt.title(f\"{title}\\nshape={sound.shape}, sr={sr}, dtype={sound.dtype}\")\n", + " plt.xlabel(\"Time (s)\")\n", + " plt.ylabel(\"Amplitude\")\n", + " plt.grid(True)\n", + " plt.tight_layout()\n", + " plt.show()\n", + "\n", + "\n", + "def plot_spectrogram(sound, sr, title, figsize=(12, 4), n_fft=2048, hop_length=256, n_mels=128, cmap='nipy_spectral'):\n", + " \"\"\"Plots the Mel spectrogram of the audio using librosa.\"\"\"\n", + " plt.figure(figsize=figsize)\n", + " mel_spectrogram = librosa.feature.melspectrogram(y=sound, sr=sr, n_fft=n_fft, hop_length=hop_length, n_mels=n_mels)\n", + " log_mel_spectrogram = librosa.power_to_db(mel_spectrogram, ref=np.max)\n", + " librosa.display.specshow(log_mel_spectrogram, sr=sr, hop_length=hop_length, x_axis='time', y_axis='mel', cmap=cmap)\n", + " plt.title(f\"{title} - Mel Spectrogram\")\n", + " plt.tight_layout()\n", + " plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "3dJKyHXznPgg" + }, + "outputs": [], + "source": [ + "# @title Event Detector Plot Helpers\n", + "\n", + "def plot_frontend_feature(\n", + " frontend_feature: np.ndarray,\n", + " title: str,\n", + " figsize: tuple[int, int] = (12, 4),\n", + " cmap: str = 'nipy_spectral',\n", + ") -\u003e None:\n", + " \"\"\"Plots the frontend spectrogram input feature.\n", + "\n", + " Args:\n", + " frontend_feature: The event detector frontend feature as a 2D NumPy array\n", + " with shape (number of time steps, number of frequency bins). The default\n", + " shape for the included event detectors is (200, 48), which represents a\n", + " 2-second audio clip with 48 frequency bins.\n", + " title: The title prefix of the plot.\n", + " figsize: Optional size of the figure.\n", + " cmap: Optional colormap to use.\n", + " \"\"\"\n", + " # Frontend features are typically rotated when fed into the model\n", + " # for spectrogram visualization it is more standard for x axis to be time\n", + " audio_spectrogram = np.rot90(frontend_feature)\n", + " plt.figure(figsize=figsize)\n", + " plt.imshow(audio_spectrogram, aspect='auto', cmap=cmap)\n", + " plt.title(f\"{title} - Frontend PCEN Mel Spectrogram\")\n", + " plt.tight_layout()\n", + " plt.show()\n", + "\n", + "def plot_detection_scores(\n", + " scores_batch: np.ndarray,\n", + " label_list: list[str],\n", + " title: str,\n", + " figsize: tuple[int, int] = (12, 4),\n", + " cmap: str = 'nipy_spectral',\n", + ") -\u003e None:\n", + " \"\"\"Plots per-label detection scores for batch sequentially with a consistent color scale.\n", + "\n", + " Args:\n", + " scores_batch: The event detection scores as a 2D NumPy array with shape\n", + " (number of clips, number of labels). Where number of labels should match\n", + " the length of the label_list, which is 8 for the included event detectors.\n", + " label_list: A list of labels representing the event detector classes.\n", + " title: The title prefix of the plot.\n", + " figsize: Optional size of the figure.\n", + " cmap: Optional colormap to use.\n", + " \"\"\"\n", + " plt.figure(figsize=figsize)\n", + " scores_img = np.transpose(scores_batch)\n", + " # Explicitly set the color limits for imshow\n", + " im = plt.imshow(scores_img, aspect='auto', cmap=cmap, vmin=0, vmax=1)\n", + " # Set up the 'y' label axis\n", + " plt.yticks(\n", + " np.arange(len(label_list)), [l.replace(' ', '\\n') for l in label_list]\n", + " )\n", + " # Add horizontal grid lines between labels\n", + " for i in range(1, scores_img.shape[0]):\n", + " plt.axhline(y=i - 0.5, color='gray', linestyle='--')\n", + " plt.grid(axis='y', which='major', color='white', alpha=0)\n", + " # Setup the 'x' time axis\n", + " n_clips = scores_img.shape[1]\n", + " plt.xticks(np.arange(n_clips), [f'Clip {i+1}' for i in range(n_clips)])\n", + " plt.xlabel(\"Time Step\")\n", + " # Add vertical grid lines between time steps\n", + " for j in range(1, n_clips):\n", + " plt.axvline(x=j - 0.5, color='gray', linestyle='--')\n", + " plt.title(f\"{title} - Sound Event Detections\")\n", + " # Add colorbar with a consistent scale from 0 to 1\n", + " plt.colorbar(im, ticks=[0, 0.2, 0.4, 0.6, 0.8, 1.0])\n", + " plt.tight_layout()\n", + " plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "KYbz904QXio2" + }, + "outputs": [], + "source": [ + "# @title Load Audio and Generate HeAR Embeddings\n", + "%%time\n", + "\n", + "# Audio display options\n", + "SHOW_WAVEFORM = False\n", + "SHOW_SPECTROGRAM = False\n", + "SHOW_PLAYER = True\n", + "SHOW_DETECTION_SCORES = True\n", + "COLORMAP = \"Blues\"\n", + "\n", + "# Keep clips with high detection scores for these respiratory labels\n", + "# then embed the clips using HeAR. Ignore clips with no detections\n", + "LABELS_TO_EMBED = ['Cough', 'Snore', 'Breathe', 'Sneeze']\n", + "# Assert that all labels to embed are actually present in the main label list\n", + "assert all(label in LABEL_LIST for label in LABELS_TO_EMBED)\n", + "\n", + "# Clips of length CLIP_DURATION seconds are extracted from the audio file\n", + "# using a sliding window. Adjecent clips are overlapped by CLIP_OVERLAP_PERCENT.\n", + "CLIP_OVERLAP_PERCENT = 10\n", + "\n", + "# Labels must have score above this threshold to be considered a detection\n", + "DETECTION_THRESHOLD = 0.9\n", + "\n", + "frame_length = int(CLIP_DURATION * SAMPLE_RATE)\n", + "frame_step = int(frame_length * (1 - CLIP_OVERLAP_PERCENT / 100))\n", + "hear_embeddings = {}\n", + "for file_key, file_url in files_map.items():\n", + " hear_embeddings[file_key] = {}\n", + " print(f\"\\nLoading file: {file_key} from {file_url}\")\n", + " audio, sample_rate = librosa.load(file_key, sr=SAMPLE_RATE, mono=True)\n", + "\n", + " # Display full audio file (optional).\n", + " if SHOW_WAVEFORM:\n", + " plot_waveform(audio, sample_rate, title=file_key, color='blue')\n", + " if SHOW_SPECTROGRAM:\n", + " plot_spectrogram(\n", + " audio, sample_rate, title=file_key, n_fft=2*1024, hop_length=64, n_mels=256, cmap=COLORMAP)\n", + " if SHOW_PLAYER:\n", + " display(Audio(data=audio, rate=sample_rate))\n", + "\n", + " # Segment an audio array into fixed length overlapping clips.\n", + " if len(audio) \u003c frame_length:\n", + " audio = np.pad(audio, (0, frame_length - len(audio)), mode='constant')\n", + " audio_clip_batch = tf.signal.frame(audio, frame_length, frame_step )\n", + " print(f\"Number of audio clips in batch: {len(audio_clip_batch)}.\")\n", + "\n", + " # Perform detector inference on the audio_clip_batch\n", + " # The model will generate the input feature, then infer the detection\n", + " print(f\"Running batched {EVENT_DETECTOR} model inference on audio clips.\")\n", + " detection_scores_batch = event_detector(audio_clip_batch)[\"scores\"].numpy()\n", + " print(\"Computed batch probability scores with shape:\",\n", + " f\"{detection_scores_batch.shape} from input audio clips batch with\",\n", + " f\"shape: {audio_clip_batch.shape}\"\n", + " )\n", + " hear_embeddings[file_key]['detections'] = detection_scores_batch\n", + "\n", + " if SHOW_DETECTION_SCORES:\n", + " plot_detection_scores(detection_scores_batch, LABEL_LIST, title=f'{file_key}: {EVENT_DETECTOR}', cmap=COLORMAP)\n", + "\n", + " # Filter clips for HeAR inference based on if in ANY detection scores for\n", + " # 'LABELS_TO_EMBED' are above the 'DETECTION_THRESHOLD'.\n", + " print(f\"Filtering clips based on detections for labels: {LABELS_TO_EMBED}\")\n", + " embed_hear_clips = []\n", + " for clip_i, scores in enumerate(detection_scores_batch):\n", + " for label_index, label in enumerate(LABEL_LIST):\n", + " if label in LABELS_TO_EMBED and scores[label_index] \u003e DETECTION_THRESHOLD:\n", + " embed_hear_clips.append(audio_clip_batch[clip_i])\n", + " break\n", + "\n", + " # Perform HeAR batch inference to extract the associated clip embedding.\n", + " # Only run inference on 'embed_hear_clips' which have have a high detection\n", + " # score for one of the 'LABELS_TO_EMBED'.\n", + " if len(embed_hear_clips):\n", + " print(f\"Computing HeAR embedding for batch of\",\n", + " f\"{len(embed_hear_clips)} selected clips.\")\n", + " hear_embedding_batch = hear_infer(x=np.asarray(embed_hear_clips))[ 'output_0'].numpy()\n", + " print(f\"Embedding batch shape: {hear_embedding_batch.shape},\",\n", + " f\"data type: {hear_embedding_batch.dtype}\")\n", + " else:\n", + " hear_embedding_batch = np.array([])\n", + " print(f\"None of the {len(audio_clip_batch)} clips in {file_key} have\",\n", + " f\"detections above the threshold: {DETECTION_THRESHOLD} for the\",\n", + " f\"labels: {LABELS_TO_EMBED}\")\n", + " hear_embeddings[file_key]['embeddings'] = hear_embedding_batch" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "PFLdq8WI3AJT" + }, + "outputs": [], + "source": [ + "# @title Use HeAR embeddings to find most similar file to query file\n", + "from scipy.spatial import distance\n", + "\n", + "# Set up query file and make sure if the query file_key exists in the dictionary and has embeddings.\n", + "query_file_key = 'Cough_1.ogg'\n", + "assert query_file_key in hear_embeddings and len(hear_embeddings[query_file_key]['embeddings'])\n", + "\n", + "# Get the average embedding for the query file and compare similarity to the average embedding for the other files.\n", + "query_embedding = np.mean(hear_embeddings[query_file_key]['embeddings'], axis=0)\n", + "similarities = {}\n", + "for file_key, model_outputs in hear_embeddings.items():\n", + " # Skip comparing file_key to itself or comparing to keys without HeAR embeddings.\n", + " if file_key == query_file_key or not len(model_outputs['embeddings']):\n", + " continue\n", + "\n", + " # Compute cosine similarity between the query file and the current file.\n", + " current_embedding = np.mean(model_outputs['embeddings'], axis=0)\n", + " similarities[file_key] = 1 - distance.cosine(query_embedding, current_embedding)\n", + "\n", + "# Find the top N most similar entries\n", + "N = 3\n", + "top_N_similar = dict(sorted(similarities.items(), key=lambda item: item[1], reverse=True)[:N])\n", + "print(f\"\\nTop {N} most similar entries to '{query_file_key}':\")\n", + "for key, similarity in top_N_similar.items():\n", + " print(f\" {key}: {similarity:.3f}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "q4LoPrUCs0F4" + }, + "source": [ + "# Extract Batch Frontend Spectrogram Features\n", + "\n", + "We provide a standalone, frozen frontend spectrogram model for efficient generation of [PCEN Mel-Spectrogram](https://research.google/pubs/trainable-frontend-for-robust-and-far-field-keyword-spotting/) features, which are used by the event detectors. This model is comprised of non-trainable TensorFlow operations, prioritizing portability, scalability, and optimization at the cost of configurability.\n", + "\n", + "The frontend can also be used as a standalone feature extractor for generating large amounts of training data for finetuning or retraining additional models as demonstrated below.\n", + "\n", + "The input must be 2-seconds and corresponding output feature will have shape `(200, 48)`. Typically, input clips will be segmented with some amount of overlap to avoid distorting sounds near the boundry." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "DLDc2OFAngcT" + }, + "outputs": [], + "source": [ + "# @title Plot Example Frontend Features\n", + "example_file_key = 'Sneezing.ogg'\n", + "audio, sample_rate = librosa.load(example_file_key, sr=SAMPLE_RATE, mono=True)\n", + "print(f\"Loaded {example_file_key} with duration {len(audio)/sample_rate:0.2f}s\")\n", + "\n", + "# Extract example clips with 50% overlap, then generate frontend feature batch.\n", + "frame_length = int(CLIP_DURATION * SAMPLE_RATE)\n", + "audio_clip_batch = tf.signal.frame(audio, frame_length, frame_length // 2 )\n", + "frontend_feature_batch = frontend_model(audio_clip_batch)\n", + "print(f\"Generated input feature batch from {example_file_key} with shape: {frontend_feature_batch.shape}\")\n", + "for clip_i, frontend_feature in enumerate(frontend_feature_batch):\n", + " plot_frontend_feature(frontend_feature.numpy(), title=f'Clip: {clip_i}', cmap=COLORMAP, figsize=(8, 3))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "_9e_m3Wf0ClK" + }, + "source": [ + "# Convert Event Detectors to TFLite\n", + "\n", + "The included event detectors are based on [MobileNet-V3](https://huggingface.co/docs/timm/en/models/mobilenet-v3) and are designed to be run on-device with low latency and power requirements while also allowing for finetuning and quick training.\n", + "\n", + "Converting the event detector models to [LiteRT](https://ai.google.dev/edge/litert) (formally TFLite) allows for more optimal performance on specific hardware such as mobile devices or compute constrained servers.\n", + "\n", + "In order to save compute in realtime or large scale applications, these hardware optimized event detectors can be used as a gate for HeAR, only embedding clips detected to contain a health sound of interest (as shown above).\n", + "\n", + "The code below demonstrates converting the loaded event detector to a TensorFlow LiteRT model, there are many more options available to quantize or furthur optimize the converted models, see the [documentation](https://ai.google.dev/edge/litert/models/convert_tf) for more information." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "I5UimMm15FMy" + }, + "outputs": [], + "source": [ + "# @title TFLite Conversion\n", + "%%time\n", + "\n", + "def convert_to_tflite(\n", + " model: tf.keras.Model,\n", + " quantize: bool = False,\n", + ") -\u003e bytes:\n", + " \"\"\"Converts a SavedModel model to a TensorFlow Lite (TFLite) model.\n", + "\n", + " Args:\n", + " model: The model to convert.\n", + " quantize: If True, apply dynamic range quantization to optimize the model.\n", + "\n", + " Returns:\n", + " The raw byte representation of the converted TFLite model.\n", + " \"\"\"\n", + " converter = tf.lite.TFLiteConverter.from_keras_model(model)\n", + " converter.target_spec.supported_ops = [\n", + " tf.lite.OpsSet.TFLITE_BUILTINS,\n", + " tf.lite.OpsSet.SELECT_TF_OPS, # needed for frontend ops\n", + " ]\n", + " # See documentation for quantization options beyond dyanmic range quantization.\n", + " # https://ai.google.dev/edge/litert/models/post_training_quantization\n", + " if quantize:\n", + " converter.optimizations = [tf.lite.Optimize.DEFAULT]\n", + " return converter.convert()\n", + "\n", + "# Convert event detector to TFLite and save result.\n", + "event_detector_lite = convert_to_tflite(event_detector, quantize=False)\n", + "tflite_output_path ='event_detector.tflite'\n", + "with open(tflite_output_path, 'wb') as f:\n", + " f.write(event_detector_lite)\n", + "print(f\"Saved TFLite model to: {tflite_output_path}\")\n", + "\n", + "# Initalize TFLite Model and print I/O specification.\n", + "tflite_interp = tf.lite.Interpreter(model_content=event_detector_lite)\n", + "print(f\"Input details:\\n {tflite_interp.get_input_details()[0]}\")\n", + "print(f\"Output details: \\n {tflite_interp.get_output_details()[0]}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "PHTxQttKYNpa" + }, + "source": [ + "# Next steps\n", + "\n", + "Explore the other [notebooks](https://github.com/google-health/hear/blob/master/notebooks) to learn what else you can do with the model." + ] + } + ], + "metadata": { + "colab": { + "last_runtime": { + "build_target": "//experimental/health_foundation_models/colab:colab_deps", + "kind": "private" + }, + "name": "hear_event_detector_demo.ipynb", + "private_outputs": true, + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/notebooks/quick_start_with_hugging_face.ipynb b/notebooks/quick_start_with_hugging_face.ipynb index 16728bc..716a340 100644 --- a/notebooks/quick_start_with_hugging_face.ipynb +++ b/notebooks/quick_start_with_hugging_face.ipynb @@ -32,12 +32,12 @@ " \u003c/td\u003e \n", " \u003ctd style=\"text-align: center\"\u003e\n", " \u003ca href=\"https://github.com/google-health/hear/blob/master/notebooks/quick_start_with_hugging_face.ipynb\"\u003e\n", - " \u003cimg alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"\u003e\u003cbr\u003e View on GitHub\n", + " \u003cimg alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"\u003e\u003cbr\u003e View on GitHub\n", " \u003c/a\u003e\n", " \u003c/td\u003e\n", " \u003ctd style=\"text-align: center\"\u003e\n", " \u003ca href=\"https://huggingface.co/google/hear\"\u003e\n", - " \u003cimg alt=\"HuggingFace logo\" src=\"https://huggingface.co/front/assets/huggingface_logo-noborder.svg\" width=\"32px\"\u003e\u003cbr\u003e View on HuggingFace\n", + " \u003cimg alt=\"Hugging Face logo\" src=\"https://huggingface.co/front/assets/huggingface_logo-noborder.svg\" width=\"32px\"\u003e\u003cbr\u003e View on Hugging Face\n", " \u003c/a\u003e\n", " \u003c/td\u003e\n", "\u003c/tr\u003e\u003c/tbody\u003e\u003c/table\u003e\n", diff --git a/notebooks/quick_start_with_hugging_face_pytorch.ipynb b/notebooks/quick_start_with_hugging_face_pytorch.ipynb new file mode 100644 index 0000000..838e06c --- /dev/null +++ b/notebooks/quick_start_with_hugging_face_pytorch.ipynb @@ -0,0 +1,264 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "1fgVWTMK9SNz" + }, + "source": [ + "~~~\n", + "Copyright 2025 Google LLC\n", + "\n", + "Licensed under the Apache License, Version 2.0 (the \"License\");\n", + "you may not use this file except in compliance with the License.\n", + "You may obtain a copy of the License at\n", + "\n", + " https://www.apache.org/licenses/LICENSE-2.0\n", + "\n", + "Unless required by applicable law or agreed to in writing, software\n", + "distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "See the License for the specific language governing permissions and\n", + "limitations under the License.\n", + "~~~\n", + "\n", + "\n", + "# Quick start with Hugging Face (PyTorch model)\n", + "\n", + "\u003ctable\u003e\u003ctbody\u003e\u003ctr\u003e\n", + " \u003ctd style=\"text-align: center\"\u003e\n", + " \u003ca href=\"https://colab.research.google.com/github/google-health/hear/blob/master/notebooks/quick_start_with_hugging_face_pytorch.ipynb\"\u003e\n", + " \u003cimg alt=\"Google Colab logo\" src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" width=\"32px\"\u003e\u003cbr\u003e Run in Google Colab\n", + " \u003c/a\u003e\n", + " \u003c/td\u003e \n", + " \u003ctd style=\"text-align: center\"\u003e\n", + " \u003ca href=\"https://github.com/google-health/hear/blob/master/notebooks/quick_start_with_hugging_face_pytorch.ipynb\"\u003e\n", + " \u003cimg alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"\u003e\u003cbr\u003e View on GitHub\n", + " \u003c/a\u003e\n", + " \u003c/td\u003e\n", + " \u003ctd style=\"text-align: center\"\u003e\n", + " \u003ca href=\"https://huggingface.co/google/hear-pytorch\"\u003e\n", + " \u003cimg alt=\"Hugging Face logo\" src=\"https://huggingface.co/front/assets/huggingface_logo-noborder.svg\" width=\"32px\"\u003e\u003cbr\u003e View on Hugging Face\n", + " \u003c/a\u003e\n", + " \u003c/td\u003e\n", + "\u003c/tr\u003e\u003c/tbody\u003e\u003c/table\u003e\n", + "\n", + "This Colab notebook provides a basic usage example of the HeAR encoder that generates a machine learning representation (known as \"embeddings\") from health-related sounds (2-second audio clips sampled at 16kHz). These embeddings can be used to develop custom machine learning models for health acoustic use-cases with less data and compute compared to traditional model development methods.\n", + "\n", + " Learn more about embeddings and their benefits at [this page](https://developers.google.com/health-ai-developer-foundations/hear)." + ] + }, + { + "metadata": { + "id": "EvQJsuAZ50-Z" + }, + "cell_type": "markdown", + "source": [ + "## Install dependencies" + ] + }, + { + "metadata": { + "id": "qdrXvhfo6dMV" + }, + "cell_type": "code", + "source": [ + "! git clone https://github.com/Google-Health/hear.git\n", + "! pip install --upgrade --quiet transformers==4.50.3" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "metadata": { + "id": "_uZFXCSuqr1V" + }, + "source": [ + "## Authenticate with HuggingFace, skip if you have a HF_TOKEN secret" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "9XpCla68-Iol" + }, + "outputs": [], + "source": [ + "from huggingface_hub.utils import HfFolder\n", + "\n", + "if HfFolder.get_token() is None:\n", + " from huggingface_hub import notebook_login\n", + " notebook_login()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "SWckxRXiqc3L" + }, + "source": [ + "## Load and play cough audio recording" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "z2vEauhtvTan" + }, + "outputs": [], + "source": [ + "SAMPLE_RATE = 16000 # Samples per second (Hz)\n", + "CLIP_DURATION = 2 # Duration of the audio clip in seconds\n", + "CLIP_LENGTH = SAMPLE_RATE * CLIP_DURATION # Total number of samples" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "h1Lw92M1b7Dd" + }, + "outputs": [], + "source": [ + "!wget -nc https://upload.wikimedia.org/wikipedia/commons/b/be/Woman_coughing_three_times.wav" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "5W7BMKk9cerB" + }, + "outputs": [], + "source": [ + "from scipy.io import wavfile\n", + "\n", + "# Load file\n", + "with open('Woman_coughing_three_times.wav', 'rb') as f:\n", + " original_sampling_rate, audio_array = wavfile.read(f)\n", + "\n", + "print(f\"Sample Rate: {original_sampling_rate} Hz\")\n", + "print(f\"Data Shape: {audio_array.shape}\")\n", + "print(f\"Data Type: {audio_array.dtype}\")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "1l-eFOh0mq1r" + }, + "outputs": [], + "source": [ + "from IPython.display import Audio, display\n", + "import importlib\n", + "audio_utils = importlib.import_module(\n", + " \"hear.python.data_processing.audio_utils\"\n", + ")\n", + "resample_audio_and_convert_to_mono = audio_utils.resample_audio_and_convert_to_mono\n", + "\n", + "\n", + "audio_array = resample_audio_and_convert_to_mono(\n", + " audio_array=audio_array, \n", + " sampling_rate=original_sampling_rate,\n", + " new_sampling_rate=SAMPLE_RATE,\n", + ")\n", + "display(Audio(audio_array, rate=SAMPLE_RATE))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "fJIdGmpFqhi8" + }, + "source": [ + "## Compute embeddings" + ] + }, + { + "metadata": { + "id": "MIRuwbfz2a-l" + }, + "cell_type": "code", + "source": [ + "from transformers import AutoModel\n", + "\n", + "\n", + "# Load the model directly from Hugging Face Hub\n", + "loaded_model = AutoModel.from_pretrained(\"google/hear-pytorch\")" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "Fhk4GKkwsAMj" + }, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import torch\n", + "\n", + "preprocess_audio = audio_utils.preprocess_audio\n", + "\n", + "# This index corresponds to a cough and was determined by hand. In practice, you\n", + "# would need a detector.\n", + "START = 0\n", + "\n", + "# Add batch dimension\n", + "input_tensor = np.expand_dims(audio_array[START: START + CLIP_LENGTH], axis=0)\n", + "\n", + "# Call inference\n", + "infer = lambda audio_array: loaded_model.forward(\n", + " preprocess_audio(audio_array), return_dict=True, output_hidden_states=True)\n", + "output = infer(torch.Tensor(input_tensor))\n", + "\n", + "# Extract the embedding vector\n", + "embedding_vector = np.asarray(output.pooler_output.detach()).flatten()\n", + "print(\"Size of embedding vector:\", len(embedding_vector))\n", + "\n", + "# Plot the embedding vector\n", + "plt.figure(figsize=(12, 4))\n", + "plt.plot(embedding_vector)\n", + "plt.title('Embedding Vector')\n", + "plt.xlabel('Index')\n", + "plt.ylabel('Value')\n", + "plt.grid(True)\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "PHTxQttKYNpa" + }, + "source": [ + "# Next steps\n", + "\n", + "Explore the other [notebooks](https://github.com/google-health/hear/blob/master/notebooks) to learn what else you can do with the model." + ] + } + ], + "metadata": { + "colab": { + "name": "quick_start_with_hugging_face_pytorch.ipynb", + "private_outputs": true, + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/notebooks/quick_start_with_model_garden.ipynb b/notebooks/quick_start_with_model_garden.ipynb index 24f541e..fbd0fea 100644 --- a/notebooks/quick_start_with_model_garden.ipynb +++ b/notebooks/quick_start_with_model_garden.ipynb @@ -40,7 +40,7 @@ " \n", " \n", " \n", - " \"GitHub
View on GitHub\n", + " \"GitHub
View on GitHub\n", " \n", " \n", "" diff --git a/notebooks/train_data_efficient_classifier.ipynb b/notebooks/train_data_efficient_classifier.ipynb index f47200e..4542bd5 100644 --- a/notebooks/train_data_efficient_classifier.ipynb +++ b/notebooks/train_data_efficient_classifier.ipynb @@ -32,12 +32,12 @@ " \u003c/td\u003e \n", " \u003ctd style=\"text-align: center\"\u003e\n", " \u003ca href=\"https://github.com/google-health/hear/blob/master/notebooks/train_data_efficient_classifier.ipynb\"\u003e\n", - " \u003cimg alt=\"GitHub logo\" src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" width=\"32px\"\u003e\u003cbr\u003e View on GitHub\n", + " \u003cimg alt=\"GitHub logo\" src=\"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\" width=\"32px\"\u003e\u003cbr\u003e View on GitHub\n", " \u003c/a\u003e\n", " \u003c/td\u003e\n", " \u003ctd style=\"text-align: center\"\u003e\n", " \u003ca href=\"https://huggingface.co/google/hear\"\u003e\n", - " \u003cimg alt=\"HuggingFace logo\" src=\"https://huggingface.co/front/assets/huggingface_logo-noborder.svg\" width=\"32px\"\u003e\u003cbr\u003e View on HuggingFace\n", + " \u003cimg alt=\"Hugging Face logo\" src=\"https://huggingface.co/front/assets/huggingface_logo-noborder.svg\" width=\"32px\"\u003e\u003cbr\u003e View on Hugging Face\n", " \u003c/a\u003e\n", " \u003c/td\u003e\n", "\u003c/tr\u003e\u003c/tbody\u003e\u003c/table\u003e\n", diff --git a/python/data_processing/audio_utils.py b/python/data_processing/audio_utils.py new file mode 100644 index 0000000..7b15211 --- /dev/null +++ b/python/data_processing/audio_utils.py @@ -0,0 +1,516 @@ +# +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Utility functions for audio processing.""" + +import math +from typing import Callable + +import numpy as np +from scipy import signal +import torch +import torch.nn.functional as F + + +def _enclosing_power_of_two(value: int) -> int: + """Calculates the smallest power of 2 greater than or equal to `value`.""" + return int(2 ** math.ceil(math.log2(value))) if value > 0 else 1 + + +def _compute_stft( + signals: torch.Tensor, + frame_length: int, + frame_step: int, + fft_length: int | None = None, + window_fn: Callable[[int], torch.Tensor] | None = torch.hann_window, + pad_end: bool = True, +) -> torch.Tensor: + """Computes the Short-time Fourier Transform of `signals`. + + Args: + signals: A `[..., samples]` `torch.Tensor` of real-valued signals. + frame_length: The window length in samples. + frame_step: The number of samples to step. + fft_length: The size of the FFT to apply. If not provided, uses the smallest + power of 2 enclosing `frame_length`. + window_fn: A callable that takes a window length and returns a + `[window_length]` `torch.Tensor` of samples. If set to `None`, no + windowing is used. + pad_end: Whether to pad the end of `signals` with zeros when the provided + frame length and step produces a frame that lies partially past its end. + + Returns: + A `[..., frames, fft_unique_bins]` `torch.Tensor` of `complex64` + STFT values where `fft_unique_bins` is `fft_length // 2 + 1` (the unique + components of the FFT). + + Raises: + ValueError: If `signals` is not at least rank 1, `frame_length` is + not scalar, or `frame_step` is not scalar. + """ + if signals.ndim < 1: + raise ValueError( + f'Input signals must have rank at least 1, got rank {signals.ndim}' + ) + if not isinstance(frame_length, int): + raise ValueError( + f'frame_length must be an integer scalar, got type {type(frame_length)}' + ) + if not isinstance(frame_step, int): + raise ValueError( + f'frame_step must be an integer scalar, got type {type(frame_step)}' + ) + + if fft_length is None: + fft_length = _enclosing_power_of_two(frame_length) + elif not isinstance(fft_length, int): + raise ValueError( + 'fft_length must be an integer scalar or None, got type' + f' {type(fft_length)}' + ) + + if pad_end: + n_frames = ( + math.ceil(signals.shape[-1] / frame_step) + if signals.shape[-1] > 0 + else 0 + ) + padded_length = ( + max(0, (n_frames - 1) * frame_step + frame_length) + if n_frames > 0 + else frame_length + ) + padding_needed = max(0, padded_length - signals.shape[-1]) + if padding_needed > 0: + signals = F.pad(signals, (0, padding_needed)) + + framed_signals = signals.unfold(-1, frame_length, frame_step) + + if framed_signals.shape[-2] == 0: + return torch.empty( + *signals.shape[:-1], + 0, + fft_length // 2 + 1, + dtype=torch.complex64, + device=signals.device, + ) + + # Optionally window the framed signals. + if window_fn is not None: + window = ( + window_fn(frame_length) + .to(framed_signals.device) + .to(framed_signals.dtype) + ) # Ensure window is on the same device and dtype + framed_signals = framed_signals * window + + # torch.fft.rfft produces the (fft_length/2 + 1) unique components of the + # FFT of the real windowed signals in framed_signals. + return torch.fft.rfft(framed_signals, n=fft_length, dim=-1) + + +def _ema( + inputs: torch.Tensor, + num_channels: int, + smooth_coef: float, + initial_state: torch.Tensor | None = None, +) -> torch.Tensor: + """Exponential Moving Average (EMA). + + Args: + inputs: A 3D tensor of shape (batch_size, timesteps, input_dim). This is the + input sequence to the EMA RNN. + num_channels: The number of channels/units in the EMA. + smooth_coef: The smoothing coefficient (alpha) for the EMA. + initial_state: (Optional) A 2D tensor of shape (batch_size, num_channels) + representing the initial state for the EMA. If None, it defaults to zeros. + + Returns: + A 3D tensor of shape (batch_size, timesteps, num_channels) representing + the EMA output sequence. + """ + batch_size, timesteps, _ = inputs.shape + + if initial_state is None: + ema_state = torch.zeros( + (batch_size, num_channels), dtype=torch.float32, device=inputs.device + ) + else: + ema_state = initial_state + + identity_kernel_gain = smooth_coef + identity_recurrent_gain = 1.0 - smooth_coef + + identity_kernel = ( + torch.eye(num_channels, dtype=torch.float32, device=inputs.device) + * identity_kernel_gain + ) + identity_recurrent_kernel = ( + torch.eye(num_channels, dtype=torch.float32, device=inputs.device) + * identity_recurrent_gain + ) + + output_sequence = [] + + start = initial_state is not None + if start: + output_sequence.append(ema_state) + for t in range(start, timesteps): + current_input = inputs[:, t, :] # Shape (batch_size, num_channels) + + # EMA update formula: + output = torch.matmul(current_input, identity_kernel) + torch.matmul( + ema_state, identity_recurrent_kernel + ) + + ema_state = output + output_sequence.append(output) + + # Shape (batch_size, timesteps, num_channels) + output_sequence = torch.stack(output_sequence, dim=1) + return output_sequence + + +def _pcen_function( + inputs: torch.Tensor, + num_channels: int = 128, + alpha: float = 0.8, + smooth_coef: float = 0.04, + delta: float = 2.0, + root: float = 2.0, + floor: float = 1e-8, +) -> torch.Tensor: + """Per-Channel Energy Normalization as a function. + + This applies a fixed normalization by an exponential moving + average smoother, and a compression. + See https://arxiv.org/abs/1607.05666 for more details. + + Args: + inputs: A `[..., num_channels]` `torch.Tensor` of input signals. + num_channels: Number of channels + alpha: Exponent of EMA smoother + smooth_coef: Smoothing coefficient of EMA + delta: Bias added before compression + root: One over exponent applied for compression (r in the paper) + floor: Offset added to EMA smoother + + Returns: + A `[..., num_channels]` `torch.Tensor` of PCEN values. + """ + alpha_param = torch.ones(num_channels) * alpha + delta_param = torch.ones(num_channels) * delta + root_param = torch.ones(num_channels) * root + + delta_param = delta_param.to(inputs.device).to(inputs.dtype) + alpha_param = ( + torch.minimum(alpha_param, torch.ones_like(alpha_param)) + .to(inputs.device) + .to(inputs.dtype) + ) + root_param = ( + torch.maximum(root_param, torch.ones_like(root_param)) + .to(inputs.device) + .to(inputs.dtype) + ) + ema_smoother = _ema( + inputs, + num_channels=num_channels, + smooth_coef=smooth_coef, + # Handle cases where input is 2D or 3D + initial_state=inputs[:, 0] if inputs.ndim > 1 else None, + ).to(inputs.device) + + one_over_root = 1.0 / root_param + output = ( + inputs / (floor + ema_smoother) ** alpha_param + delta_param + ) ** one_over_root - delta_param**one_over_root + return output + + +def _hertz_to_mel(frequencies_hertz: torch.Tensor) -> torch.Tensor: + """Scale filter frequencies to mel scale. + + https://en.wikipedia.org/wiki/Mel_scale + + Args: + frequencies_hertz: A float Tensor of frequencies in Hertz. + + Returns: + A Tensor of the same shape but in mel-scale. + """ + return 2595.0 * torch.log10(1.0 + frequencies_hertz / 700.0) + + +def _linear_to_mel_weight_matrix( + device: torch.device, + num_mel_bins: int = 128, + num_spectrogram_bins: int = 201, + sample_rate: float = 16000, + lower_edge_hertz: float = 0.0, + upper_edge_hertz: float = 8000.0, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Returns a matrix to warp linear scale spectrograms to the mel scale. + + This function is a PyTorch version of the TensorFlow function with the same + name. See the TensorFlow documentation for detailed explanation and usage + https://www.tensorflow.org/api_docs/python/tf/signal/linear_to_mel_weight_matrix + + Args: + device: The device to use for the result matrix. + num_mel_bins: How many bands in the resulting mel spectrum. + num_spectrogram_bins: How many bins there are in the source spectrogram + data. + sample_rate: Samples per second of the input signal used to create the + spectrogram. + lower_edge_hertz: Lower bound on the frequencies to be included in the mel + spectrum. + upper_edge_hertz: The desired top edge of the highest frequency band. + dtype: The torch.dtype of the result matrix. Must be a floating point type. + + Returns: + A torch.Tensor of shape `[num_spectrogram_bins, num_mel_bins]`. + + Raises: + ValueError: If `num_mel_bins`/`num_spectrogram_bins`/`sample_rate` are not + positive, `lower_edge_hertz` is negative, frequency edges are incorrectly + ordered, `upper_edge_hertz` is larger than the Nyquist frequency. + """ + + if num_mel_bins <= 0: + raise ValueError(f'num_mel_bins must be positive. Got: {num_mel_bins}.') + if num_spectrogram_bins <= 0: + raise ValueError( + f'num_spectrogram_bins must be positive. Got: {num_spectrogram_bins}.' + ) + if sample_rate <= 0: + raise ValueError(f'sample_rate must be positive. Got: {sample_rate}.') + if lower_edge_hertz < 0.0: + raise ValueError( + f'lower_edge_hertz must be non-negative. Got: {lower_edge_hertz}.' + ) + if lower_edge_hertz >= upper_edge_hertz: + raise ValueError( + 'lower_edge_hertz must be smaller than upper_edge_hertz. Got: ' + f'lower_edge_hertz={lower_edge_hertz}, ' + f'upper_edge_hertz={upper_edge_hertz}.' + ) + if upper_edge_hertz > sample_rate / 2.0: + raise ValueError( + 'upper_edge_hertz must not be larger than the Nyquist frequency' + f'({sample_rate / 2.0}). Got: upper_edge_hertz={upper_edge_hertz}.' + ) + + sample_rate_tensor = torch.tensor(sample_rate, dtype=dtype) + lower_edge_hertz_tensor = torch.tensor( + lower_edge_hertz, dtype=dtype, device=device + ) + upper_edge_hertz_tensor = torch.tensor( + upper_edge_hertz, dtype=dtype, device=device + ) + zero = torch.tensor(0.0, dtype=dtype, device=device) + + # HTK excludes the spectrogram DC bin. + bands_to_zero = 1 + nyquist_hertz = sample_rate_tensor / 2.0 + linear_frequencies = torch.linspace( + zero, nyquist_hertz, num_spectrogram_bins, dtype=dtype, device=device + )[bands_to_zero:] + spectrogram_bins_mel = _hertz_to_mel(linear_frequencies).unsqueeze(1) + + # Compute num_mel_bins triples of (lower_edge, center, upper_edge). + # The center of each band is the lower and upper edge of the adjacent bands. + # Accordingly, we divide [lower_edge_hertz, upper_edge_hertz] into + # num_mel_bins + 2 pieces. + band_edges_mel = torch.linspace( + _hertz_to_mel(lower_edge_hertz_tensor), + _hertz_to_mel(upper_edge_hertz_tensor), + num_mel_bins + 2, + dtype=dtype, + device=device, + ) + # Create frames of size 3 with stride 1 + band_edges_mel = band_edges_mel.unfold(0, 3, 1) + + # Split the triples up and reshape them into [1, num_mel_bins] tensors. + lower_edge_mel = band_edges_mel[:, 0].unsqueeze(0) + center_mel = band_edges_mel[:, 1].unsqueeze(0) + upper_edge_mel = band_edges_mel[:, 2].unsqueeze(0) + + # Calculate lower and upper slopes for every spectrogram bin. + # Line segments are linear in the mel domain, not Hertz. + lower_slopes = (spectrogram_bins_mel - lower_edge_mel) / ( + center_mel - lower_edge_mel + ) + upper_slopes = (upper_edge_mel - spectrogram_bins_mel) / ( + upper_edge_mel - center_mel + ) + + # Intersect the line segments with each other and zero. + mel_weights_matrix = torch.maximum( + zero, torch.minimum(lower_slopes, upper_slopes) + ) + + # Re-add the zeroed lower bins we sliced out above. + return F.pad( + mel_weights_matrix, (0, 0, bands_to_zero, 0), mode='constant', value=0.0 + ) + + +def _mel_pcen( + x: torch.Tensor, +) -> torch.Tensor: + """Melspec followed by pcen.""" + x = x.float() + # Scale to -1, 1 range + x -= torch.min(x) + x = x / (torch.max(x) + 1e-8) + x = (x * 2) - 1 + + frame_length = 16 * 25 + frame_step = 160 + + stft = _compute_stft( + x, + frame_length=frame_length, + fft_length=frame_length, + frame_step=frame_step, + window_fn=torch.hann_window, + pad_end=True, + ) + spectrograms = torch.square(torch.abs(stft)) + + mel_transform = _linear_to_mel_weight_matrix(x.device) + mel_spectrograms = torch.matmul(spectrograms, mel_transform) + return _pcen_function(mel_spectrograms) + + +def _torch_resize_bilinear_tf_compat(images, size): + """PyTorch implementation of tf.image.resize. + + Internally matches the numerical output of: + ``` + tf.image.resize( + # TF input needs HWC/BHWC format + tf_permuted_input, + size, + method=tf.image.ResizeMethod.BILINEAR, + preserve_aspect_ratio=False, + antialias=False + ) + ``` + + Args: + images: Input tensor, shape [C, H, W] or [B, C, H, W]. Should be float32 + or convertible to float32 for numerical matching. + size: Target size as [new_height, new_width]. + + Returns: + Resized tensor, shape [C, new_H, new_W] or [B, C, new_H, new_W], dtype + torch.float32. + """ + original_dims = images.dim() + new_height, new_width = size + + if original_dims not in [3, 4]: + raise ValueError('Input tensor must be 3D [C, H, W] or 4D [B, C, H, W]') + + if len(size) != 2: + raise ValueError( + 'size must be a tuple or list of 2 integers (new_height, new_width)' + ) + + images = images.to(torch.float32) + + was_3d = False + if original_dims == 3: + images = images.unsqueeze(0) # Shape: [1, C, H, W] + was_3d = True + + resized_images_bchw = F.interpolate( + images, # Shape: [B, C, H, W] + size=(new_height, new_width), + mode='bilinear', + align_corners=False, + antialias=False, + ) + # Output shape: [B, C, new_H, new_W] + + # Remove batch dimension if original input was 3D (1CNHWC -> CNHWC) + if was_3d: + # Shape: [C, new_H, new_W] + resized_images_bchw = resized_images_bchw.squeeze(0) + + return resized_images_bchw + + +def preprocess_audio(audio: torch.Tensor) -> torch.Tensor: + """Preprocesses audio. + + Args: + audio: A `[..., samples]` `torch.Tensor` of real-valued signals. Represents + batched audio clips of duration 2s sampled at 16kHz. + + Returns: + A `[..., 1, 192, 128]` `torch.Tensor` of mel-pcen values. + """ + if audio.ndim != 2: + raise ValueError( + f'Input audio must have rank 2, got rank {audio.ndim}' + ) + if audio.shape[1] < 32000: + n = 32000 - audio.shape[1] + audio = torch.nn.functional.pad(audio, pad=(0, n), mode='constant', value=0) + elif audio.shape[1] > 32000: + raise ValueError( + f'Input audio must have 32000 samples, got {audio.shape[1]}' + ) + spectrogram = _mel_pcen(audio) + # Add a channel dimension. Torch images have format [B, C, H, W]. + spectrogram = torch.unsqueeze(spectrogram, dim=1) + return _torch_resize_bilinear_tf_compat(spectrogram, size=(192, 128)) + + +def resample_audio_and_convert_to_mono( + audio_array: np.ndarray, + sampling_rate: float, + new_sampling_rate: float, +) -> np.ndarray: + """Resamples an audio array to 16kHz and converts it to mono. + + Args: + audio_array: A numpy array representing the audio data. + sampling_rate: The original sampling rate of the audio. + new_sampling_rate: Target sampling rate. + + Returns: + resampled_audio_mono: A numpy array representing the resampled mono audio at + 16kHz. + """ + # Convert to mono if it's multi-channel + if audio_array.ndim > 1: + audio_mono = np.mean(audio_array, axis=1) + else: + audio_mono = audio_array + + # Resample + original_sample_count = audio_mono.shape[0] + new_sample_count = int( + round(original_sample_count * (new_sampling_rate / sampling_rate)) + ) + resampled_audio_mono = signal.resample(audio_mono, new_sample_count) + + return resampled_audio_mono diff --git a/python/serving/requirements.in b/python/serving/requirements.in index 8c6fc5a..2410d2f 100644 --- a/python/serving/requirements.in +++ b/python/serving/requirements.in @@ -5,3 +5,4 @@ tensorflow~=2.18.0 -r ../data_processing/requirements.in -r serving_framework/requirements.in +-r serving_framework/tensorflow/requirements.in diff --git a/python/serving/server_gunicorn.py b/python/serving/server_gunicorn.py index de070df..a554f3c 100644 --- a/python/serving/server_gunicorn.py +++ b/python/serving/server_gunicorn.py @@ -28,6 +28,7 @@ from serving.serving_framework import inline_prediction_executor from serving.serving_framework import server_gunicorn +from serving.serving_framework.tensorflow import server_model_runner from serving import predictor @@ -52,7 +53,8 @@ def main(argv: Sequence[str]) -> None: logging.info('Launching gunicorn application.') server_gunicorn.PredictionApplication( inline_prediction_executor.InlinePredictionExecutor( - predictor_instance.predict + predictor_instance.predict, + server_model_runner.ServerModelRunner ), health_check=health_checker, options=options, diff --git a/python/serving/serving_framework/inline_prediction_executor.py b/python/serving/serving_framework/inline_prediction_executor.py index a3e25d2..0be41bd 100644 --- a/python/serving/serving_framework/inline_prediction_executor.py +++ b/python/serving/serving_framework/inline_prediction_executor.py @@ -30,7 +30,6 @@ from serving.serving_framework import model_runner from serving.serving_framework import server_gunicorn -from serving.serving_framework import server_model_runner class InlinePredictionExecutor(server_gunicorn.PredictionExecutor): @@ -49,9 +48,11 @@ def __init__( predictor: Callable[ [dict[str, Any], model_runner.ModelRunner], dict[str, Any] ], + model_runner_source: Callable[[], model_runner.ModelRunner] ): self._predictor = predictor self._model_runner = None + self._model_runner_source = model_runner_source @override def start(self) -> None: @@ -61,7 +62,7 @@ def start(self) -> None: which needs to be done post-fork. """ # Safer to instantiate the RPC stub post-fork. - self._model_runner = server_model_runner.ServerModelRunner() + self._model_runner = self._model_runner_source() def predict(self, input_json: dict[str, Any]) -> dict[str, Any]: """Executes the given request payload.""" diff --git a/python/serving/serving_framework/inline_prediction_executor_test.py b/python/serving/serving_framework/inline_prediction_executor_test.py index af21676..96ede9a 100644 --- a/python/serving/serving_framework/inline_prediction_executor_test.py +++ b/python/serving/serving_framework/inline_prediction_executor_test.py @@ -12,39 +12,64 @@ # See the License for the specific language governing permissions and # limitations under the License. +from collections.abc import Mapping, Set from unittest import mock +import numpy as np + from absl.testing import absltest from serving.serving_framework import inline_prediction_executor -from serving.serving_framework import server_model_runner +from serving.serving_framework import model_runner + + +class DummyModelRunner(model_runner.ModelRunner): + """Dummy model runner for testing.""" + + def run_model_multiple_output( + self, + model_input: Mapping[str, np.ndarray] | np.ndarray, + *, + model_name: str = "default", + model_version: int | None = None, + model_output_keys: Set[str], + ) -> Mapping[str, np.ndarray]: + del model_name, model_version, model_output_keys + return {"output_0": np.ones((1, 2), dtype=np.float32)} class InlinePredictionExecutorTest(absltest.TestCase): def test_predict_requires_start(self): predictor = mock.MagicMock() - executor = inline_prediction_executor.InlinePredictionExecutor(predictor) + executor = inline_prediction_executor.InlinePredictionExecutor( + predictor, DummyModelRunner + ) with self.assertRaises(RuntimeError): executor.predict({"placeholder": "input"}) def test_execute_catches_predictor_exception(self): predictor = mock.MagicMock(side_effect=Exception("test error")) - executor = inline_prediction_executor.InlinePredictionExecutor(predictor) + executor = inline_prediction_executor.InlinePredictionExecutor( + predictor, DummyModelRunner + ) executor.start() with self.assertRaises(RuntimeError): executor.execute({"placeholder": "input"}) def test_execute_calls_predictor(self): predictor = mock.MagicMock(return_value={"placeholder": "output"}) - executor = inline_prediction_executor.InlinePredictionExecutor(predictor) mock_model_runner = mock.create_autospec( - server_model_runner.ServerModelRunner, instance=True + DummyModelRunner, instance=True + ) + mock_model_runner_class = mock.create_autospec( + DummyModelRunner, autospec=True + ) + mock_model_runner_class.return_value = mock_model_runner + executor = inline_prediction_executor.InlinePredictionExecutor( + predictor, mock_model_runner_class ) - with mock.patch.object( - server_model_runner, "ServerModelRunner", autospec=True - ) as mock_model_runner_class: - mock_model_runner_class.return_value = mock_model_runner - executor.start() + + executor.start() self.assertEqual( executor.execute({"placeholder": "input"}), {"placeholder": "output"}, diff --git a/python/serving/serving_framework/requirements.in b/python/serving/serving_framework/requirements.in index a79e0d9..33fed85 100644 --- a/python/serving/serving_framework/requirements.in +++ b/python/serving/serving_framework/requirements.in @@ -1,14 +1,11 @@ -absl-py~=2.1.0 +absl-py>=2.1.0 flask~=3.0.3 grpcio~=1.68.1 grpcio-status~=1.68.1 gunicorn~=23.0.0 -numpy~=2.0.2 +numpy<=2.0.2 # bypassing faulty version restriction in tritonclient requests~=2.32.3 # TODO: b/375469331 - Enable testing with most current requests-mock release. requests-mock==1.9.3 -tensorflow~=2.18.0 -tensorflow-serving-api~=2.18.0 -tensorflow-io-gcs-filesystem>=0.23.1 setuptools~=75.6.0 typing-extensions~=4.12.2 diff --git a/python/serving/serving_framework/tensorflow/__init__.py b/python/serving/serving_framework/tensorflow/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python/serving/serving_framework/inline_model_runner.py b/python/serving/serving_framework/tensorflow/inline_model_runner.py similarity index 100% rename from python/serving/serving_framework/inline_model_runner.py rename to python/serving/serving_framework/tensorflow/inline_model_runner.py diff --git a/python/serving/serving_framework/inline_model_runner_test.py b/python/serving/serving_framework/tensorflow/inline_model_runner_test.py similarity index 99% rename from python/serving/serving_framework/inline_model_runner_test.py rename to python/serving/serving_framework/tensorflow/inline_model_runner_test.py index 2f6029a..337ac6d 100644 --- a/python/serving/serving_framework/inline_model_runner_test.py +++ b/python/serving/serving_framework/tensorflow/inline_model_runner_test.py @@ -19,7 +19,7 @@ from absl.testing import absltest from absl.testing import parameterized -from serving.serving_framework import inline_model_runner +from serving.serving_framework.tensorflow import inline_model_runner def tensor_equal(first: tf.Tensor, second: tf.Tensor) -> bool: diff --git a/python/serving/serving_framework/tensorflow/requirements.in b/python/serving/serving_framework/tensorflow/requirements.in new file mode 100644 index 0000000..b455760 --- /dev/null +++ b/python/serving/serving_framework/tensorflow/requirements.in @@ -0,0 +1,3 @@ +tensorflow~=2.18.0 +tensorflow-serving-api~=2.18.0 +tensorflow-io-gcs-filesystem>=0.23.1 \ No newline at end of file diff --git a/python/serving/serving_framework/server_model_runner.py b/python/serving/serving_framework/tensorflow/server_model_runner.py similarity index 100% rename from python/serving/serving_framework/server_model_runner.py rename to python/serving/serving_framework/tensorflow/server_model_runner.py diff --git a/python/serving/serving_framework/server_model_runner_test.py b/python/serving/serving_framework/tensorflow/server_model_runner_test.py similarity index 99% rename from python/serving/serving_framework/server_model_runner_test.py rename to python/serving/serving_framework/tensorflow/server_model_runner_test.py index c13b3d8..162ca00 100644 --- a/python/serving/serving_framework/server_model_runner_test.py +++ b/python/serving/serving_framework/tensorflow/server_model_runner_test.py @@ -18,7 +18,7 @@ import tensorflow as tf from absl.testing import absltest -from serving.serving_framework import server_model_runner +from serving.serving_framework.tensorflow import server_model_runner from tensorflow_serving.apis import predict_pb2 from tensorflow_serving.apis import prediction_service_pb2_grpc