Added PB Geomodel

Author: Josef-Haupt

docs/birdnet.geo_models.v3_0.rst

@@ -20,6 +20,14 @@ birdnet.geo\_models.v3\_0.tf module
    :show-inheritance:
    :undoc-members:
 
+birdnet.geo\_models.v3\_0.pb module
+------------------------------------
+
+.. automodule:: birdnet.geo_models.v3_0.pb
+   :members:
+   :show-inheritance:
+   :undoc-members:
+
 Module contents
 ---------------
 

src/birdnet/geo/models/v3_0/pb.py

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+import shutil
+import tempfile
+import zipfile
+from pathlib import Path
+
+from ordered_set import OrderedSet
+
+from birdnet.core.backends import (
+  PBBackend,
+  VersionedGeoBackendProtocol,
+)
+from birdnet.geo.models.v3_0.model import GeoDownloaderBaseV3_0
+from birdnet.globals import (
+  MODEL_PRECISION_FP32,
+  MODEL_PRECISIONS,
+)
+from birdnet.utils.helper import (
+  check_protobuf_model_files_exist,
+  download_file_tqdm,
+  get_species_from_file,
+)
+from birdnet.utils.local_data import get_lang_dir, get_model_path
+
+
+class GeoPBDownloaderV3_0(GeoDownloaderBaseV3_0):
+  @classmethod
+  def _get_paths(cls) -> tuple[Path, Path]:
+    model_path = get_model_path("geo", "3.0", "pb", MODEL_PRECISION_FP32)
+    lang_dir = get_lang_dir("geo", "3.0", "pb")
+    return model_path, lang_dir
+
+  @classmethod
+  def _check_geo_model_available(cls) -> bool:
+    model_path, lang_dir = cls._get_paths()
+
+    model_is_downloaded = True
+    model_is_downloaded &= model_path.is_dir()
+    model_is_downloaded &= check_protobuf_model_files_exist(model_path)
+
+    model_is_downloaded &= lang_dir.is_dir()
+    for lang in cls.AVAILABLE_LANGUAGES:
+      model_is_downloaded &= (lang_dir / f"{lang}.txt").is_file()
+
+    return model_is_downloaded
+
+  @classmethod
+  def _download_model(cls) -> None:
+    dl_url = "TODO"  # TODO: update with real zenodo URL once published
+    dl_size = 0  # TODO: update with real download size once published
+
+    with tempfile.TemporaryDirectory(prefix="birdnet_download") as temp_dir:
+      zip_download_path = Path(temp_dir) / "download.zip"
+      download_file_tqdm(
+        dl_url,
+        zip_download_path,
+        download_size=dl_size,
+        description="Downloading geo model v3.0 (pb)",
+      )
+
+      print("Extracting...")
+      extract_dir = Path(temp_dir) / "extracted"
+
+      with zipfile.ZipFile(zip_download_path, "r") as zip_ref:
+        zip_ref.extractall(extract_dir)
+
+      geo_model_dl_dir = extract_dir / "meta-model"  # TODO: update folder name if different
+      species_dl_dir = extract_dir / "labels"
+
+      geo_model_dir, geo_lang_dir = cls._get_paths()
+      geo_model_dir.parent.mkdir(parents=True, exist_ok=True)
+      shutil.rmtree(geo_model_dir, ignore_errors=True)
+      shutil.move(geo_model_dl_dir, geo_model_dir)
+
+      geo_lang_dir.parent.mkdir(parents=True, exist_ok=True)
+      shutil.rmtree(geo_lang_dir, ignore_errors=True)
+      shutil.move(species_dl_dir, geo_lang_dir)
+      print("Extracted.")
+
+  @classmethod
+  def get_model_path_and_labels(
+    cls,
+    lang: str,
+  ) -> tuple[Path, OrderedSet[str]]:
+    if not cls._check_geo_model_available():
+      cls._download_model()
+    assert cls._check_geo_model_available()
+
+    model_dir, langs_path = cls._get_paths()
+
+    lang_file = langs_path / f"{lang}.txt"
+    if not lang_file.is_file():
+      raise ValueError(f"Language does not exist: {lang}")
+
+    labels = get_species_from_file(lang_file, encoding="utf8")
+    return model_dir, labels
+
+
+class GeoPBBackendFP32V3_0(PBBackend, VersionedGeoBackendProtocol):
+  def __init__(
+    self,
+    model_path: Path,
+    device_name: str,
+    half_precision: bool,
+  ) -> None:
+    super().__init__(model_path, device_name, half_precision)
+
+  @classmethod
+  def input_key(cls) -> str:
+    return "MNET_INPUT"  # TODO: update if key name differs in v3.0
+
+  @classmethod
+  def prediction_signature_name(cls) -> str:
+    return "serving_default"
+
+  @classmethod
+  def prediction_key(cls) -> str:
+    return "MNET_CLASS_ACTIVATION"  # TODO: update if key name differs in v3.0
+
+  @classmethod
+  def supports_encoding(cls) -> bool:
+    return False
+
+  @classmethod
+  def encoding_signature_name(cls) -> str | None:
+    return None
+
+  @classmethod
+  def encoding_key(cls) -> str | None:
+    return None
+
+  @classmethod
+  def precision(cls) -> MODEL_PRECISIONS:
+    return MODEL_PRECISION_FP32

