diff --git a/headmic.py b/headmic.py index b41896c..87e01f6 100644 --- a/headmic.py +++ b/headmic.py @@ -465,15 +465,19 @@ async def startup(): # --- Sound classifier (optional) --- model_dir = Path(__file__).parent / "models" + edgetpu_model_path = model_dir / "yamnet_edgetpu.tflite" model_path = model_dir / "yamnet.tflite" class_map_path = model_dir / "yamnet_class_map.csv" - if model_path.exists() and class_map_path.exists(): + # Prefer Edge TPU model if available + use_edgetpu = edgetpu_model_path.exists() + active_model = edgetpu_model_path if use_edgetpu else model_path + if active_model.exists() and class_map_path.exists(): try: from sound_id import SoundClassifier - sound_classifier = SoundClassifier(str(model_path), str(class_map_path)) + sound_classifier = SoundClassifier(str(active_model), str(class_map_path), use_edgetpu=use_edgetpu) sound_ring_buffer = collections.deque(maxlen=31) state.sound_classification_enabled = True - logger.info("Sound classification enabled (YAMNet)") + logger.info("Sound classification enabled (YAMNet %s)", "Edge TPU" if use_edgetpu else "CPU") sc_thread = threading.Thread(target=sound_classifier_loop, daemon=True) sc_thread.start()