burn-subs

#!/usr/bin/env python3
"""
burn-subs: Burn English transcripts into MP4 videos using Whisper and ffmpeg.
"""

import argparse
import os
import sys
import subprocess
import tempfile
from pathlib import Path


def check_dependencies():
    """Check that required dependencies are available."""
    try:
        import whisper
        import torch
    except ImportError as e:
        print(f"Error: Missing dependency - {e}", file=sys.stderr)
        print("Please run ./install.sh first.", file=sys.stderr)
        sys.exit(1)

    # Check ffmpeg
    try:
        subprocess.run(["ffmpeg", "-version"], capture_output=True, check=True)
    except (subprocess.CalledProcessError, FileNotFoundError):
        print("Error: ffmpeg is not installed or not in PATH.", file=sys.stderr)
        print("Please run ./install.sh first.", file=sys.stderr)
        sys.exit(1)


def validate_input(video_path: str) -> Path:
    """Validate the input video file."""
    path = Path(video_path)

    if not path.exists():
        print(f"Error: Video file not found: {video_path}", file=sys.stderr)
        sys.exit(1)

    if not path.is_file():
        print(f"Error: Not a file: {video_path}", file=sys.stderr)
        sys.exit(1)

    if path.suffix.lower() != ".mp4":
        print(f"Error: Input must be an MP4 file, got: {path.suffix}", file=sys.stderr)
        sys.exit(1)

    # Verify it's a valid video with ffprobe
    try:
        result = subprocess.run(
            ["ffprobe", "-v", "error", "-select_streams", "v:0",
             "-show_entries", "stream=codec_type", "-of", "csv=p=0", str(path)],
            capture_output=True, text=True, check=True
        )
        if "video" not in result.stdout:
            raise ValueError("No video stream found")
    except (subprocess.CalledProcessError, ValueError):
        print(f"Error: Invalid or corrupted video file: {video_path}", file=sys.stderr)
        sys.exit(1)

    return path


def validate_directory(dir_path: str) -> tuple[Path, list[Path]]:
    """Validate directory and return (path, list of .mp4 files)."""
    path = Path(dir_path)
    if not path.exists():
        print(f"Error: Directory not found: {dir_path}", file=sys.stderr)
        sys.exit(1)
    if not path.is_dir():
        print(f"Error: Not a directory: {dir_path}", file=sys.stderr)
        sys.exit(1)
    mp4s = sorted(path.glob("*.mp4")) + sorted(path.glob("*.MP4"))
    mp4s = list(dict.fromkeys(p.resolve() for p in mp4s))
    valid = []
    for p in mp4s:
        if not p.is_file():
            continue
        try:
            result = subprocess.run(
                ["ffprobe", "-v", "error", "-select_streams", "v:0",
                 "-show_entries", "stream=codec_type", "-of", "csv=p=0", str(p)],
                capture_output=True, text=True, check=True
            )
            if "video" not in result.stdout:
                raise ValueError("No video stream")
        except (subprocess.CalledProcessError, FileNotFoundError, ValueError):
            print(f"Warning: Skipping invalid video: {p}", file=sys.stderr)
            continue
        valid.append(p)
    if not valid:
        print("Error: No valid MP4 videos found in directory.", file=sys.stderr)
        sys.exit(1)
    return path, valid


def get_device():
    """Detect and return the appropriate device for Whisper."""
    import torch

    if torch.cuda.is_available():
        # Check VRAM
        vram_gb = torch.cuda.get_device_properties(0).total_memory / (1024**3)
        print(f"GPU detected: {torch.cuda.get_device_name(0)} ({vram_gb:.1f}GB VRAM)")
        return "cuda"
    else:
        print("No GPU detected, using CPU (this will be slower)")
        return "cpu"


def load_whisper_model(device: str):
    """Load Whisper model (downloads on first run)."""
    import whisper
    print("Loading Whisper model (this may download on first run)...")
    return whisper.load_model("large-v3", device=device)


def transcribe_audio(video_path: Path, device: str, model=None):
    """Transcribe audio from video using Whisper."""
    if model is None:
        model = load_whisper_model(device)
    print("Transcribing audio...")
    result = model.transcribe(
        str(video_path),
        language="en",
        task="transcribe",
        verbose=False,
        # Add these:
        condition_on_previous_text=False,  # Reduces hallucinations
        compression_ratio_threshold=2.4,    # Filter suspicious segments
        logprob_threshold=-1.0,             # Filter low-confidence segments
        no_speech_threshold=0.6,            # Adjust speech detection sensitivity
    )

    if not result.get("segments"):
        print("Error: Transcription failed - no speech detected.", file=sys.stderr)
        sys.exit(1)

    return result


