360 lines
15 KiB
Python
360 lines
15 KiB
Python
import argparse
|
|
import os
|
|
import sys
|
|
from dotenv import load_dotenv
|
|
|
|
# Load environment variables from central .env_files directory
|
|
script_dir = os.path.dirname(os.path.abspath(__file__))
|
|
env_path = os.path.abspath(os.path.join(script_dir, '../../.env_files/.env.aitranscribe'))
|
|
|
|
if os.path.exists(env_path):
|
|
load_dotenv(env_path)
|
|
else:
|
|
# Fallback: check local .env
|
|
local_env = os.path.join(script_dir, '.env')
|
|
if os.path.exists(local_env):
|
|
load_dotenv(local_env)
|
|
else:
|
|
# Last resort: just try loading generic (cwd)
|
|
load_dotenv()
|
|
|
|
from extractor import extract_audio, embed_subtitles
|
|
from transcriber import transcribe_audio, save_as_srt, load_whisper_model
|
|
from translator import translate_with_auto_fallback
|
|
from utils import validate_and_repair_srt, check_srt_duration_match, GracefulKiller, ensure_ollama_running, check_service_availability, check_path_permissions, LANGUAGE_MAP
|
|
from diarizer import diarize_audio, merge_diarization_with_transcript
|
|
import tracker
|
|
from tracker import JobStatus
|
|
|
|
def save_srt_with_speakers(segments, output_path):
|
|
"""Helper to save SRT with speaker labels prepended to text."""
|
|
def format_timestamp(seconds: float):
|
|
whole_seconds = int(seconds)
|
|
milliseconds = int((seconds - whole_seconds) * 1000)
|
|
hours = whole_seconds // 3600
|
|
minutes = (whole_seconds % 3600) // 60
|
|
seconds = whole_seconds % 60
|
|
return f"{hours:02d}:{minutes:02d}:{seconds:02d},{milliseconds:03d}"
|
|
|
|
with open(output_path, "w", encoding="utf-8") as f:
|
|
for i, segment in enumerate(segments, start=1):
|
|
start = format_timestamp(segment["start"])
|
|
end = format_timestamp(segment["end"])
|
|
text = segment["text"].strip()
|
|
speaker = segment.get("speaker", "")
|
|
|
|
if speaker and speaker != "Unknown":
|
|
text = f"[{speaker}]: {text}"
|
|
|
|
f.write(f"{i}\n")
|
|
f.write(f"{start} --> {end}\n")
|
|
f.write(f"{text}\n\n")
|
|
print(f"SRT saved to: {output_path}")
|
|
|
|
def process_file(file_path, args, source_lang=None, loaded_model=None, service_status=None):
|
|
tracker.logger.info(f"=== Processing: {file_path} ===")
|
|
|
|
# Initialize Job
|
|
job = tracker.get_job(file_path)
|
|
|
|
if job.status == JobStatus.COMPLETED and not args.force:
|
|
tracker.logger.info("Job already completed. Skipping.")
|
|
return
|
|
|
|
tracker.update_job_status(file_path, JobStatus.PROCESSING)
|
|
|
|
try:
|
|
# 1. Extract Audio
|
|
tracker.update_step(file_path, "step_extract", "processing")
|
|
audio_path = extract_audio(file_path)
|
|
tracker.update_step(file_path, "step_extract", "done")
|
|
|
|
# 2. Transcribe (Generate SRT)
|
|
tracker.update_step(file_path, "step_transcribe", "processing")
|
|
transcript_file = os.path.splitext(file_path)[0] + ".srt"
|
|
transcript_exists = os.path.exists(transcript_file) and not args.force
|
|
|
|
final_srt_path = transcript_file
|
|
detected_iso = None
|
|
|
|
if transcript_exists:
|
|
tracker.logger.info(f"Transcript exists: {transcript_file}. Skipping transcription.")
|
|
with open(transcript_file, "r", encoding="utf-8") as f:
|
|
srt_content = f.read()
|
|
else:
|
|
# Use loaded_model if available
|
|
result = transcribe_audio(audio_path, model_size=args.model, language=source_lang, loaded_model=loaded_model)
|
|
segments = result["segments"]
|
|
detected_iso = result.get("language")
|
|
|
|
if args.diarize:
|
|
hf_token = args.hf_token or os.getenv("HF_TOKEN")
|
|
if hf_token:
|
|
tracker.logger.info("Running Speaker Diarization...")
|
|
diar_segments = diarize_audio(audio_path, hf_token=hf_token)
|
|
if diar_segments:
|
|
segments = merge_diarization_with_transcript(segments, diar_segments)
|
|
tracker.logger.info("Diarization merged into transcript.")
|
|
else:
|
|
tracker.logger.warning("Warning: --diarize requested but HF_TOKEN not provided. Skipping.")
|
|
|
|
if args.diarize:
|
|
save_srt_with_speakers(segments, transcript_file)
|
|
else:
|
|
save_as_srt(result, transcript_file)
|
|
|
|
validate_and_repair_srt(transcript_file)
|
|
with open(transcript_file, "r", encoding="utf-8") as f:
|
|
srt_content = f.read()
|
|
tracker.update_step(file_path, "step_transcribe", "done")
|
|
|
|
# 3. Translate
|
|
tracker.update_step(file_path, "step_translate", "processing")
|
|
|
|
base_translated = os.path.splitext(file_path)[0] + f".{args.lang}.srt"
|
|
deep_translated = os.path.splitext(file_path)[0] + f".{args.lang}.deep_translate.srt"
|
|
local_translated = os.path.splitext(file_path)[0] + f".{args.lang}.local_llm.srt"
|
|
|
|
# Determine output path logic
|
|
target_path_gemini = base_translated
|
|
target_path_deep = deep_translated
|
|
target_path_local = local_translated
|
|
|
|
translated_file = None
|
|
translation_success = False
|
|
method_used = "None"
|
|
|
|
# Check if source language matches target language
|
|
target_iso = LANGUAGE_MAP.get(args.lang)
|
|
if detected_iso and target_iso and detected_iso == target_iso:
|
|
tracker.logger.info(f"Source language '{detected_iso}' matches target '{target_iso}'. Skipping translation.")
|
|
final_srt_path = transcript_file
|
|
translation_success = True
|
|
method_used = "Source Match"
|
|
|
|
# Check existing (if not already handled by match)
|
|
elif (os.path.exists(base_translated) or os.path.exists(deep_translated) or os.path.exists(local_translated)) and not args.force:
|
|
if os.path.exists(local_translated):
|
|
translated_file = local_translated
|
|
method_used = "Local LLM (Existing)"
|
|
elif os.path.exists(deep_translated):
|
|
translated_file = deep_translated
|
|
method_used = "DeepTranslate (Existing)"
|
|
else:
|
|
translated_file = base_translated
|
|
method_used = "Gemini (Existing)"
|
|
|
|
tracker.logger.info(f"Translation exists: {translated_file} ({method_used}). Skipping translation.")
|
|
final_srt_path = translated_file
|
|
translation_success = True
|
|
else:
|
|
if srt_content:
|
|
res_content, method = translate_with_auto_fallback(
|
|
srt_content,
|
|
target_language=args.lang,
|
|
prefer_deep=args.prefer_deep,
|
|
prefer_local=args.prefer_local,
|
|
available_services=service_status
|
|
)
|
|
|
|
if res_content:
|
|
# Save based on method used
|
|
if "DeepTranslate" in method:
|
|
save_path = target_path_deep
|
|
elif "Local LLM" in method:
|
|
save_path = target_path_local
|
|
else:
|
|
save_path = target_path_gemini
|
|
|
|
with open(save_path, "w", encoding="utf-8") as f:
|
|
f.write(res_content)
|
|
|
|
tracker.logger.info(f"Translation saved to: {save_path} ({method})")
|
|
validate_and_repair_srt(save_path)
|
|
|
|
# Duration Check
|
|
is_valid_duration, msg = check_srt_duration_match(transcript_file, save_path)
|
|
if is_valid_duration:
|
|
tracker.logger.info(f"Validation: {msg}")
|
|
final_srt_path = save_path
|
|
translation_success = True
|
|
method_used = method
|
|
else:
|
|
tracker.logger.error(f"VALIDATION FAILED: {msg}")
|
|
tracker.logger.error("Marking translation as failed due to incomplete coverage.")
|
|
|
|
redo_file = os.path.join(os.path.dirname(file_path), "redo_queue.txt")
|
|
with open(redo_file, "a", encoding="utf-8") as rf:
|
|
rf.write(f"{file_path} | {msg}\n")
|
|
|
|
translation_success = False
|
|
else:
|
|
tracker.logger.error("TRANSLATION FAILED (All methods attempted).")
|
|
tracker.update_step(file_path, "step_translate", "failed")
|
|
translation_success = False
|
|
|
|
if translation_success:
|
|
tracker.update_step(file_path, "step_translate", "done")
|
|
tracker.logger.info(f"Translation Method: {method_used}")
|
|
|
|
# 4. Embed Subtitles
|
|
tracker.update_step(file_path, "step_embed", "processing")
|
|
should_embed = args.embed
|
|
if args.embed and not translation_success:
|
|
tracker.logger.warning("SAFETY HALT: Translation failed. Skipping embedding/deletion.")
|
|
should_embed = False
|
|
|
|
if should_embed:
|
|
success_embed = embed_subtitles(file_path, final_srt_path)
|
|
|
|
if success_embed and args.delete_source:
|
|
base, ext = os.path.splitext(file_path)
|
|
expected_output = f"{base}.subbed{ext}"
|
|
|
|
if os.path.exists(expected_output):
|
|
try:
|
|
os.remove(file_path)
|
|
tracker.logger.info(f"SOURCE DELETED: {file_path}")
|
|
except OSError as e:
|
|
tracker.logger.error(f"Error deleting source: {e}")
|
|
else:
|
|
tracker.logger.error(f"SAFETY ABORT: Output '{expected_output}' not found.")
|
|
tracker.update_step(file_path, "step_embed", "done")
|
|
|
|
# 5. Cleanup
|
|
if args.cleanup:
|
|
try:
|
|
os.remove(audio_path)
|
|
tracker.logger.info(f"Cleanup: Removed {audio_path}")
|
|
except OSError as e:
|
|
tracker.logger.warning(f"Warning: Could not remove audio: {e}")
|
|
|
|
# Mark Complete
|
|
if translation_success:
|
|
tracker.update_job_status(file_path, JobStatus.COMPLETED)
|
|
else:
|
|
tracker.update_job_status(file_path, JobStatus.FAILED, error="Translation failed")
|
|
|
|
except Exception as e:
|
|
tracker.logger.exception(f"Job Failed for {file_path}")
|
|
tracker.update_job_status(file_path, JobStatus.FAILED, error=str(e))
|
|
return
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="AI Video Transcriber & Translator")
|
|
parser.add_argument("inputs", nargs='*', help="Path(s) to video file or directory")
|
|
parser.add_argument("--model", default="auto", choices=["auto", "tiny", "base", "small", "medium", "large"], help="Whisper model size (default: auto)")
|
|
parser.add_argument("--lang", default="English", help="Target language for translation (default: English)")
|
|
parser.add_argument("--source-lang", help="Source language of the audio (e.g., 'fr', 'es'). If omitted, you will be prompted.")
|
|
parser.add_argument("--force", action="store_true", help="Overwrite existing files")
|
|
parser.add_argument("--cleanup", action="store_true", help="Delete temporary .wav file")
|
|
parser.add_argument("--embed", action="store_true", help="Embed subtitles (Soft Subs)")
|
|
parser.add_argument("--diarize", action="store_true", help="Enable speaker diarization")
|
|
parser.add_argument("--hf-token", help="HuggingFace Token")
|
|
parser.add_argument("--delete-source", action="store_true", help="Delete original file after embedding")
|
|
parser.add_argument("--retry-failed", action="store_true", help="Retry FAILED jobs from DB")
|
|
parser.add_argument("--prefer-deep", action="store_true", help="Prefer DeepTranslate (Free) over Gemini")
|
|
parser.add_argument("--prefer-local", action="store_true", help="Prefer Local LLM (Ollama) over cloud APIs")
|
|
|
|
args = parser.parse_args()
|
|
|
|
if not os.getenv("GEMINI_API_KEY"):
|
|
print("Warning: GEMINI_API_KEY environment variable not set. Translation step will fail.")
|
|
|
|
source_lang = args.source_lang
|
|
|
|
if args.retry_failed:
|
|
print("Retrying failed jobs from database...")
|
|
failed_files = tracker.get_failed_jobs()
|
|
if not failed_files:
|
|
print("No failed jobs found.")
|
|
return
|
|
|
|
if not source_lang:
|
|
print("\n--- Audio Configuration ---")
|
|
user_input = input("Enter source language (e.g. 'French'). Enter for Auto: ").strip()
|
|
source_lang = user_input if user_input else None
|
|
|
|
# Load model for retries too
|
|
loaded_model = load_whisper_model(args.model)
|
|
service_status = check_service_availability()
|
|
|
|
for file_path in failed_files:
|
|
if os.path.exists(file_path):
|
|
process_file(file_path, args, source_lang, loaded_model=loaded_model, service_status=service_status)
|
|
else:
|
|
print(f"Skipping missing file: {file_path}")
|
|
return
|
|
|
|
if not args.inputs:
|
|
parser.print_help()
|
|
sys.exit(1)
|
|
|
|
if not source_lang:
|
|
print("\n--- Audio Configuration ---")
|
|
user_input = input("Enter source language (e.g. 'French'). Enter for Auto: ").strip()
|
|
source_lang = user_input if user_input else None
|
|
print(f"Selected: {source_lang if source_lang else 'Auto-detect'}")
|
|
|
|
# --- Ensure Ollama is Running ---
|
|
ensure_ollama_running()
|
|
# --------------------------------
|
|
|
|
# --- Check Service Health ---
|
|
service_status = check_service_availability()
|
|
# ----------------------------
|
|
|
|
# --- Check Path Permissions ---
|
|
valid_inputs = []
|
|
print("Checking Input Permissions...")
|
|
for inp in args.inputs:
|
|
ok, msg = check_path_permissions(inp)
|
|
print(msg)
|
|
if ok:
|
|
valid_inputs.append(inp)
|
|
|
|
if not valid_inputs:
|
|
print("\n❌ Error: No valid inputs with read/write permissions found. Exiting.")
|
|
return
|
|
# ------------------------------
|
|
|
|
# --- Load Model Once ---
|
|
loaded_model = load_whisper_model(args.model)
|
|
# -----------------------
|
|
|
|
# Initialize Graceful Exit Handler
|
|
killer = GracefulKiller()
|
|
|
|
video_extensions = ('.mp4', '.mkv', '.mov', '.avi', '.webm', '.flv', '.wmv', '.m4v')
|
|
|
|
for input_path in valid_inputs:
|
|
if killer.kill_now:
|
|
break
|
|
|
|
if os.path.isfile(input_path):
|
|
process_file(input_path, args, source_lang, loaded_model=loaded_model, service_status=service_status)
|
|
elif os.path.isdir(input_path):
|
|
found = False
|
|
for root, dirs, files in os.walk(input_path):
|
|
if killer.kill_now:
|
|
break
|
|
|
|
for file in files:
|
|
if killer.kill_now:
|
|
break
|
|
|
|
if file.lower().endswith(video_extensions):
|
|
found = True
|
|
process_file(os.path.join(root, file), args, source_lang, loaded_model=loaded_model, service_status=service_status)
|
|
if not found:
|
|
print(f"No video files found in {input_path}")
|
|
else:
|
|
print(f"Error: Invalid input path '{input_path}'")
|
|
|
|
if killer.kill_now:
|
|
print("\n🛑 Process stopped by user. Progress saved in database.")
|
|
else:
|
|
print("\n✅ All jobs finished.")
|
|
|
|
if __name__ == "__main__":
|
|
main() |