VideoCaptioningTranslation/main.py

import os
import sys
import argparse
import time
import shutil
import config
from translator import extract_audio, transcribe_audio, translate_segments, write_srt, burn_subtitles

def check_dependencies():
    """Verify that FFmpeg is installed and accessible in system path"""
    if not shutil.which("ffmpeg"):
        print("[Error] FFmpeg is not installed or not in your system PATH.")
        print("Please install FFmpeg to run this tool. On macOS: 'brew install ffmpeg'")
        sys.exit(1)

def main():
    check_dependencies()
    
    parser = argparse.ArgumentParser(description="Video Captioning & Translation Tool using OpenAI API + Local FFmpeg")
    
    parser.add_argument("-v", "--video", type=str, required=True, help="Path to the input video file")
    parser.add_argument("-s", "--source-lang", type=str, default=config.DEFAULT_SOURCE_LANG, 
                        help=f"Source language of the video audio (default: {config.DEFAULT_SOURCE_LANG})")
    parser.add_argument("-t", "--target-lang", type=str, default=config.DEFAULT_TARGET_LANG, 
                        help=f"Target language for the captions (default: {config.DEFAULT_TARGET_LANG})")
    parser.add_argument("-o", "--output", type=str, default=None, 
                        help="Path for the output captioned video (default: <video_name>_<target>_captioned.mp4)")
    parser.add_argument("--whisper-model", type=str, default=config.DEFAULT_WHISPER_MODEL,
                        help=f"OpenAI Whisper model name (default: {config.DEFAULT_WHISPER_MODEL})")
    parser.add_argument("--model", type=str, default=config.DEFAULT_OPENAI_MODEL,
                        help=f"OpenAI GPT model name for translation (default: {config.DEFAULT_OPENAI_MODEL})")
    parser.add_argument("--keep-temp", action="store_true", default=False,
                        help="Keep intermediate audio (.wav) and subtitle (.srt) files")
    parser.add_argument("--style", type=str, choices=list(config.SUBTITLE_PRESETS.keys()), default=config.DEFAULT_STYLE,
                        help=f"Subtitle style preset: {', '.join(config.SUBTITLE_PRESETS.keys())} (default: {config.DEFAULT_STYLE})")
    parser.add_argument("--custom-style", type=str, default=None,
                        help="Raw ASS style override (e.g. 'FontName=Arial,FontSize=24,PrimaryColour=&H0000FFFF')")
    parser.add_argument("--max-chars", type=int, default=0,
                        help="Force subtitle line length limit by splitting segments (default: 0 for no splitting)")
    
    args = parser.parse_args()
    
    video_path = os.path.abspath(args.video)
    if not os.path.exists(video_path):
        print(f"[Error] Input video file does not exist: {args.video}")
        sys.exit(1)
        
    # Define output file path if not specified
    video_dir = os.path.dirname(video_path)
    video_name, video_ext = os.path.splitext(os.path.basename(video_path))
    
    if args.output:
        output_path = os.path.abspath(args.output)
    else:
        output_path = os.path.join(video_dir, f"{video_name}_{args.target_lang}_captioned{video_ext}")
        
    # Use space-free temporary names to prevent FFmpeg subtitle filter parsing errors
    temp_audio_path = os.path.join(video_dir, f"temp_audio_{int(time.time())}.wav")
    temp_srt_path = os.path.join(video_dir, f"temp_subtitles_{int(time.time())}.srt")
    
    # User-friendly SRT path for when keep-temp is enabled
    final_srt_path = os.path.join(video_dir, f"{video_name}_{args.target_lang}.srt")
    
    print("\n" + "="*50)
    print("      VIDEO CAPTIONING & TRANSLATION TOOL (OPENAI)")
    print("="*50)
    print(f"Input Video:      {video_path}")
    print(f"Source Language:  {args.source_lang}")
    print(f"Target Language:  {args.target_lang}")
    print(f"Output Video:     {output_path}")
    print(f"Whisper Model:    {args.whisper_model}")
    print(f"Translation Model:{args.model}")
    print("="*50 + "\n")
    
    pipeline_start = time.time()
    
    # Step 1: Extract Audio
    step_start = time.time()
    if not extract_audio(video_path, temp_audio_path):
        print("[Error] Failed to extract audio. Exiting.")
        sys.exit(1)
    audio_time = time.time() - step_start
    
    # Step 2: Transcribe Audio via OpenAI API
    step_start = time.time()
    try:
        segments = transcribe_audio(temp_audio_path, args.whisper_model)
    except Exception as e:
        print(f"[Error] ASR failed: {e}")
        if not args.keep_temp and os.path.exists(temp_audio_path):
            os.remove(temp_audio_path)
        sys.exit(1)
    transcribe_time = time.time() - step_start
    
    if not segments:
        print("[Warning] No speech detected in video audio.")
        shutil.copy2(video_path, output_path)
        print(f"[Finished] Output copied to {output_path}")
        if not args.keep_temp and os.path.exists(temp_audio_path):
            os.remove(temp_audio_path)
        sys.exit(0)
        
    # Step 3: Translate segments via OpenAI API
    step_start = time.time()
    try:
        translated_segments = translate_segments(
            segments=segments,
            source_lang=args.source_lang,
            target_lang=args.target_lang,
            model=args.model
        )
    except Exception as e:
        print(f"[Error] Translation failed: {e}")
        if not args.keep_temp and os.path.exists(temp_audio_path):
            os.remove(temp_audio_path)
        sys.exit(1)
    translation_time = time.time() - step_start
    
    # Optional: Adjust subtitle cadence by splitting long segments
    if args.max_chars > 0:
        from translator import split_long_segments
        translated_segments = split_long_segments(translated_segments, args.max_chars)
    
    # Step 4: Write Subtitles
    write_srt(translated_segments, temp_srt_path)
    
    # Step 5: Burn Subtitles locally using FFmpeg
    step_start = time.time()
    burn_success = burn_subtitles(
        video_path=video_path, 
        srt_path=temp_srt_path, 
        output_path=output_path, 
        style_preset=args.style, 
        custom_style=args.custom_style
    )
    burn_time = time.time() - step_start
    
    # Clean up temp files
    if args.keep_temp:
        friendly_audio_path = os.path.join(video_dir, f"{video_name}_temp_audio.wav")
        if os.path.exists(temp_audio_path):
            shutil.move(temp_audio_path, friendly_audio_path)
            print(f"[Cleanup] Saved audio track copy at: {friendly_audio_path}")
        if os.path.exists(temp_srt_path):
            shutil.move(temp_srt_path, final_srt_path)
            print(f"[Cleanup] Saved subtitle file copy at: {final_srt_path}")
    else:
        print("[Cleanup] Removing temporary files...")
        if os.path.exists(temp_audio_path):
            os.remove(temp_audio_path)
        if os.path.exists(temp_srt_path):
            os.remove(temp_srt_path)
            
    if burn_success:
        print("\n" + "="*50)
        print(f"SUCCESS: Captioned video generated!")
        print(f"Saved at: {output_path}")
        print("-"*50)
        print(f"Timing Breakdown:")
        print(f"  - Audio Extraction: {audio_time:.1f}s")
        print(f"  - OpenAI ASR:       {transcribe_time:.1f}s")
        print(f"  - OpenAI Translate: {translation_time:.1f}s")
        print(f"  - FFmpeg Burning:   {burn_time:.1f}s")
        print(f"  - Total Elapsed:    {time.time() - pipeline_start:.1f}s")
        print("="*50 + "\n")
    else:
        print("\n[Error] Failed to generate captioned video during subtitle burning phase.")
        sys.exit(1)

if __name__ == "__main__":
    main()