diff --git a/headmic.py b/headmic.py index 4e54783..2888c26 100644 --- a/headmic.py +++ b/headmic.py @@ -469,8 +469,8 @@ async def startup(): edgetpu_model_path = model_dir / "yamnet_edgetpu.tflite" model_path = model_dir / "yamnet.tflite" class_map_path = model_dir / "yamnet_class_map.csv" - # Edge TPU opt-in via env var (segfaults with some libedgetpu versions) - use_edgetpu = os.environ.get("USE_EDGETPU", "").lower() in ("1", "true", "yes") and edgetpu_model_path.exists() + # Auto-detect Edge TPU — probe in subprocess catches segfaults safely + 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: diff --git a/sound_id.py b/sound_id.py index 4b9ea40..da631c0 100644 --- a/sound_id.py +++ b/sound_id.py @@ -85,6 +85,29 @@ CATEGORY_GROUPS = { class SoundClassifier: + @staticmethod + def _probe_edgetpu(model_path: str) -> bool: + """Test Edge TPU in a subprocess to catch segfaults safely.""" + import subprocess, sys + try: + result = subprocess.run( + [sys.executable, "-c", + "import ai_edge_litert.interpreter as tfl; " + f"d = tfl.load_delegate('libedgetpu.so.1'); " + f"i = tfl.Interpreter(model_path='{model_path}', experimental_delegates=[d]); " + "i.allocate_tensors(); " + "print('ok')"], + capture_output=True, text=True, timeout=10 + ) + if result.returncode == 0 and "ok" in result.stdout: + logger.info("Edge TPU probe: OK") + return True + logger.warning("Edge TPU probe failed: %s", result.stderr.strip() or f"exit {result.returncode}") + return False + except Exception as e: + logger.warning("Edge TPU probe error: %s", e) + return False + def __init__(self, model_path, class_map_path, use_edgetpu=False): # Load class names self._class_names = [] @@ -105,14 +128,21 @@ class SoundClassifier: import ai_edge_litert.interpreter as tfl if use_edgetpu: - delegate = tfl.load_delegate("libedgetpu.so.1") - self._interp = tfl.Interpreter( - model_path=str(model_path), - experimental_delegates=[delegate], - ) - logger.info("YAMNet loaded on Edge TPU") - else: - self._interp = tfl.Interpreter(model_path=str(model_path)) + if self._probe_edgetpu(model_path): + delegate = tfl.load_delegate("libedgetpu.so.1") + self._interp = tfl.Interpreter( + model_path=str(model_path), + experimental_delegates=[delegate], + ) + logger.info("YAMNet loaded on Edge TPU") + else: + logger.warning("Edge TPU probe failed (segfault or error) — falling back to CPU") + use_edgetpu = False + + if not use_edgetpu: + # Use CPU model (swap edgetpu model path for CPU model if needed) + cpu_path = str(model_path).replace("_edgetpu.tflite", ".tflite") + self._interp = tfl.Interpreter(model_path=cpu_path) logger.info("YAMNet loaded on CPU") self._interp.allocate_tensors()