recorder.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""영상 쓰려면 recorder.py 상단의 USE_VIDEO = False를 True로 바꾸거나
stt_daemon.py에서 record_session(..., use_video=True)로 호출
또는 프런트에서 RECORD_AV 명령을 보내도록 하면 됨
"""
import os, sys, time, warnings
import numpy as np
import sounddevice as sd
import soundfile as sf
import webrtcvad

warnings.filterwarnings("ignore", category=UserWarning, module="webrtcvad")

# ===== 설정 =====
FRAME_DURATION_MS = 30
INITIAL_PERIOD = 1.0
GRACE_PERIOD = 2.0
RMS_THRESHOLD = 600
VAD_MODE = 1

BASE_DIR = "/home/tdd_jimin/tdd"
AUDIO_PATH = os.path.join(BASE_DIR, "audio.wav")
VIDEO_PATH = os.path.join(BASE_DIR, "video.mp4")

# 기본: 오디오만. ↓ 주석 해제하면 영상도 함께 녹화
# USE_VIDEO = True
USE_VIDEO = False

VIDEO_FPS = 20
VIDEO_SIZE = (640, 480)
VIDEO_CODEC = "mp4v"  # OpenCV fourcc

def log(msg):  # stderr 로깅
    print(msg, file=sys.stderr, flush=True)

def ensure_paths():
    os.makedirs(BASE_DIR, exist_ok=True)

def pick_supported_vad_rate(device_sr: float) -> int:
    supported = [8000, 16000, 32000, 48000]
    return min(supported, key=lambda x: abs(x - device_sr))

def choose_input_device(idx_hint: int | None):
    devices = sd.query_devices()
    input_indices = [i for i, d in enumerate(devices) if d.get("max_input_channels", 0) > 0]
    if not input_indices:
        raise RuntimeError("입력 가능한 오디오 장치를 찾지 못했습니다.")
    if idx_hint is not None:
        try:
            sd.query_devices(idx_hint, kind='input')
            return idx_hint
        except Exception:
            pass
    default_in = sd.default.device[0] if isinstance(sd.default.device, (list, tuple)) else sd.default.device
    if isinstance(default_in, int) and default_in in input_indices:
        return default_in
    return input_indices[0]

def open_input_stream_with_fallback(samplerates, device_index, frame_size):
    last_err = None
    for sr in samplerates:
        try:
            stream = sd.InputStream(
                samplerate=sr, channels=1, dtype='int16',
                blocksize=frame_size if frame_size else int(sr * FRAME_DURATION_MS / 1000),
                device=device_index
            )
            stream.start()
            return stream, sr
        except Exception as e:
            last_err = e
    raise RuntimeError(f"오디오 입력 스트림을 열 수 없습니다: {last_err}")

def record_session(
    out_audio=AUDIO_PATH,
    out_video=VIDEO_PATH,
    sd_input_device_index=0,
    use_video=USE_VIDEO
):
    """
    무음 종료 규칙으로 하나의 세션을 녹음
    오디오 파일 경로와 (영상 사용 시) 비디오 파일 경로를 반환
    
    """
    ensure_paths()

    # 오디오 장치/VAD 준비
    device_index = choose_input_device(sd_input_device_index)
    di = sd.query_devices(device_index, kind='input')
    dev_sr = float(di['default_samplerate'])
    vad_sr = pick_supported_vad_rate(dev_sr)
    vad = webrtcvad.Vad(VAD_MODE)

    sr_candidates = [vad_sr, 16000, 48000, 32000, 8000]
    frame_size_candidate = int(vad_sr * FRAME_DURATION_MS / 1000)

    # (선택) 비디오 준비
    cap = None
    writer = None
    if use_video:
        try:
            import cv2
            from picamera2 import Picamera2
            log("[Video] 초기화 중...")
            picam2 = Picamera2()
            config = picam2.create_video_configuration(
                main={"size": VIDEO_SIZE},
                controls={"FrameRate": VIDEO_FPS}
            )
            picam2.configure(config)
            picam2.start()

            fourcc = cv2.VideoWriter_fourcc(*VIDEO_CODEC)
            writer = cv2.VideoWriter(out_video, fourcc, VIDEO_FPS, VIDEO_SIZE)

            def capture_write():
                frame = picam2.capture_array()  # RGB
                frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
                writer.write(frame)

            video_cleanup = lambda: (picam2.stop(), writer.release())
        except Exception as e:
            log(f"[Video] 초기화 실패(오디오만 진행): {e}")
            use_video = False
            capture_write = None
            video_cleanup = lambda: None
    else:
        capture_write = None
        video_cleanup = lambda: None

    # 오디오 스트림 및 WAV 파일 오픈
    start_ts = time.time()
    silence_start = None
    with sf.SoundFile(out_audio, mode='w', samplerate=vad_sr, channels=1, subtype='PCM_16') as wf:
        stream, use_sr = open_input_stream_with_fallback(sr_candidates, device_index, frame_size_candidate)
        FRAME_SIZE = int(use_sr * FRAME_DURATION_MS / 1000)
        log("[Audio] 녹음 시작")
        try:
            while True:
                data, overflowed = stream.read(FRAME_SIZE)
                if overflowed:
                    log("[Audio] 오버플로 발생")
                mono = data.flatten()
                wf.write(mono)

                # 비디오 프레임 기록(선택)
                if use_video and capture_write:
                    capture_write()

                now = time.time()
                elapsed = now - start_ts
                if elapsed < INITIAL_PERIOD:
                    continue

                rms = float(np.sqrt(np.mean(np.square(mono.astype(np.float32)))))
                is_speech = vad.is_speech(mono.tobytes(), use_sr)
                log(f"[Audio] rms: {rms}")
                if (not is_speech) and (rms <= RMS_THRESHOLD):
                    if silence_start is None:
                        silence_start = now
                    elif now - silence_start >= GRACE_PERIOD:
                        log("[Audio] 무음 지속 → 종료")
                        break
                else:
                    silence_start = None
        finally:
            try:
                stream.stop(); stream.close()
            except Exception:
                pass
            video_cleanup()

    return out_audio, (out_video if use_video else None)

if __name__ == "__main__":
    # 단독 실행 테스트: 기본(오디오만) 녹음
    a, v = record_session()
    print(json.dumps({"audio": a, "video": v}), flush=True)