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 argparse
import signal
import subprocess
import tempfile
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)

    # 嘗試以 CUT 場景分段處理（降低長片記憶體使用）
    cut_scenes = []
    cut_path = output_path.replace(".asr.json", ".cut.json")
    if os.path.exists(cut_path):
        try:
            with open(cut_path) as f:
                cut_data = json.load(f)
            scenes = cut_data.get("scenes", [])
            if scenes:
                cut_scenes = [(s["start_time"], s["end_time"]) for s in scenes]
                print(f"[ASR] Loaded {len(cut_scenes)} cut scenes for segmented transcription", file=sys.stderr)
        except Exception as e:
            print(f"[ASR] Failed to load cut scenes: {e}", file=sys.stderr)

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

    model = WhisperModel(MODEL_SIZE, device="cpu", compute_type="int8")

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

    results = []
    total_segments = 0

    if cut_scenes:
        # 分段處理：對每個場景萃取音訊並轉錄
        import subprocess
        import tempfile
        import json
        temp_dir = tempfile.mkdtemp(prefix="asr_cut_")
        transcript_language = None

        # 建立 scene lookup: 給定時間點，找是哪個 scene
        import bisect
        scene_starts = [s[0] for s in cut_scenes]
        def find_scene_idx(t):
            i = bisect.bisect_right(scene_starts, t) - 1
            return max(0, i)

        # 逐段處理，每段結果即時寫入 .asr.tmp
        tmp_path = output_path + ".tmp"
        all_segments = []

        for idx, (start_t, end_t) in enumerate(cut_scenes):
            seg_wav = os.path.join(temp_dir, f"seg_{idx:04d}.wav")
            # 用 ffmpeg 萃取出該段音訊
            cmd = ["ffmpeg", "-y", "-v", "quiet", "-i", video_path,
                   "-ss", str(start_t), "-to", str(end_t),
                   "-ar", "16000", "-ac", "1", seg_wav]
            subprocess.run(cmd, check=False, capture_output=True)

            if not os.path.exists(seg_wav) or os.path.getsize(seg_wav) < 100:
                continue  # 跳過空音訊

            try:
                seg_result, seg_info = model.transcribe(
                    seg_wav, beam_size=5,
                    vad_filter=True,
                    vad_parameters=dict(min_silence_duration_ms=500, speech_pad_ms=200),
                )
                if transcript_language is None:
                    transcript_language = seg_info.language

                scene_segments = []
                for segment in seg_result:
                    seg_start = start_t + segment.start
                    seg_end = start_t + segment.end
                    scene_idx = find_scene_idx((seg_start + seg_end) / 2)
                    scene_segments.append({
                        "start": seg_start,
                        "end": seg_end,
                        "text": segment.text.strip(),
                        "scene_number": scene_idx + 1,
                    })
                    total_segments += 1

                # 當前 scene 結果寫入 .asr.tmp
                all_segments.extend(scene_segments)
                with open(tmp_path, "w") as f:
                    json.dump({"language": transcript_language or "", "segments": all_segments}, f)

                if total_segments % 100 == 0:
                    if publisher:
                        publisher.progress("asr", total_segments, 0, f"Segment {total_segments}")
            except Exception as e:
                print(f"[ASR] Segment {idx} failed: {e}", file=sys.stderr)

            # 清理暫存 WAV
            try: os.remove(seg_wav)
            except: pass

        try: os.rmdir(temp_dir)
        except: pass

        info_language = transcript_language or "unknown"
        print(f"[ASR] Segmented transcription complete: {total_segments} segments", file=sys.stderr)
    else:
        # 無 CUT 資料，直接轉錄（原有流程）
        segments, info = transcribe_with_fallback(model, video_path, publisher)
        info_language = info.language

        tmp_path = output_path + ".tmp"
        all_segments = []
        for segment in segments:
            all_segments.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}")
        with open(tmp_path, "w") as f:
            json.dump({"language": info_language, "segments": all_segments}, f)

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

    # rename .tmp → .json
    os.rename(tmp_path, output_path)

    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)