VideoCaptioningTranslation/translator.py

import os
import subprocess
import time
from tqdm import tqdm
from openai import OpenAI
import config

def format_timestamp(seconds):
    """Convert float seconds to SRT timestamp format: HH:MM:SS,mmm"""
    hours = int(seconds // 3600)
    minutes = int((seconds % 3600) // 60)
    secs = int(seconds % 60)
    milliseconds = int(round((seconds % 1) * 1000))
    if milliseconds == 1000:
        secs += 1
        milliseconds = 0
    return f"{hours:02d}:{minutes:02d}:{secs:02d},{milliseconds:03d}"

def extract_audio(video_path, audio_path):
    """Extract mono 16kHz audio from a video using FFmpeg"""
    print(f"[Audio] Extracting audio from {os.path.basename(video_path)}...")
    cmd = [
        "ffmpeg", "-y",
        "-i", video_path,
        "-ar", "16000",
        "-ac", "1",
        "-c:a", "pcm_s16le",
        audio_path
    ]
    try:
        subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=True)
        print("[Audio] Audio extraction complete.")
        return True
    except subprocess.CalledProcessError as e:
        print(f"[Audio Error] FFmpeg extraction failed: {e.stderr}")
        return False

def transcribe_audio(audio_path, model_name=config.DEFAULT_WHISPER_MODEL):
    """Transcribe audio file using OpenAI's Whisper API and return segments with timestamps"""
    if not config.OPENAI_API_KEY:
        raise ValueError("OPENAI_API_KEY is not configured in .env file.")
        
    print(f"[ASR] Sending audio to OpenAI Whisper API using model '{model_name}'...")
    client = OpenAI(api_key=config.OPENAI_API_KEY)
    
    start_time = time.time()
    with open(audio_path, "rb") as audio_file:
        transcript = client.audio.transcriptions.create(
            model=model_name,
            file=audio_file,
            response_format="verbose_json"
        )
    duration = time.time() - start_time
    print(f"[ASR] Transcription complete in {duration:.2f} seconds.")
    
    # Extract segments
    raw_segments = []
    if hasattr(transcript, "segments"):
        raw_segments = transcript.segments
    elif isinstance(transcript, dict):
        raw_segments = transcript.get("segments", [])
    else:
        raw_segments = getattr(transcript, "segments", [])
        
    segments = []
    for seg in raw_segments:
        if isinstance(seg, dict):
            segments.append({
                "start": seg.get("start", 0.0),
                "end": seg.get("end", 0.0),
                "text": seg.get("text", "").strip()
            })
        else:
            segments.append({
                "start": getattr(seg, "start", 0.0),
                "end": getattr(seg, "end", 0.0),
                "text": getattr(seg, "text", "").strip()
            })
            
    print(f"[ASR] Transcribed {len(segments)} segments.")
    return segments

class OpenAIAPITranslator:
    def __init__(self, model=config.DEFAULT_OPENAI_MODEL, api_key=None):
        key = api_key or config.OPENAI_API_KEY
        if not key:
            raise ValueError("OpenAI API key must be provided or set in environment.")
        self.client = OpenAI(api_key=key)
        self.model = model

    def translate_text(self, text, source_lang, target_lang, context_history=None):
        """Translate a single block of text using OpenAI GPT model, with optional context history"""
        source_name = config.SUPPORTED_LANGUAGES.get(source_lang, source_lang)
        target_name = config.SUPPORTED_LANGUAGES.get(target_lang, target_lang)
        
        system_prompt = (
            f"You are an expert, professional subtitle translator.\n"
            f"Your task is to translate the given text from {source_name} to {target_name}.\n\n"
            f"Hard Rules:\n"
            f"1. Preserve the meaning, tone, and slang style accurately.\n"
            f"2. Output ONLY the translated text. Do not include any notes, explanations, introductory/outro remarks, quotes, or formatting.\n"
            f"3. Make it natural and fitting for screen subtitles (short, punchy sentences where appropriate).\n"
            f"4. If you cannot translate it (e.g. it's nonsensical), output the original text unchanged."
        )
        
        user_content = ""
        if context_history:
            user_content += "--- CONTEXT ---\n"
            for prev_src, prev_trans in context_history:
                user_content += f"Previous original: \"{prev_src}\"\n"
                user_content += f"Previous translation: \"{prev_trans}\"\n\n"
            user_content += "--- CURRENT SEGMENT TO TRANSLATE ---\n"
            
        user_content += f"Original text: \"{text}\"\nTranslation:"

        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_content}
        ]

        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=messages,
                temperature=0.3,
                max_tokens=256
            )
            translated = response.choices[0].message.content.strip()
            # Remove surrounding quotes if model added them
            if translated.startswith('"') and translated.endswith('"'):
                translated = translated[1:-1]
            return translated
        except Exception as e:
            print(f"\n[OpenAI Error] Translation failed: {e}")
            return text

