Skip to content

Commit

Permalink
Add yamnet and vggish configs to convenience config builder fn.
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 681892178
  • Loading branch information
sdenton4 authored and copybara-github committed Oct 3, 2024
1 parent f709582 commit 657c89c
Showing 1 changed file with 14 additions and 0 deletions.
14 changes: 14 additions & 0 deletions chirp/projects/zoo/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,20 @@ def get_preset_model_config(preset_name):
# Note: The v2_1 class list is appropriate for Birdnet 2.1, 2.2, and 2.3.
model_config.class_list_name = 'birdnet_v2_1'
model_config.num_tflite_threads = 4
elif preset_name == 'yamnet':
model_key = 'tfhub_model'
embedding_dim = 1024
model_config.sample_rate = 16000
model_config.model_url = 'https://tfhub.dev/google/yamnet/1'
model_config.embedding_index = 1
model_config.logits_index = 0
elif preset_name == 'vggish':
model_key = 'tfhub_model'
embedding_dim = 128
model_config.sample_rate = 16000
model_config.model_url = 'https://tfhub.dev/google/vggish/1'
model_config.embedding_index = -1
model_config.logits_index = -1
else:
raise ValueError('Unsupported model preset: %s' % preset_name)
return model_key, embedding_dim, model_config
Expand Down

0 comments on commit 657c89c

Please sign in to comment.