Pet-Monitor/pet_audio.py

"""
Detección de ladridos sobre el audio RTSP de la cámara Tapo.

La cámara publica su micrófono dentro del mismo stream RTSP que el vídeo, en
formato G.711 A-law a 8 kHz mono. OpenCV no sabe decodificar audio, así que
aquí se abre una segunda conexión RTSP con PyAV (la cámara admite varias en
paralelo, verificado con el vídeo funcionando a la vez).

El audio se remuestrea a los 16 kHz que espera YAMNet y se analiza en ventanas
de 0.975 s. Medido en la Pi 5: ~4.6 ms por ventana, unas 200 veces más rápido
que el tiempo real, así que no compite de forma apreciable con la NPU.
"""

import os
import threading
import time

import av
import numpy as np
import onnxruntime as ort

from pet_common import mask_secrets, secrets_of

SAMPLE_RATE = 16000
WINDOW_SAMPLES = 15600  # 0.975 s, la ventana que espera YAMNet
HOP_SAMPLES = 7800      # medio solape entre ventanas consecutivas

# Índices en AudioSet (verificados contra models/yamnet_classes.csv)
BARK_CLASS = 70
CANINE_GROUP = (69, 70, 71, 72, 73, 74, 75, 117)  # Dog, Bark, Yip, Howl, Bow-wow, Growling, Whimper (dog), Canidae

MODEL_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models", "yamnet.onnx")


class BarkDetector:
    """
    YAMNet en ONNX: puntúa una ventana de audio como ladrido.

    Se dispara por la clase `Bark` sola, no por el grupo canino: con el mismo
    material, `Bark` separa perros de gatos con un margen 6 veces mayor que la
    suma del grupo, porque las clases genéricas (`Animal`, `Whimper`) también
    se activan con maullidos. El grupo se calcula solo como contexto al avisar.
    """

    def __init__(self, model_path=MODEL_PATH):
        if not os.path.exists(model_path):
            raise FileNotFoundError(
                f"No se encontró {model_path}.\n"
                "Descárgalo con:  bash scripts/fetch_models.sh"
            )
        self.session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
        self.input_name = self.session.get_inputs()[0].name

    def analyze(self, samples):
        """
        samples: array float32 mono a 16 kHz de exactamente 0.975 s.
        Devuelve (puntuación de ladrido, suma del grupo canino), ambas en [0, 1+] .
        """
        scores = np.asarray(self.session.run(None, {self.input_name: samples})[0]).reshape(-1)
        return float(scores[BARK_CLASS]), float(scores[list(CANINE_GROUP)].sum())


class AudioListener:
    """
    Hilo que escucha el audio RTSP y avisa cuando detecta un ladrido.

    Llama a `on_bark(puntuacion, grupo, cuando)` como máximo una vez cada
    `cooldown` segundos, para no inundar de avisos un ladrido sostenido.
    """

    def __init__(
        self,
        url,
        detector,
        on_bark=None,
        threshold=0.15,
        cooldown=15.0,
        reconnect_delay=2.0,
        max_reconnect_delay=30.0,
        socket_timeout_us=5_000_000,
    ):
        self._url = url
        self._detector = detector
        self._on_bark = on_bark
        self.threshold = threshold
        self.cooldown = cooldown
        self._reconnect_delay = reconnect_delay
        self._max_reconnect_delay = max_reconnect_delay
        self._socket_timeout_us = socket_timeout_us
        self._secrets = secrets_of(url)
        self._stop = threading.Event()
        self._thread = None
        self._last_alert = 0.0

        # Estadísticas consultables desde el bucle principal
        self.windows = 0
        self.alerts = 0
        self.reconnects = 0
        self.connected = False
        self.max_bark = 0.0
        self.last_error = None

    def start(self):
        self._thread = threading.Thread(target=self._loop, name="audio-listener", daemon=True)
        self._thread.start()
        return self

    def stop(self):
        self._stop.set()
        if self._thread is not None:
            self._thread.join(timeout=5.0)

    def take_max_bark(self):
        """Puntuación máxima de ladrido desde la última consulta (para las estadísticas)."""
        valor, self.max_bark = self.max_bark, 0.0
        return valor

    # --- interno ---------------------------------------------------------
    def _loop(self):
        delay = self._reconnect_delay
        while not self._stop.is_set():
            try:
                self._captura()
            except Exception as exc:  # noqa: BLE001 - el hilo nunca debe morir
                mensaje = mask_secrets(exc, self._secrets)
                self.last_error = mensaje
                self.connected = False
                if self._stop.is_set():
                    break
                print(f"[WARN] Audio RTSP interrumpido ({mensaje}); reintentando en {delay:.0f}s", flush=True)
                self._stop.wait(delay)
                delay = min(delay * 2, self._max_reconnect_delay)
                self.reconnects += 1
            else:
                delay = self._reconnect_delay

    def _captura(self):
        """Una sesión de escucha completa; termina cuando se para o falla la conexión."""
        container = av.open(
            self._url,
            options={"rtsp_transport": "tcp", "timeout": str(self._socket_timeout_us)},
        )
        try:
            audio = next((s for s in container.streams if s.type == "audio"), None)
            if audio is None:
                raise RuntimeError("la cámara no publica pista de audio")

            self.connected = True
            print("[INFO] Audio RTSP conectado.", flush=True)

            resampler = av.AudioResampler(format="s16", layout="mono", rate=SAMPLE_RATE)
            buffer = np.empty(0, dtype=np.int16)

            for packet in container.demux(audio):
                if self._stop.is_set():
                    return
                for frame in packet.decode():
                    for remuestreado in resampler.resample(frame):
                        buffer = np.concatenate([buffer, remuestreado.to_ndarray().reshape(-1)])
                while len(buffer) >= WINDOW_SAMPLES:
                    self._procesa(buffer[:WINDOW_SAMPLES])
                    buffer = buffer[HOP_SAMPLES:]
        finally:
            self.connected = False
            container.close()

    def _procesa(self, ventana):
        samples = ventana.astype(np.float32) / 32768.0
        ladrido, grupo = self._detector.analyze(samples)
        self.windows += 1
        if ladrido > self.max_bark:
            self.max_bark = ladrido

        ahora = time.time()
        if ladrido >= self.threshold and ahora - self._last_alert >= self.cooldown:
            self._last_alert = ahora
            self.alerts += 1
            if self._on_bark is not None:
                self._on_bark(ladrido, grupo, ahora)