440 lines
23 KiB
Python
440 lines
23 KiB
Python
# vad_core.py
|
||
import pyaudio
|
||
import webrtcvad # type: ignore
|
||
import collections
|
||
import threading
|
||
import wave
|
||
import time
|
||
import os
|
||
import struct
|
||
import numpy as np
|
||
import logging
|
||
import sys
|
||
from datetime import datetime
|
||
from app.core.status import status_tracker, SystemStatus
|
||
|
||
class VADPolling:
|
||
def __init__(self, config, on_speech_end_callback=None):
|
||
self.config = config
|
||
self.vad = webrtcvad.Vad()
|
||
self.vad.set_mode(config['mode'])
|
||
self.ring_buffer = collections.deque(maxlen=config['ring_buffer_size'])
|
||
self.history = collections.deque(maxlen=config.get('history_maxlen', 100))
|
||
self.lock = threading.Lock()
|
||
|
||
# Параметры шумоподавления и автоусиления
|
||
self.noise_reduction_enabled = config.get('noise_reduction_enabled', True)
|
||
self.noise_floor_buffer = collections.deque(maxlen=50) # Буфер для оценки уровня шума
|
||
self.noise_floor = 0.0
|
||
self.auto_gain_enabled = config.get('auto_gain_enabled', True)
|
||
self.target_rms = config.get('target_rms', 2000.0) # Целевой уровень сигнала
|
||
self.current_gain = 1.0
|
||
|
||
self.frame_size = int(config['sample_rate'] * config['frame_duration_ms'] / 1000)
|
||
|
||
# 🔹 Размер предбуфера (в фреймах)
|
||
self.prebuffer_size = int(config.get('prebuffer_duration_sec', 1.0) * config['sample_rate'] / self.frame_size)
|
||
self.prebuffer = collections.deque(maxlen=self.prebuffer_size)
|
||
|
||
# 🔹 Основной буфер записи (используется только во время записи)
|
||
self.recording_buffer = []
|
||
|
||
# 🔹 Максимальная длина записи
|
||
self.max_record_duration_sec = config.get('max_record_duration_sec', 5.0)
|
||
self.max_record_frames = int(self.max_record_duration_sec * config['sample_rate'] / self.frame_size)
|
||
|
||
# 🔹 Фильтрация по энергии и длительности
|
||
self.energy_threshold = config.get('energy_threshold', 1000)
|
||
self.min_speech_duration_ms = config.get('min_speech_duration_ms', 300)
|
||
self.adaptive_energy = config.get('adaptive_energy', True)
|
||
|
||
# 🔹 Адаптивный порог энергии
|
||
self.background_energy_buffer = collections.deque(maxlen=100)
|
||
self.current_energy_threshold = self.energy_threshold
|
||
self.adaptive_energy_std_factor = config.get('adaptive_energy_std_factor', 2.5) # Множитель для стандартного отклонения
|
||
|
||
# 🔹 Счетчики для фильтрации
|
||
self.speech_start_time = None
|
||
self.continuous_speech_frames = 0
|
||
self.min_speech_frames = int(self.min_speech_duration_ms / config['frame_duration_ms'])
|
||
|
||
# 🔹 Хранение последнего состояния
|
||
self.last_result = {
|
||
'speech_ratio': 0.0,
|
||
'speech_detected': False,
|
||
'energy': 0.0,
|
||
'energy_threshold': float(self.current_energy_threshold),
|
||
'speech_threshold': float(self.config['speech_threshold']),
|
||
'continuous_frames': 0,
|
||
|
||
}
|
||
|
||
# 🔹 Состояния записи
|
||
self.last_speech_detected = False
|
||
self.prev_speech_detected = False
|
||
self.is_recording = False
|
||
|
||
# 🔹 Настройка записи
|
||
self.record_enabled = config.get('record_on_speech', False)
|
||
self.record_stop_delay_sec = config.get('record_stop_delay_sec', 1.0)
|
||
self.output_filename = "recordings/VAD/debug_recording.wav"
|
||
|
||
# 🔹 Таймер остановки записи
|
||
self.stop_timer = None
|
||
|
||
# 🔹 Callback для обработки речи после завершения записи
|
||
self.on_speech_end_callback = on_speech_end_callback
|
||
|
||
# 🔹 Удаляем старый файл при запуске, если запись включена
|
||
if self.record_enabled:
|
||
os.makedirs("recordings", exist_ok=True)
|
||
if os.path.exists(self.output_filename):
|
||
os.remove(self.output_filename)
|
||
logging.getLogger('VADApp').info(f"Старый файл записей удалён: {self.output_filename}")
|
||
|
||
# 🔹 Запуск фонового потока
|
||
self.audio = pyaudio.PyAudio()
|
||
self.running = True
|
||
self.vad_thread = threading.Thread(target=self._vad_loop, daemon=True)
|
||
self.vad_thread.start()
|
||
|
||
def _calculate_energy(self, frame):
|
||
"""Вычисляет RMS энергию аудио фрейма с шумоподавлением и автоусилением.
|
||
Возвращает кортеж: (обработанный_фрейм, исходная_энергия)
|
||
"""
|
||
logger = logging.getLogger('VADApp')
|
||
try:
|
||
# Преобразуем байты в массив int16
|
||
audio_data = struct.unpack(f'{len(frame)//2}h', frame)
|
||
audio_array = np.array(audio_data, dtype=np.float32)
|
||
|
||
# Вычисляем энергию ИСХОДНОГО сигнала (для VAD)
|
||
base_energy = np.sqrt(np.mean(audio_array**2))
|
||
|
||
# Обновляем оценку уровня шума
|
||
if self.noise_reduction_enabled and base_energy < self.current_energy_threshold * 0.5:
|
||
current_noise = np.mean(np.abs(audio_array))
|
||
if self.noise_floor == 0.0:
|
||
# Первая инициализация
|
||
self.noise_floor = current_noise
|
||
else:
|
||
# Плавное обновление с настраиваемой скоростью
|
||
alpha = float(self.config.get('noise_update_alpha', 0.05))
|
||
self.noise_floor = (1 - alpha) * self.noise_floor + alpha * current_noise
|
||
|
||
# ========== ШУМОПОДАВЛЕНИЕ ==========
|
||
if self.noise_reduction_enabled and self.noise_floor > 0:
|
||
reduction_factor = float(self.config.get('noise_reduction_factor', 1.5))
|
||
noise_threshold = self.noise_floor * reduction_factor
|
||
audio_array = np.where(
|
||
np.abs(audio_array) > noise_threshold,
|
||
np.sign(audio_array) * (np.abs(audio_array) - self.noise_floor),
|
||
0
|
||
)
|
||
|
||
# ========== АВТОУСИЛЕНИЕ ==========
|
||
if self.auto_gain_enabled:
|
||
current_rms = np.sqrt(np.mean(audio_array**2))
|
||
|
||
# Вычисляем минимальный порог RMS относительно уровня шума
|
||
min_rms_factor = float(self.config.get('min_rms_threshold_factor', 2.0))
|
||
min_rms_threshold = self.noise_floor * min_rms_factor if self.noise_floor > 0 else 100.0
|
||
|
||
if current_rms > min_rms_threshold:
|
||
# Вычисляем целевой коэффициент усиления
|
||
target_gain = self.target_rms / current_rms
|
||
|
||
# Ограничиваем максимальное усиление
|
||
max_gain = float(self.config.get('max_gain', 5.0))
|
||
target_gain = min(max_gain, max(0.1, target_gain))
|
||
|
||
# Плавное изменение усиления с настраиваемой скоростью
|
||
alpha = float(self.config.get('gain_smoothing_alpha', 0.1))
|
||
self.current_gain = (1 - alpha) * self.current_gain + alpha * target_gain
|
||
|
||
# Применяем усиление
|
||
audio_array = audio_array * self.current_gain
|
||
|
||
# Преобразуем обратно в байты
|
||
audio_array = np.clip(audio_array, -32768, 32767).astype(np.int16)
|
||
processed_frame = struct.pack(f'{len(audio_array)}h', *audio_array)
|
||
|
||
# Возвращаем обработанный фрейм И исходную энергию
|
||
return processed_frame, base_energy
|
||
|
||
except Exception as e:
|
||
logger.error(f"_calculate_energy: {e}")
|
||
return frame, 0.0 # В случае ошибки возвращаем исходный фрейм
|
||
|
||
def _update_adaptive_threshold(self, energy, is_speech_candidate):
|
||
"""Обновляет адаптивный порог энергии на основе фонового шума"""
|
||
if not self.adaptive_energy:
|
||
return
|
||
|
||
# Если не похоже на речь, добавляем в буфер фоновой энергии
|
||
if not is_speech_candidate:
|
||
self.background_energy_buffer.append(energy)
|
||
|
||
# Обновляем порог как среднее + отклонение фонового шума
|
||
if len(self.background_energy_buffer) >= 10:
|
||
bg_mean = np.mean(self.background_energy_buffer)
|
||
bg_std = np.std(self.background_energy_buffer)
|
||
# Порог = среднее + adaptive_energy_std_factor стандартных отклонения
|
||
self.current_energy_threshold = max(bg_mean + self.adaptive_energy_std_factor * bg_std, self.energy_threshold)
|
||
|
||
def _vad_loop(self):
|
||
logger = logging.getLogger('VADApp')
|
||
stream = None
|
||
try:
|
||
stream = self.audio.open(
|
||
format=pyaudio.paInt16,
|
||
channels=1,
|
||
rate=self.config['sample_rate'],
|
||
input=True,
|
||
frames_per_buffer=self.frame_size
|
||
)
|
||
|
||
logger.info("Инициализация VAD завершена. Начинаем обработку...")
|
||
|
||
while self.running:
|
||
try:
|
||
frame = stream.read(self.frame_size, exception_on_overflow=False)
|
||
|
||
# Базовая VAD проверка
|
||
is_speech_vad = self.vad.is_speech(frame, self.config['sample_rate'])
|
||
|
||
# Вычисляем энергию сигнала
|
||
processed_frame, energy = self._calculate_energy(frame)
|
||
|
||
# Обновляем адаптивный порог
|
||
self._update_adaptive_threshold(energy, is_speech_vad)
|
||
|
||
with self.lock:
|
||
# Комбинированная проверка речи (VAD + энергия)
|
||
is_speech = is_speech_vad and (energy > self.current_energy_threshold)
|
||
|
||
# Обновляем кольцевой буфер и историю
|
||
self.ring_buffer.append(1 if is_speech else 0)
|
||
speech_ratio = sum(self.ring_buffer) / len(self.ring_buffer)
|
||
speech_confirmed = speech_ratio > self.config['speech_threshold']
|
||
self.history.append(speech_ratio)
|
||
|
||
# Обновляем счетчик連続ных фреймов речи
|
||
if speech_confirmed:
|
||
self.continuous_speech_frames += 1
|
||
else:
|
||
self.continuous_speech_frames = 0
|
||
|
||
# Финальное подтверждение речи (фильтр по длительности)
|
||
speech_detected = self.continuous_speech_frames >= self.min_speech_frames
|
||
|
||
# Обновляем last_result
|
||
self.last_result = {
|
||
'speech_ratio': speech_ratio,
|
||
'speech_detected': speech_detected,
|
||
'energy': energy,
|
||
'energy_threshold': self.current_energy_threshold,
|
||
'speech_threshold': self.config['speech_threshold'],
|
||
'continuous_frames': self.continuous_speech_frames,
|
||
}
|
||
|
||
# Логика начала/окончания записи
|
||
if speech_detected:
|
||
if not self.is_recording:
|
||
self._start_recording()
|
||
self._cancel_stop_recording() # Отменяем таймер, если речь продолжается
|
||
else:
|
||
if self.is_recording and not self.stop_timer:
|
||
self._schedule_stop_recording()
|
||
|
||
# Добавляем фрейм в буфер
|
||
if not self.is_recording:
|
||
self.prebuffer.append(processed_frame)
|
||
else:
|
||
self.recording_buffer.append(processed_frame)
|
||
# Проверка на максимальную длину записи
|
||
if len(self.recording_buffer) >= self.max_record_frames:
|
||
logger.info(f"Максимальная длина записи достигнута ({self.max_record_duration_sec} сек)")
|
||
self._stop_recording()
|
||
|
||
# Обновляем состояния
|
||
self.prev_speech_detected = speech_confirmed
|
||
self.last_speech_detected = speech_confirmed
|
||
|
||
# Вывод статуса
|
||
recording_status = "🔴 ЗАПИСЬ" if self.is_recording else "⚪ ПРЕДБУФЕР"
|
||
timer_status = " [⏱️]" if self.stop_timer else ""
|
||
prebuf_len = len(self.prebuffer) if not self.is_recording else 0
|
||
rec_len = len(self.recording_buffer) if self.is_recording else 0
|
||
|
||
status_line = ( f"Статус: {status_tracker.get_status()} |"
|
||
f"VAD: {recording_status}{timer_status} |"
|
||
f"Речь: {speech_ratio:.2f} | Энергия: {energy:.0f}/{self.current_energy_threshold:.0f} | "
|
||
f"Gain: {self.current_gain:.2f}x | Шум: {self.noise_floor:.0f} | "
|
||
f"Кадры: {self.continuous_speech_frames:>3} | Буферы: П{prebuf_len:>3}/З{rec_len:>4}")
|
||
if logger.isEnabledFor(logging.DEBUG):
|
||
sys.stdout.write(f"\r{status_line}")
|
||
sys.stdout.flush()
|
||
else:
|
||
logger.debug(status_line.replace('\n', ' '))
|
||
|
||
except Exception as e:
|
||
if self.running:
|
||
logger.error(f"VAD Read Error: {e}")
|
||
break
|
||
|
||
except Exception as e:
|
||
logger.error(f"VAD Stream Error: {e}")
|
||
finally:
|
||
if stream:
|
||
try:
|
||
stream.stop_stream()
|
||
stream.close()
|
||
except:
|
||
pass
|
||
|
||
def _start_recording(self):
|
||
"""Начинаем запись: копируем предбуфер в основной буфер и переключаемся в режим записи"""
|
||
logger = logging.getLogger('VADApp')
|
||
self.is_recording = True
|
||
self.recording_buffer = list(self.prebuffer)
|
||
self.prebuffer.clear()
|
||
logger.info(f"Запись начата (предбуфер: {len(self.recording_buffer)} фреймов)")
|
||
|
||
status_tracker.set_status(SystemStatus.RECORDING)
|
||
|
||
def _schedule_stop_recording(self):
|
||
"""Запускает таймер на остановку записи через N секунд (только если таймер еще не запущен)"""
|
||
logger = logging.getLogger('VADApp')
|
||
if self.stop_timer is not None:
|
||
return
|
||
|
||
def delayed_stop():
|
||
with self.lock:
|
||
if self.is_recording:
|
||
self._stop_recording()
|
||
|
||
self.stop_timer = threading.Timer(self.record_stop_delay_sec, delayed_stop)
|
||
self.stop_timer.start()
|
||
logger.info(f"Таймер остановки запущен ({self.record_stop_delay_sec} сек)")
|
||
|
||
def _cancel_stop_recording(self):
|
||
"""Отменяет таймер остановки записи"""
|
||
logger = logging.getLogger('VADApp')
|
||
if self.stop_timer:
|
||
self.stop_timer.cancel()
|
||
self.stop_timer = None
|
||
logger.info("Таймер остановки отменён (речь возобновилась)")
|
||
|
||
def _stop_recording(self):
|
||
"""Останавливаем запись: обрабатываем буфер для распознавания и/или сохранения в файл"""
|
||
logger = logging.getLogger('VADApp')
|
||
if not self.is_recording:
|
||
return
|
||
|
||
if self.on_speech_end_callback and self.recording_buffer:
|
||
try:
|
||
audio_data = b''.join(self.recording_buffer)
|
||
self.on_speech_end_callback(audio_data, self.config['sample_rate'])
|
||
except Exception as e:
|
||
logger.error(f"Ошибка вызова callback: {e}")
|
||
|
||
if self.recording_buffer and self.record_enabled:
|
||
try:
|
||
with wave.open(self.output_filename, 'wb') as wf:
|
||
wf.setnchannels(1)
|
||
wf.setsampwidth(2)
|
||
wf.setframerate(self.config['sample_rate'])
|
||
wf.writeframes(b''.join(self.recording_buffer))
|
||
|
||
duration = len(self.recording_buffer) * self.config['frame_duration_ms'] / 1000
|
||
file_size = os.path.getsize(self.output_filename) / 1024
|
||
logger.info(f"Запись сохранена: {self.output_filename}")
|
||
logger.info(f"Длительность: {duration:.2f} сек, Размер: {file_size:.1f} KB, Фреймов: {len(self.recording_buffer)}")
|
||
except Exception as e:
|
||
logger.error(f"Ошибка сохранения записи: {e}")
|
||
|
||
self.recording_buffer.clear()
|
||
status_tracker.set_status(SystemStatus.LISTENING)
|
||
self.is_recording = False
|
||
self.stop_timer = None
|
||
logger.info("Возврат к режиму предбуфера")
|
||
|
||
def get_latest_status(self):
|
||
"""Возвращает последнее состояние VAD (для polling)"""
|
||
with self.lock:
|
||
return self.last_result.copy()
|
||
|
||
def update_config(self, new_config):
|
||
logger = logging.getLogger('VADApp')
|
||
with self.lock:
|
||
self.config = new_config
|
||
self.vad.set_mode(new_config['mode'])
|
||
self.ring_buffer = collections.deque(maxlen=new_config['ring_buffer_size'])
|
||
self.history = collections.deque(maxlen=new_config.get('history_maxlen', 100))
|
||
self.frame_size = int(new_config['sample_rate'] * new_config['frame_duration_ms'] / 1000)
|
||
|
||
# Обновляем параметры фильтрации
|
||
self.energy_threshold = new_config.get('energy_threshold', 1000)
|
||
self.min_speech_duration_ms = new_config.get('min_speech_duration_ms', 300)
|
||
self.adaptive_energy = new_config.get('adaptive_energy', True)
|
||
self.adaptive_energy_std_factor = new_config.get('adaptive_energy_std_factor', 2.5)
|
||
self.min_speech_frames = int(self.min_speech_duration_ms / new_config['frame_duration_ms'])
|
||
|
||
# Обновляем параметры шумоподавления и автоусиления
|
||
self.noise_reduction_enabled = new_config.get('noise_reduction_enabled', True)
|
||
self.auto_gain_enabled = new_config.get('auto_gain_enabled', True)
|
||
self.target_rms = new_config.get('target_rms', 2000.0)
|
||
|
||
# Обновляем размер предбуфера и пересоздаем его
|
||
old_prebuffer_data = list(self.prebuffer) if not self.is_recording else []
|
||
self.prebuffer_size = int(new_config.get('prebuffer_duration_sec', 1.0) * new_config['sample_rate'] / self.frame_size)
|
||
self.prebuffer = collections.deque(old_prebuffer_data, maxlen=self.prebuffer_size)
|
||
|
||
# Обновляем максимальную длину записи
|
||
self.max_record_duration_sec = new_config.get('max_record_duration_sec', 5.0)
|
||
self.max_record_frames = int(self.max_record_duration_sec * new_config['sample_rate'] / self.frame_size)
|
||
|
||
# Обновляем настройки записи
|
||
self.record_enabled = new_config.get('record_on_speech', False)
|
||
self.record_stop_delay_sec = new_config.get('record_stop_delay_sec', 1.0)
|
||
|
||
logger.info("Конфигурация VAD обновлена")
|
||
|
||
def close(self):
|
||
"""Корректное завершение работы с очисткой всех ресурсов"""
|
||
logger = logging.getLogger('VADApp')
|
||
logger.info("Начинаем корректное завершение VAD...")
|
||
self.running = False
|
||
|
||
# Даем время потоку завершиться
|
||
if hasattr(self, 'vad_thread') and self.vad_thread.is_alive():
|
||
self.vad_thread.join(timeout=3.0)
|
||
|
||
with self.lock:
|
||
# Отменяем таймер, если активен
|
||
if self.stop_timer:
|
||
self.stop_timer.cancel()
|
||
self.stop_timer = None
|
||
|
||
# Сохраняем текущую запись, если активна
|
||
if self.is_recording:
|
||
logger.info("Сохраняем активную запись...")
|
||
self._stop_recording()
|
||
|
||
# Очищаем все буферы для освобождения памяти
|
||
self.prebuffer.clear()
|
||
self.recording_buffer.clear()
|
||
self.ring_buffer.clear()
|
||
self.history.clear()
|
||
self.background_energy_buffer.clear()
|
||
|
||
# Завершаем audio
|
||
if hasattr(self, 'audio'):
|
||
try:
|
||
self.audio.terminate()
|
||
except:
|
||
pass # Игнорируем ошибки при завершении
|
||
|
||
logger.info("Корректное завершение VAD завершено")
|