input.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import time
import sounddevice as sd
import soundfile as sf
import cv2
import multiprocessing as mp
from picamera2 import Picamera2
import numpy as np
import webrtcvad
from faster_whisper import WhisperModel
import warnings
import subprocess
warnings.filterwarnings("ignore", category=UserWarning, module="webrtcvad")

# ──────────────────────────────────────────────────────────────────────────────
# [1] 설정 값
# ──────────────────────────────────────────────────────────────────────────────
FRAME_DURATION_MS = 30        # 블록 길이 (ms)
INITIAL_PERIOD = 1.0          # 녹음 시작 후 무조건 기록하는 시간 (초)
GRACE_PERIOD = 2.0            # 이후 연속 무음이 이 시간(초) 이상 지속되면 녹음 종료
RMS_THRESHOLD = 500           # RMS 임계치
VAD_MODE = 1                  # WebRTC VAD 모드 (0~3)

BASE_DIR = "/home/tdd_jimin/tdd"
TEMP_AUDIO_PATH = os.path.join(BASE_DIR, "audio.wav")
TEMP_TXT_PATH   = os.path.join(BASE_DIR, "texts.txt")
TEMP_VIDEO_PATH = os.path.join(BASE_DIR, "video.mp4")
FINAL_VIDEO_PATH= os.path.join(BASE_DIR, "final_output.mp4")

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

SD_INPUT_DEVICE_INDEX = 0     # 가능하면 자동 대체함
USE_VIDEO = False             # 비디오를 켜려면 True
# ──────────────────────────────────────────────────────────────────────────────

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

