momentry_core/scripts/asr_processor.py

#!/opt/homebrew/bin/python3.11
"""
ASR Processor - faster-whisper small model (Production)

Version: 2.1
Model: small (int8 quantization, CPU)
Reason: small 模型在準確率和速度間取得最佳平衡
        經實驗驗證，最少要使用 small 才可以較好的處理多語種及台灣腔國語

Configuration:
- Model: faster-whisper/small
- Device: CPU (MPS not supported by faster_whisper)
- Compute: int8
- Beam size: 5
- VAD filter: enabled (min_silence=500ms, speech_pad=200ms)
- Audio fallback: ffmpeg extraction for PyAV-incompatible streams (v2.1)
"""
import sys
import json
import os
import time
import argparse
import signal
import subprocess
import tempfile
from datetime import datetime
from faster_whisper import WhisperModel

PROCESSOR_VERSION = "2.1"
MODEL_SIZE = "small"
DEVICE = "cpu"
COMPUTE_TYPE = "int8"

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from redis_publisher import RedisPublisher


def signal_handler(signum, frame):
    print(f"ASR: Received signal {signum}, exiting...")
    sys.exit(1)


def has_audio_stream(video_path):
    """Check if video file has audio stream using ffprobe."""
    try:
        cmd = [
            "ffprobe",
            "-v",
            "error",
            "-select_streams",
            "a",
            "-show_entries",
            "stream=codec_type",
            "-of",
            "csv=p=0",
            video_path,
        ]
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        return bool(result.stdout.strip())
    except subprocess.CalledProcessError:
        return False
    except FileNotFoundError:
        print("WARNING: ffprobe not found, assuming audio exists")
        return True


def extract_audio_with_ffmpeg(video_path):
    """Extract audio from video to WAV using ffmpeg.
    
    Returns path to temporary WAV file. Caller is responsible for cleanup.
    """
    wav_path = tempfile.mktemp(suffix=".wav", prefix="asr_audio_")
    cmd = [
        "ffmpeg",
        "-y",
        "-i", video_path,
        "-vn",
        "-acodec", "pcm_s16le",
        "-ar", "16000",
        "-ac", "1",
        wav_path,
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        sys.stderr.write(f"ASR: ffmpeg extraction failed: {result.stderr}\n")
        sys.stderr.flush()
        return None
    return wav_path


def transcribe_with_fallback(model, video_path, publisher=None):
    """Transcribe video with fallback to ffmpeg-extracted WAV.
    
    First tries direct transcription (PyAV). If PyAV fails to decode,
    falls back to ffmpeg audio extraction then transcription.
    """
    # Try direct transcription first
    try:
        if publisher:
            publisher.info("asr", "Direct transcription attempt...")
        return model.transcribe(
            video_path,
            beam_size=5,
            vad_filter=True,
            vad_parameters=dict(min_silence_duration_ms=500, speech_pad_ms=200),
        )
    except Exception as e:
        error_str = str(e)
        # Check if it's a PyAV/av decoding error
        is_pyav_error = any(
            keyword in error_str.lower()
            for keyword in ["av.error", "avcodec", "decode", "packet"]
        )
        
        if not is_pyav_error:
            raise  # Re-raise non-PyAV errors
        
        if publisher:
            publisher.info("asr", "PyAV decode failed, falling back to ffmpeg extraction...")
        sys.stderr.write("ASR: PyAV decode error detected, falling back to ffmpeg extraction\n")
        sys.stderr.flush()
        
        wav_path = extract_audio_with_ffmpeg(video_path)
        if wav_path is None:
            raise RuntimeError("Failed to extract audio with ffmpeg")
        
        try:
            if publisher:
                publisher.info("asr", "Transcribing extracted WAV audio...")
            segments, info = model.transcribe(
                wav_path,
                beam_size=5,
                vad_filter=True,
                vad_parameters=dict(min_silence_duration_ms=500, speech_pad_ms=200),
            )
            return segments, info
        finally:
            # Clean up temporary WAV file
            try:
                os.remove(wav_path)
            except OSError:
                pass


def run_asr(video_path, output_path, uuid: str = ""):
    # Set up signal handlers
    signal.signal(signal.SIGTERM, signal_handler)
    signal.signal(signal.SIGINT, signal_handler)

    publisher = RedisPublisher(uuid) if uuid else None
    if publisher:
        publisher.info("asr", "ASR_START")

    # Check for audio stream
    if not has_audio_stream(video_path):
        if publisher:
            publisher.info("asr", "No audio stream detected, skipping transcription")
        output = {"language": "", "language_probability": 0.0, "segments": []}
        with open(output_path, "w") as f:
            json.dump(output, f, indent=2)
        if publisher:
            publisher.complete("asr", "0 segments (no audio)")
        sys.stderr.write("ASR: No audio stream, skipping transcription\n")
        sys.stderr.flush()
        sys.exit(0)

    if publisher:
        publisher.info("asr", "Loading Whisper model...")

    # Use small model with CPU (MPS not supported by faster_whisper)
    # small 模型在準確率和速度間取得最佳平衡
    model = WhisperModel("small", device="cpu", compute_type="int8")

    if publisher:
        publisher.info("asr", f"Transcribing: {video_path}")

    # Transcribe with VAD filter for better accuracy, with PyAV fallback
    segments, info = transcribe_with_fallback(model, video_path, publisher)

    if publisher:
        publisher.info("asr", f"ASR_LANGUAGE:{info.language}")

    results = []
    total_segments = 0

    for segment in segments:
        results.append(
            {"start": segment.start, "end": segment.end, "text": segment.text.strip()}
        )
        total_segments += 1
        if total_segments % 100 == 0:
            if publisher:
                publisher.progress(
                    "asr", total_segments, 0, f"Segment {total_segments}"
                )

    output = {
        "language": info.language,
        "language_probability": info.language_probability,
        "segments": results,
    }

    with open(output_path, "w") as f:
        json.dump(output, f, indent=2)

    if publisher:
        publisher.complete("asr", f"{len(results)} segments")

    sys.stderr.write(
        f"ASR: Transcription complete, {len(results)} segments written to {output_path}\n"
    )
    sys.stderr.flush()
    sys.exit(0)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="ASR Transcription")
    parser.add_argument("video_path", help="Path to video file")
    parser.add_argument("output_path", help="Output JSON path")
    parser.add_argument("--uuid", "-u", help="UUID for Redis progress", default="")
    args = parser.parse_args()

    run_asr(args.video_path, args.output_path, args.uuid)