src/birdnet/model_loader.py

@@ -51,6 +51,7 @@
 from birdnet.geo.models.v3_0.model import (
   GeoModelV3_0,
 )
+from birdnet.geo.models.v3_0.pb import GeoPBBackendFP32V3_0, GeoPBDownloaderV3_0
 from birdnet.geo.models.v3_0.tf import (
   GeoTFBackendFP16V3_0,
   GeoTFBackendFP32V3_0,
@@ -715,9 +716,20 @@ def _load_geo_model_V3_0(
       },
     )
   elif backend == MODEL_BACKEND_PB:
-    raise ValueError(
-      "The geo model v3.0 does not support the 'pb' backend. "
-      "Use 'tf' instead."
+    if precision != MODEL_PRECISION_FP32:
+      raise ValueError(
+        f"Unsupported model precision for geo model: {precision}. "
+        f"Currently supported precision is: {MODEL_PRECISION_FP32}."
+      )
+
+    model_kwargs = _validate_kwargs_allowed(model_kwargs, None)
+
+    model_path, species_list = GeoPBDownloaderV3_0.get_model_path_and_labels(lang)
+    return GeoModelV3_0.load(
+      model_path,
+      species_list,
+      backend_type=GeoPBBackendFP32V3_0,
+      backend_kwargs={},
     )
   else:
     raise AssertionError()
@@ -755,9 +767,21 @@ def _load_custom_geo_model_V3_0(
       check_validity=check_validity,
     )
   elif backend == MODEL_BACKEND_PB:
-    raise ValueError(
-      "The geo model v3.0 does not support the 'pb' backend. "
-      "Use 'tf' instead."
+    if precision != MODEL_PRECISION_FP32:
+      raise ValueError(
+        f"Unsupported model precision for geo model: {precision}. "
+        f"Currently supported precision is: {MODEL_PRECISION_FP32}."
+      )
+
+    model = _validate_pb_model_folder(model)
+    model_kwargs = _validate_kwargs_allowed(model_kwargs, None)
+
+    return GeoModelV3_0.load_custom(
+      model,
+      species_list,
+      backend_type=GeoPBBackendFP32V3_0,
+      backend_kwargs={},
+      check_validity=check_validity,
     )
   else:
     raise AssertionError()

src/birdnet_tests/geo_models/v3_0/model_py/test_predict/test_geo_predict_model.py

@@ -139,3 +139,49 @@ def test_tf_fp32_gpu() -> None:
   assert result.model_path == model.model_path.absolute()
   assert result.model_version == "3.0"
   assert result.model_precision == "fp32"
+
+
+# --- PB ---
+
+
+def test_pb_cpu() -> None:
+  model = load("geo", "3.0", "pb", precision="fp32")
+  result = model.predict(20, 50, week=1, min_confidence=0.03, half_precision=False)
+
+  assert result.latitude == 20
+  assert result.longitude == 50
+  assert result.week == 1
+  assert result.model_path == model.model_path.absolute()
+  assert result.model_version == "3.0"
+  assert result.model_precision == "fp32"
+
+
+@pytest.mark.gpu
+def test_pb_gpu() -> None:
+  ensure_gpu_or_skip()
+
+  model = load("geo", "3.0", "pb", precision="fp32")
+  result = model.predict(
+    20, 50, week=1, min_confidence=0.03, half_precision=False, device="GPU"
+  )
+
+  assert result.latitude == 20
+  assert result.longitude == 50
+  assert result.week == 1
+  assert result.model_path == model.model_path.absolute()
+  assert result.model_version == "3.0"
+  assert result.model_precision == "fp32"
+
+
+def test_pb_cpu_half() -> None:
+  model = load("geo", "3.0", "pb", precision="fp32")
+  result = model.predict(
+    20, 50, week=1, min_confidence=0.03, half_precision=True, device="CPU"
+  )
+
+  assert result.latitude == 20
+  assert result.longitude == 50
+  assert result.week == 1
+  assert result.model_path == model.model_path.absolute()
+  assert result.model_version == "3.0"
+  assert result.model_precision == "fp32"

src/birdnet_tests/geo_models/v3_0/pb_py/__init__.py

@@ -0,0 +1 @@
+#

src/birdnet_tests/geo_models/v3_0/pb_py/test_download_geo_pb_model.py

@@ -0,0 +1,11 @@
+from birdnet.geo.models.v3_0.pb import GeoPBDownloaderV3_0
+
+
+def xtest_double_download() -> None:
+  # takes too long to run normally
+  GeoPBDownloaderV3_0._download_model()
+  GeoPBDownloaderV3_0._download_model()
+
+
+if __name__ == "__main__":
+  xtest_double_download()

src/birdnet_tests/model_loader_py/geo/test_load_custom_v3_0.py

@@ -3,6 +3,7 @@
 import pytest
 
 from birdnet.geo.models.v3_0.model import GeoModelV3_0
+from birdnet.geo.models.v3_0.pb import GeoPBDownloaderV3_0
 from birdnet.geo.models.v3_0.tf import GeoTFDownloaderV3_0
 from birdnet.model_loader import load_custom
 from birdnet.utils.local_data import get_lang_dir, get_model_path
@@ -108,3 +109,53 @@ def test_tf_type_with_precisions_is_correct() -> None:
     )
     is GeoModelV3_0
   )