def pick_supported_vad_rate(device_sr: float) -> int:
    """
    WebRTC-VAD 지원 샘플레이트(8000, 16000, 32000, 48000) 중
    device_sr에 가장 가까운 값을 반환.
    """
    supported = [8000, 16000, 32000, 48000]
    best = min(supported, key=lambda x: abs(x - device_sr))
    return best

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:
            info = sd.query_devices(idx_hint, kind='input')
            return idx_hint
        except Exception:
            pass  # fallthrough to default

    # 기본 입력 장치가 지정되어 있으면 사용
    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: list[int], device_index: int, frame_size: int):
    """
    주어진 샘플레이트 후보들에 대해 순차적으로 InputStream 열기 시도.
    성공하면 (stream, use_sr) 반환.
    """
    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 audio_process(stop_event: mp.Event):
    """
    1) 입력 디바이스의 기본 샘플레이트를 읽어서 VAD가 지원하는 rate로 맞춤.
    2) 해당 샘플레이트(use_sr)를 사용해 프레임 크기를 계산하고 오디오 녹음.
    3) 초기 INITIAL_PERIOD 동안은 무조건 기록 → 그 이후 RMS+VAD 검사.
    4) (VAD=False AND RMS≤RMS_THRESHOLD) 연속 GRACE_PERIOD 이상일 때 녹음 종료.
    """
    try:
        # 디렉터리 보장
        ensure_paths()

        # ─ 1) 디바이스/샘플레이트 선택 ─
        device_index = choose_input_device(SD_INPUT_DEVICE_INDEX)
        device_info = sd.query_devices(device_index, kind='input')
        dev_default_sr = float(device_info['default_samplerate'])
        vad_target_sr = pick_supported_vad_rate(dev_default_sr)

        # 샘플레이트 후보: VAD 권장값 우선 → 그 외 지원 목록
        sr_candidates = [vad_target_sr, 16000, 48000, 32000, 8000]

        FRAME_SIZE_CANDIDATE = int(vad_target_sr * FRAME_DURATION_MS / 1000)

        print(f"[AudioProc] 입력장치 #{device_index} '{device_info['name']}' "
              f"기본SR={dev_default_sr:.1f}Hz → VAD 대상SR={vad_target_sr}Hz")

        # VAD 객체
        vad = webrtcvad.Vad(VAD_MODE)

        start_ts = time.time()
        silence_start = None

        # WAV 파일 오픈
        with sf.SoundFile(
            TEMP_AUDIO_PATH, mode='w',
            samplerate=vad_target_sr, channels=1,
            subtype='PCM_16'
        ) as wf:
            # InputStream을 후보 SR로 열기
            stream, use_sr = open_input_stream_with_fallback(sr_candidates, device_index, FRAME_SIZE_CANDIDATE)
            print(f"[AudioProc] 실제 입력 SR={use_sr}Hz")

            # 프레임 크기 재계산(실제 SR 기준)
            FRAME_SIZE = int(use_sr * FRAME_DURATION_MS / 1000)
            print(f"[AudioProc] 프레임 크기={FRAME_SIZE}샘플, INITIAL_PERIOD={INITIAL_PERIOD}s")

            try:
                while True:
                    data, overflowed = stream.read(FRAME_SIZE)  # shape=(FRAME_SIZE, 1), dtype=int16
                    if overflowed:
                        print("[AudioProc] 경고: 오디오 버퍼 오버플로가 발생했습니다.")

                    # WAV 기록 (WF의 SR은 vad_target_sr이라 다를 수 있음 -> 리샘플 없이 기록)
                    mono_data = data.flatten()
                    wf.write(mono_data)

                    now = time.time()
                    elapsed = now - start_ts

                    # 초기 구간: 무조건 기록
                    if elapsed < INITIAL_PERIOD:
                        print(f"[{elapsed:.2f}s] 초기 기록 중…")
                        continue

                    # RMS & VAD 검사
                    rms = np.sqrt(np.mean(np.square(mono_data.astype(np.float32))))
                    pcm_bytes = mono_data.tobytes()
                    # VAD는 현재 입력 SR 기준으로 동작
                    is_speech = vad.is_speech(pcm_bytes, use_sr)

                    vad_state = "음성" if is_speech else "무음"
                    rms_state = "음성" if rms > RMS_THRESHOLD else "무음"
                    print(f"[{elapsed:.2f}s] VAD={vad_state}, RMS={rms:.1f}({rms_state})")

                    # 의도대로: (VAD=False AND RMS<=THRESHOLD) 이 연속 GRACE_PERIOD
                    if (not is_speech) and (rms <= RMS_THRESHOLD):
                        if silence_start is None:
                            silence_start = now
                            print(f"    → 무음 시작 시점 기록: {elapsed:.2f}s")
                        else:
                            if now - silence_start >= GRACE_PERIOD:
                                print(f"[AudioProc] 무음 연속 {GRACE_PERIOD}s → 녹음 종료")
                                break
                    else:
                        silence_start = None

                # 루프 종료 → 파일 닫히며 flush
            finally:
                try:
                    stream.stop()
                    stream.close()
                except Exception:
                    pass

        stop_event.set()
        print("[AudioProc] 종료 신호 set.")
    except Exception as e:
        print(f"[AudioProc] 치명적 오류: {e}")
        stop_event.set()

def video_process(stop_event: mp.Event):
    """
    Picamera2로 카메라 영상을 녹화하며, stop_event가 set되면 중단.
    """
    try:
        print("[VideoProc] Picamera2 초기화 (자식 프로세스).")
        picam2 = Picamera2()
        config = picam2.create_video_configuration(
            main={"size": VIDEO_SIZE},
            controls={"FrameRate": VIDEO_FPS}
        )
        picam2.configure(config)
        picam2.start()
        print(f"[VideoProc] 비디오 녹화 시작 (FPS={VIDEO_FPS}, Size={VIDEO_SIZE})…")

        fourcc = cv2.VideoWriter_fourcc(*VIDEO_CODEC)
        out = cv2.VideoWriter(
            TEMP_VIDEO_PATH,
            fourcc,
            VIDEO_FPS,
            VIDEO_SIZE
        )

        while not stop_event.is_set():
            frame = picam2.capture_array()  # RGB
            frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
            out.write(frame)

        picam2.stop()
        out.release()
        print("[VideoProc] 비디오 녹화 종료. 저장:", TEMP_VIDEO_PATH)
    except Exception as e:
        print(f"[VideoProc] 오류: {e}")