def translate_segments(segments, source_lang, target_lang, model=config.DEFAULT_OPENAI_MODEL):
    """Translate Whisper segments list using OpenAIAPITranslator with a rolling context window"""
    print(f"[Translation] Translating segments using OpenAI model '{model}' from '{source_lang}' to '{target_lang}'...")
    translator = OpenAIAPITranslator(model=model)
    
    translated_segments = []
    context_window = []  # Keep a rolling window of (original, translation) pairs (max 2)
    
    for segment in tqdm(segments, desc="Translating", unit="segment"):
        text = segment.get("text", "").strip()
        if not text:
            translated_segments.append({
                "start": segment["start"],
                "end": segment["end"],
                "text": ""
            })
            continue
            
        translated_text = translator.translate_text(
            text=text,
            source_lang=source_lang,
            target_lang=target_lang,
            context_history=context_window
        )
        
        translated_segments.append({
            "start": segment["start"],
            "end": segment["end"],
            "text": translated_text
        })
        
        # Update rolling context
        context_window.append((text, translated_text))
        if len(context_window) > 2:
            context_window.pop(0)
            
    print(f"[Translation] Completed translation of {len(translated_segments)} segments.")
    return translated_segments

def split_long_segments(segments, max_chars=40):
    """Split segments that exceed max_chars into smaller, proportionally-timed segments"""
    if not max_chars or max_chars <= 0:
        return segments
        
    refined = []
    for seg in segments:
        text = seg["text"].strip()
        start = seg["start"]
        end = seg["end"]
        duration = end - start
        
        if len(text) <= max_chars or duration <= 0:
            refined.append(seg)
            continue
            
        words = text.split()
        chunks = []
        current_chunk = []
        current_len = 0
        
        for word in words:
            word_len = len(word)
            added_len = word_len + (1 if current_len > 0 else 0)
            
            if current_len + added_len > max_chars and current_chunk:
                chunks.append(" ".join(current_chunk))
                current_chunk = [word]
                current_len = word_len
            else:
                current_chunk.append(word)
                current_len += added_len
                
        if current_chunk:
            chunks.append(" ".join(current_chunk))
            
        # Distribute timing proportionally
        total_chars = sum(len(c) for c in chunks)
        if total_chars == 0:
            refined.append(seg)
            continue
            
        current_time = start
        for chunk in chunks:
            chunk_len = len(chunk)
            chunk_dur = duration * (chunk_len / total_chars)
            refined.append({
                "start": current_time,
                "end": current_time + chunk_dur,
                "text": chunk
            })
            current_time += chunk_dur
            
    print(f"[Cadence] Split {len(segments)} segments into {len(refined)} segments (max_chars={max_chars}).")
    return refined

def write_srt(segments, srt_path):
    """Write subtitle segments to an SRT file"""
    print(f"[Subtitles] Writing subtitles to {srt_path}...")
    with open(srt_path, "w", encoding="utf-8") as f:
        for idx, segment in enumerate(segments, 1):
            start_str = format_timestamp(segment["start"])
            end_str = format_timestamp(segment["end"])
            text = segment["text"].replace("\n", " ").strip()
            
            f.write(f"{idx}\n")
            f.write(f"{start_str} --> {end_str}\n")
            f.write(f"{text}\n\n")
    print("[Subtitles] SRT file written.")

