Initial commit: local video captioning and translation tool
This commit is contained in:
Adolfo Reyna
+356
@@ -0,0 +1,356 @@
|
||||
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
|
||||
Reference in New Issue
Block a user