def stt_transcribe(path: str) -> str:
    """
    Whisper를 사용해 audio.wav 파일을 텍스트로 변환.
    """
    print("[STT] WhisperModel 로딩 및 음성인식 시작...")
    model = WhisperModel("tiny", device="cpu", compute_type="int8")
    segments, _ = model.transcribe(path, language="ko", task="transcribe")
    text = "".join(seg.text for seg in segments)
    print("[STT] 인식 결과:", text)
    return text

def convert_to_mp4(video_path: str, audio_path: str, out_path: str):
    """
    ffmpeg로 비디오와 오디오를 합쳐 최종 MP4 생성.
    """
    print("[Mux] 오디오+비디오 합치는 중...")
    cmd = [
        "ffmpeg", "-y",
        "-i", video_path,
        "-i", audio_path,
        "-c:v", "copy",
        "-c:a", "aac",
        "-strict", "experimental",
        out_path
    ]
    try:
        subprocess.run(cmd, check=True)
    except subprocess.CalledProcessError as e:
        print(f"[Mux] ffmpeg 실패: {e}")
    print("[Mux] 최종 영상 저장:", out_path)

def wait_for_stable_file(path: str, checks: int = 3, interval: float = 0.2) -> bool:
    """
    파일 크기가 연속 checks회 동일하면 안정화된 것으로 간주.
    """
    if not os.path.exists(path):
        return False
    prev = -1
    stable = 0
    for _ in range(checks * 3):
        try:
            sz = os.path.getsize(path)
        except FileNotFoundError:
            return False
        if sz == prev and sz > 0:
            stable += 1
            if stable >= checks:
                return True
        else:
            stable = 0
        prev = sz
        time.sleep(interval)
    return False

def main():
    # Picamera2 사용 시 spawn 권장
    mp.set_start_method('spawn', force=True)

    print("=== 시스템 시작 ===")
    ensure_paths()
    stop_event = mp.Event()

    # 프로세스 시작
    p_audio = mp.Process(target=audio_process, args=(stop_event,))
    p_audio.start()

    # 비디오 유지 (삭제 금지). 필요 시 주석 해제하여 사용.
    # p_video = mp.Process(target=video_process, args=(stop_event,))
    # if USE_VIDEO:
    #     p_video.start()

    try:
        # 오디오 종료를 **확실히** 기다린 뒤 STT 실행
        p_audio.join()

        # 비디오 사용 시에는 비디오도 종료 대기
        # if USE_VIDEO:
        #     stop_event.set()
        #     p_video.join()

        # STT 실행 전 파일 안정화 확인
        if wait_for_stable_file(TEMP_AUDIO_PATH):
            text = stt_transcribe(TEMP_AUDIO_PATH)
        else:
            raise RuntimeError("오디오 파일이 생성/안정화되지 않았습니다.")

        print(text)
        with open(TEMP_TXT_PATH, "w", encoding="utf-8") as file:
            file.write(text)

        # 비디오를 녹화했다면 병합
        # if USE_VIDEO:
        #     convert_to_mp4(TEMP_VIDEO_PATH, TEMP_AUDIO_PATH, FINAL_VIDEO_PATH)

    except KeyboardInterrupt:
        print("[Main] 사용자 중단(Ctrl+C) 감지. 종료 처리 중…")
        stop_event.set()
    finally:
        # 안전 종료
        if p_audio.is_alive():
            p_audio.terminate()
            p_audio.join()
        # if USE_VIDEO and p_video.is_alive():
        #     p_video.terminate()
        #     p_video.join()

    print("=== 시스템 종료 ===")

if __name__ == "__main__":
    main()