def check_subtitles_filter_support():
    """Check if the subtitles filter is available in FFmpeg"""
    try:
        result = subprocess.run(
            ["ffmpeg", "-filters"], 
            stdout=subprocess.PIPE, 
            stderr=subprocess.PIPE, 
            text=True
        )
        return " subtitles " in result.stdout
    except Exception:
        return False

def embed_subtitles_soft(video_path, srt_path, output_path):
    """Embed subtitles as a soft track in the video using mov_text codec (requires no filter/libass)"""
    print(f"[Video] Embedding subtitles as soft track into video...")
    
    video_dir = os.path.dirname(os.path.abspath(video_path))
    srt_name = os.path.basename(srt_path)
    video_name = os.path.basename(video_path)
    output_name = os.path.basename(output_path)
    
    cmd = [
        "ffmpeg", "-y",
        "-i", video_name,
        "-i", srt_name,
        "-c:v", "copy",
        "-c:a", "copy",
        "-c:s", "mov_text",
        output_name
    ]
    
    try:
        subprocess.run(
            cmd,
            cwd=video_dir,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            check=True
        )
        print(f"[Video] Subtitles successfully soft-embedded into {output_path}")
        return True
    except subprocess.CalledProcessError as e:
        print(f"[Video Error] Soft-embedding subtitles failed: {e.stderr}")
        return False

def burn_subtitles(video_path, srt_path, output_path, style_preset="box", custom_style=None):
    """Adds subtitles to the video, either burning them (if supported) or softcoding them (as fallback)"""
    if not check_subtitles_filter_support():
        print("\n" + "!"*60)
        print("[Warning] Your current FFmpeg build does not support hardcoding (burning) subtitles.")
        print("          It is missing the 'subtitles' video filter (usually due to a missing libass library).")
        print("          We are falling back to soft-coding (embedding a subtitle track) instead.")
        print("          To fix this and enable hardcoded/burned captions, run:")
        print("              brew reinstall ffmpeg")
        print("!"*60 + "\n")
        return embed_subtitles_soft(video_path, srt_path, output_path)

    print(f"[Video] Burning subtitles from {os.path.basename(srt_path)} into video...")
    
    video_dir = os.path.dirname(os.path.abspath(video_path))
    srt_name = os.path.basename(srt_path)
    video_name = os.path.basename(video_path)
    output_name = os.path.basename(output_path)
    
    # Resolve style string
    style_str = config.SUBTITLE_PRESETS.get(style_preset, config.SUBTITLE_PRESETS["box"])
    if custom_style:
        style_str = custom_style
        
    print(f"[Video] Using style settings: {style_str}")
    
    cmd = [
        "ffmpeg", "-y",
        "-i", video_name,
        "-vf", f"subtitles={srt_name}:force_style='{style_str}'",
        "-c:a", "copy",  # Copy audio stream without re-encoding
        output_name
    ]
    
    try:
        print(f"[Video] Launching FFmpeg subtitle filter in directory: {video_dir}")
        subprocess.run(
            cmd,
            cwd=video_dir,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            check=True
        )
        print(f"[Video] Subtitles successfully burned into {output_path}")
        return True
    except subprocess.CalledProcessError as e:
        print(f"[Video Error] FFmpeg subtitle burning failed: {e.stderr}")
        print("[Video] Attempting absolute path fallback...")
        escaped_srt = os.path.abspath(srt_path).replace(":", "\\:").replace("'", "'\\''")
        fallback_cmd = [
            "ffmpeg", "-y",
            "-i", video_path,
            "-vf", f"subtitles='{escaped_srt}':force_style='{style_str}'",
            "-c:a", "copy",
            output_path
        ]
        try:
            subprocess.run(fallback_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=True)
            print(f"[Video] Subtitles successfully burned into {output_path} (via absolute path)")
            return True
        except subprocess.CalledProcessError as err:
            print(f"[Video Error] Fallback FFmpeg failed: {err.stderr}")
            return False