Upload models and script

birdnet_backbone.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:130cb8a574141ee97c49528319f9df33f14b3c047cbed86555aa7c9dc7a41417
+size 40160497

extract_backbone.py

@@ -0,0 +1,306 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Extract and validate BirdNET v2.4 ONNX backbone models.
+
+Downloads model.onnx and birdnet.onnx from HuggingFace
+(justinchuby/BirdNET-onnx), strips the classification head, and saves:
+  - model_backbone.onnx
+  - birdnet_backbone.onnx
+
+Also downloads the reference TF SavedModel from Zenodo
+(BirdNET_v2.4_protobuf) and verifies that embeddings match.
+"""
+
+import io
+import os
+import urllib.request
+import zipfile
+
+import huggingface_hub
+import numpy as np
+import onnx
+import onnx.helper
+import onnxruntime as ort
+import tensorflow as tf
+
+# Suppress TF C++ info/warning logs; only errors are shown.
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+
+# Source HuggingFace repo that hosts the full BirdNET v2.4 ONNX models.
+HF_REPO_ID = "justinchuby/BirdNET-onnx"
+
+# Zenodo URL for the BirdNET v2.4 protobuf SavedModel archive.
+ZENODO_URL = "https://zenodo.org/records/15050749/files/BirdNET_v2.4_protobuf.zip?download=1"
+
+# Sub-directory inside the Zenodo zip that contains the audio SavedModel.
+AUDIO_MODEL_ZIP_PREFIX = "audio-model/"
+
+# Internal tensor name of the global-average-pool output — the last node of
+# the backbone, immediately before the classification dense layer.
+BACKBONE_RAW_OUTPUT = "model/GLOBAL_AVG_POOL/Mean_reduced_0"
+
+# Public name exposed by the extracted backbone model.
+BACKBONE_OUTPUT = "embedding"
+
+# Expected number of audio samples fed to the model (3 s at 48 kHz).
+BIRDNET_SAMPLE_LEN = 144000
+
+# Tolerances for np.testing.assert_allclose when comparing ONNX vs TF outputs.
+# birdnet.onnx is a separate ONNX export whose weights differ slightly from the
+# reference SavedModel, so a loose tolerance is used to accommodate both variants.
+RTOL = 1e-3
+ATOL = 1e-3
+
+
+# ---------------------------------------------------------------------------
+# Download helpers
+# ---------------------------------------------------------------------------
+
+def download_onnx_models(output_dir: str) -> dict[str, str]:
+    """Download model.onnx and birdnet.onnx from HuggingFace.
+
+    Returns a dict mapping filename -> absolute local path.
+    """
+    filenames = ["model.onnx", "birdnet.onnx"]
+    paths = {}
+    for fname in filenames:
+        path = huggingface_hub.hf_hub_download(
+            repo_id=HF_REPO_ID,
+            filename=fname,
+            local_dir=output_dir,
+        )
+        paths[fname] = path
+        print(f"Downloaded {fname} -> {path}")
+    return paths
+
+
+def download_pb_model(output_dir: str) -> str:
+    """Download BirdNET_v2.4_protobuf.zip from Zenodo and extract audio-model.
+
+    The zip contains two SavedModel sub-directories; only audio-model is
+    extracted since that is the one whose embeddings signature we compare
+    against.
+
+    Returns the path to the extracted audio-model SavedModel directory.
+    """
+    audio_model_dir = os.path.join(output_dir, "audio-model")
+    if os.path.isdir(audio_model_dir):
+        print(f"Protobuf already extracted -> {audio_model_dir}")
+        return audio_model_dir
+
+    print(f"Downloading BirdNET protobuf from Zenodo...")
+    with urllib.request.urlopen(ZENODO_URL) as response:
+        data = response.read()
+    print(f"Download complete ({len(data) / 1_000_000:.1f} MB)")
+
+    with zipfile.ZipFile(io.BytesIO(data)) as zf:
+        members = [m for m in zf.namelist() if m.startswith(AUDIO_MODEL_ZIP_PREFIX)]
+        zf.extractall(output_dir, members=members)
+
+    print(f"Extracted audio-model -> {audio_model_dir}")
+    return audio_model_dir
+
+
+# ---------------------------------------------------------------------------
+# Backbone extraction
+# ---------------------------------------------------------------------------
+
+def _extract(
+    src_path: str,
+    out_path: str,
+    input_names: list[str],
+    output_names: list[str],
+    output_renames: dict[str, str] | None = None,
+) -> None:
+    """Extract a subgraph from an ONNX model using backwards BFS and save it.
+
+    Starting from `output_names`, the algorithm traces each tensor back through
+    the graph to find every node that contributes to those outputs. Nodes that
+    only serve the classification head (i.e., downstream of `output_names`) are
+    never reached and are therefore excluded from the new model.
+
+    Args:
+        src_path: Path to the source ONNX model file.
+        out_path: Destination path for the extracted subgraph.
+        input_names: Graph-level input tensor names to keep (weight initializers
+            that appear in graph.input are automatically excluded).
+        output_names: Tensor names that define the extraction boundary — the new
+            model will produce exactly these tensors as outputs.
+        output_renames: Optional mapping {old_name: new_name} applied to the
+            output tensors of the producing nodes before saving.
+    """
+    model = onnx.load(src_path)
+    renames = output_renames or {}
+
+    # Build a reverse lookup: tensor name -> the node that produces it.
+    tensor_to_node: dict = {}
+    for node in model.graph.node:
+        for out in node.output:
+            if out:
+                tensor_to_node[out] = node
+
+    # BFS backwards from the requested outputs to collect all contributing nodes.
+    visited_node_ids: set = set()
+    queue = list(output_names)
+    while queue:
+        tensor = queue.pop()
+        node = tensor_to_node.get(tensor)
+        if node is None or id(node) in visited_node_ids:
+            continue
+        visited_node_ids.add(id(node))
+        for inp in node.input:
+            if inp:
+                queue.append(inp)
+
+    # Re-filter from the original node list to preserve topological order.
+    filtered_nodes = [n for n in model.graph.node if id(n) in visited_node_ids]
+
+    # Apply any requested output renames directly on the producing nodes.
+    for node in filtered_nodes:
+        for i, out in enumerate(node.output):
+            if out in renames:
+                node.output[i] = renames[out]
+
+    # Collect only the initializers consumed by the retained nodes.
+    needed_tensors: set = set()
+    for node in filtered_nodes:
+        needed_tensors.update(i for i in node.input if i)
+    filtered_inits = [i for i in model.graph.initializer if i.name in needed_tensors]
+
+    # Keep only the declared data inputs (skip weight aliases in graph.input).
+    input_name_set = set(input_names)
+    graph_inputs = [vi for vi in model.graph.input if vi.name in input_name_set]
+
+    # Build output ValueInfoProtos.
+    existing_out = {o.name: o for o in model.graph.output}
+    graph_outputs = []
+    for name in output_names:
+        final_name = renames.get(name, name)
+        if final_name in existing_out:
+            graph_outputs.append(existing_out[final_name])
+        else:
+            graph_outputs.append(
+                onnx.helper.make_tensor_value_info(final_name, onnx.TensorProto.FLOAT, None)
+            )
+
+    new_graph = onnx.helper.make_graph(
+        filtered_nodes,
+        "backbone",
+        graph_inputs,
+        graph_outputs,
+        initializer=filtered_inits,
+    )
+    new_model = onnx.helper.make_model(new_graph)
+    new_model.ir_version = model.ir_version
+    del new_model.opset_import[:]
+    new_model.opset_import.extend(model.opset_import)
+    onnx.save(new_model, out_path)
+
+
+def _get_graph_input_names(onnx_path: str) -> list[str]:
+    """Return the true data-input tensor names for an ONNX model."""
+    model = onnx.load(onnx_path)
+    init_names = {i.name for i in model.graph.initializer}
+    return [vi.name for vi in model.graph.input if vi.name not in init_names]
+
+
+def extract_backbone(src_path: str, out_path: str) -> str:
+    """Extract the backbone subgraph from a full BirdNET ONNX model and save it.
+
+    Traces backwards from BACKBONE_RAW_OUTPUT (the global average pool tensor)
+    and renames it to BACKBONE_OUTPUT ("embedding") in the saved file.
+
+    Returns out_path for chaining.
+    """
+    input_names = _get_graph_input_names(src_path)
+    _extract(
+        src_path,
+        out_path,
+        input_names,
+        [BACKBONE_RAW_OUTPUT],
+        output_renames={BACKBONE_RAW_OUTPUT: BACKBONE_OUTPUT},
+    )
+
+    model = onnx.load(out_path)
+    print(f"Backbone saved -> {out_path}")
+    print(f"  inputs : {input_names}")
+    print(f"  outputs: {[o.name for o in model.graph.output]}")
+    return out_path
+
+
+# ---------------------------------------------------------------------------
+# Comparison helpers
+# ---------------------------------------------------------------------------
+
+def _make_audio(length: int, seed: int = 42) -> np.ndarray:
+    """Generate a reproducible Gaussian noise waveform shaped (1, length)."""
+    rng = np.random.default_rng(seed)
+    return rng.standard_normal((1, length)).astype(np.float32)
+
+
+def _onnx_embedding(onnx_path: str, audio: np.ndarray) -> np.ndarray:
+    """Run inference on a backbone ONNX model and return the embedding array."""
+    input_names = _get_graph_input_names(onnx_path)
+    sess = ort.InferenceSession(onnx_path)
+    (emb,) = sess.run([BACKBONE_OUTPUT], {input_names[0]: audio})
+    return emb
+
+
+def _pb_embedding(pb_dir: str, audio: np.ndarray) -> np.ndarray:
+    """Run inference on the BirdNET TF SavedModel and return the embedding array.
+
+    The audio-model SavedModel exposes an "embeddings" signature whose output
+    dict contains an "embeddings" key, used here as the ground-truth reference.
+    """
+    model = tf.saved_model.load(pb_dir)
+    audio_tf = tf.constant(audio)
+    return model.signatures["embeddings"](inputs=audio_tf)["embeddings"].numpy()
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+def main():
+    """End-to-end pipeline: download → extract → compare."""
+    out_dir = os.path.dirname(os.path.abspath(__file__))
+
+    # --- Step 1: download source models ---
+    print("=== Downloading models ===")
+    onnx_paths = download_onnx_models(out_dir)
+    pb_dir = download_pb_model(out_dir)
+
+    # --- Step 2: extract backbone from each ONNX variant ---
+    print("\n=== Extracting backbones ===")
+    backbone_paths = {}
+    for fname, src in onnx_paths.items():
+        stem = fname.replace(".onnx", "")
+        out_path = os.path.join(out_dir, f"{stem}_backbone.onnx")
+        backbone_paths[stem] = extract_backbone(src, out_path)
+
+    # --- Step 3: numerical comparison against the TF SavedModel reference ---
+    print("\n=== Comparing embeddings against Zenodo TF SavedModel ===")
+    audio = _make_audio(BIRDNET_SAMPLE_LEN)
+
+    pb_emb = _pb_embedding(pb_dir, audio)
+    print(f"PB embedding shape: {pb_emb.shape}")
+
+    for stem, path in backbone_paths.items():
+        onnx_emb = _onnx_embedding(path, audio)
+        diff = np.abs(onnx_emb - pb_emb)
+        print(f"\n{stem}_backbone.onnx:")
+        print(f"  ONNX embedding shape: {onnx_emb.shape}")
+        print(f"  |diff| mean={diff.mean():.6e}  max={diff.max():.6e}")
+        try:
+            np.testing.assert_allclose(onnx_emb, pb_emb, rtol=RTOL, atol=ATOL)
+            print(f"  Embeddings match PB reference with rtol={RTOL:.0e}, atol={ATOL:.0e}  PASSED")
+        except AssertionError as e:
+            print(f"  Embeddings differ from PB reference  FAILED\n  {e}")
+
+    print("\nDone.")
+
+
+if __name__ == "__main__":
+    main()

model_backbone.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:015344d3662d262e56dc52975523f1ea9b5e3852c1fd8ceb789f5cbbfba1dc25
+size 24971712

pyproject.toml

@@ -0,0 +1,13 @@
+[project]
+name = "birdnet-onnx-backbone"
+version = "0.1.0"
+description = "Backbone-only ONNX exports of BirdNET v2.4 bird sound classifier"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "huggingface-hub>=0.23.0",
+    "numpy>=2.0.0",
+    "onnx>=1.16.0",
+    "onnxruntime>=1.18.0",
+    "tensorflow-cpu>=2.16.0",
+]