def format_timestamp(seconds: float) -> str:
    """Convert seconds to SRT timestamp format (HH:MM:SS,mmm)."""
    hours = int(seconds // 3600)
    minutes = int((seconds % 3600) // 60)
    secs = int(seconds % 60)
    millis = int((seconds % 1) * 1000)
    return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"


def generate_srt(transcription: dict, output_path: Path):
    """Generate SRT subtitle file from Whisper transcription."""
    segments = transcription["segments"]

    with open(output_path, "w", encoding="utf-8") as f:
        for i, segment in enumerate(segments, 1):
            start = format_timestamp(segment["start"])
            end = format_timestamp(segment["end"])
            text = segment["text"].strip()

            f.write(f"{i}\n")
            f.write(f"{start} --> {end}\n")
            f.write(f"{text}\n\n")

    print(f"Generated subtitles: {len(segments)} segments")


def burn_subtitles(video_path: Path, srt_path: Path, output_path: Path):
    """Burn subtitles into video using ffmpeg."""
    print("Burning subtitles into video...")

    # Escape special characters in path for ffmpeg filter
    srt_escaped = str(srt_path).replace("\\", "\\\\").replace(":", "\\:").replace("'", "\\'")

    # Subtitle filter with standard styling:
    # - White text with black outline
    # - Bottom center position
    # - Standard font size (24)
    subtitle_filter = (
        f"subtitles='{srt_escaped}':"
        f"force_style='FontSize=24,PrimaryColour=&HFFFFFF,"
        f"OutlineColour=&H000000,BorderStyle=1,Outline=2,"
        f"Shadow=1,Alignment=2,MarginV=30'"
    )

    cmd = [
        "ffmpeg",
        "-i", str(video_path),
        "-vf", subtitle_filter,
        "-c:v", "libx264",      # Re-encode video with H.264
        "-crf", "18",           # High quality (lower = better, 18 is visually lossless)
        "-preset", "medium",    # Balance between speed and compression
        "-c:a", "copy",         # Copy audio without re-encoding
        "-y",                   # Overwrite output
        str(output_path)
    ]

    try:
        result = subprocess.run(cmd, capture_output=True, text=True)
        if result.returncode != 0:
            print(f"Error: ffmpeg failed:\n{result.stderr}", file=sys.stderr)
            sys.exit(1)
    except Exception as e:
        print(f"Error: Failed to burn subtitles: {e}", file=sys.stderr)
        sys.exit(1)


def process_one_video(video_path: Path, output_path: Path, device: str, model=None, save_srt: bool = False):
    """Transcribe one video and burn subtitles to output_path."""
    transcription = transcribe_audio(video_path, device, model=model)
    with tempfile.NamedTemporaryFile(mode="w", suffix=".srt", delete=False) as f:
        srt_path = Path(f.name)
    try:
        generate_srt(transcription, srt_path)
        
        if save_srt:
            srt_output_path = output_path.with_suffix(".srt")
            srt_path.replace(srt_output_path)
            srt_path = srt_output_path
        
        burn_subtitles(video_path, srt_path, output_path)
    finally:
        if srt_path.exists() and not save_srt:
            srt_path.unlink()


def main():
    parser = argparse.ArgumentParser(
        description="Burn English transcripts into MP4 videos using Whisper.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  burn-subs video.mp4              Transcribe and burn subtitles
  burn-subs video.mp4 -s           Also save SRT file
  burn-subs path/to/videos/        Process all MP4s; output to path/to/videos_transcribed/
  burn-subs path/to/videos/ -s     Also save SRT files for all videos

Output:
  File: creates {filename}_subtitled.mp4 next to input.
  Directory: creates {dirname}_transcribed/ (sibling) with each video as {stem}_subtitled.mp4.
  With --save-srt: also creates {filename}_subtitled.srt alongside output.
"""
    )
    parser.add_argument(
        "input",
        help="Input MP4 file or directory containing MP4 files"
    )
    parser.add_argument(
        "--save-srt", "-s",
        action="store_true",
        help="Save SRT subtitle file alongside the output video"
    )

    args = parser.parse_args()
    check_dependencies()
    input_path = Path(args.input).expanduser().resolve()

    if input_path.is_dir():
        in_dir, mp4s = validate_directory(args.input)
        out_dir = in_dir.parent / f"{in_dir.name}_transcribed"
        out_dir.mkdir(parents=True, exist_ok=True)
        print(f"Input dir:  {in_dir}")
        print(f"Output dir: {out_dir} ({len(mp4s)} videos)")
        print()
        device = get_device()
        print()
        model = load_whisper_model(device)
        for i, video_path in enumerate(mp4s, 1):
            print(f"[{i}/{len(mp4s)}] {video_path.name}")
            output_path = out_dir / f"{video_path.stem}_subtitled.mp4"
            process_one_video(video_path, output_path, device, model=model, save_srt=args.save_srt)
            print()
        print(f"Done! Output saved to: {out_dir}")
    else:
        video_path = validate_input(args.input)
        output_path = video_path.parent / f"{video_path.stem}_subtitled.mp4"
        print(f"Input:  {video_path}")
        print(f"Output: {output_path}")
        print()
        device = get_device()
        print()
        process_one_video(video_path, output_path, device, save_srt=args.save_srt)
        print(f"Done! Output saved to: {output_path}")


if __name__ == "__main__":
    main()