+
+
+def test_load_custom_geo_model_v3_0_pb_fp32() -> None:
+  GeoPBDownloaderV3_0.get_model_path_and_labels("en_us")
+  model = load_custom(
+    "geo",
+    "3.0",
+    "pb",
+    get_model_path("geo", "3.0", "pb", "fp32"),
+    get_lang_dir("geo", "3.0", "pb") / "en_us.txt",
+    precision="fp32",
+    check_validity=check_validity(),
+  )
+  assert isinstance(model, GeoModelV3_0)
+
+
+def test_pb_type_is_correct() -> None:
+  GeoPBDownloaderV3_0.get_model_path_and_labels("en_us")
+  assert (
+    type(
+      load_custom(
+        "geo",
+        "3.0",
+        "pb",
+        get_model_path("geo", "3.0", "pb", "fp32"),
+        get_lang_dir("geo", "3.0", "pb") / "en_us.txt",
+        check_validity=False,
+      )
+    )
+    is GeoModelV3_0
+  )
+
+
+def test_load_pb_with_custom_library_raises_error() -> None:
+  ensure_litert_or_skip()
+
+  with pytest.raises(
+    ValueError,
+    match=r"Unexpected keyword arguments: library.",
+  ):
+    GeoPBDownloaderV3_0.get_model_path_and_labels("en_us")
+    load_custom(
+      "geo",
+      "3.0",
+      "pb",  # type: ignore
+      get_model_path("geo", "3.0", "pb", "fp32"),
+      get_lang_dir("geo", "3.0", "pb") / "en_us.txt",
+      library="litert",
+      check_validity=check_validity(),
+    )  # type: ignore

src/birdnet_tests/model_loader_py/geo/test_load_v3_0.py

@@ -12,8 +12,20 @@
 def test_pb_v3_0_with_library_raises_error() -> None:
   ensure_litert_or_skip()
 
-  with pytest.raises(ValueError):
-    load("geo", "3.0", "pb", precision="fp32")
+  with pytest.raises(
+    ValueError,
+    match=r"Unexpected keyword arguments: library.",
+  ):
+    load("geo", "3.0", "pb", precision="fp32", library="litert")  # type: ignore
+
+
+@pytest.mark.load_model
+def test_v3_0_pb() -> None:
+  try:
+    model = load("geo", "3.0", "pb", precision="fp32")
+  except ReadTimeout as e:
+    pytest.fail(f"Model download timed out: {e}. Try again later.")
+  assert isinstance(model, GeoModelV3_0)
 
 
 @pytest.mark.load_model
@@ -91,3 +103,14 @@ def test_types_with_precisions_are_correct() -> None:
     type(load("geo", "3.0", "tf", precision=cast(Literal["int8"], f"int{8}")))
     is GeoModelV3_0
   )
+
+
+def test_pb_type_is_correct() -> None:
+  assert type(load("geo", "3.0", "pb")) is GeoModelV3_0
+
+
+def test_pb_type_with_precision_is_correct() -> None:
+  assert (
+    type(load("geo", "3.0", "pb", precision=cast(Literal["fp32"], f"fp{32}")))
+    is GeoModelV3_0
+  )