Spaces:
Build error
Build error
File size: 4,902 Bytes
6159702 3a12b87 6159702 3a12b87 6159702 3a12b87 6159702 3a12b87 6159702 3a12b87 6159702 3a12b87 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | import os
import subprocess
import tempfile
from pathlib import Path
from typing import Optional
import streamlit as st
from faster_whisper import WhisperModel
from googletrans import Translator
# ---------------- CONFIG ----------------
MODEL_NAME = os.environ.get("WHISPER_MODEL", "large-v2")
DEVICE = "cuda" if (os.environ.get("CUDA_VISIBLE_DEVICES") or False) else "cpu"
_model = None
def get_model():
global _model
if _model is None:
compute_type = "float16" if DEVICE.startswith("cuda") else "int8"
_model = WhisperModel(MODEL_NAME, device=DEVICE, compute_type=compute_type)
return _model
def extract_audio(input_video_path: str, output_audio_path: str):
cmd = [
"ffmpeg", "-y", "-i", input_video_path,
"-vn", "-acodec", "pcm_s16le", "-ar", "16000", "-ac", "1",
output_audio_path,
]
subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
def segments_to_srt(segments):
def fmt_time(s):
h = int(s // 3600)
m = int((s % 3600) // 60)
sec = s % 60
return f"{h:02d}:{m:02d}:{sec:06.3f}".replace('.', ',')
srt_lines = []
for i, seg in enumerate(segments, start=1):
start = fmt_time(seg["start"])
end = fmt_time(seg["end"])
text = seg["text"].strip()
srt_lines.append(f"{i}\n{start} --> {end}\n{text}\n")
return "\n".join(srt_lines)
def transcribe_and_translate(video_file: str, target_lang: Optional[str], burn_subs: bool):
model = get_model()
tempdir = Path(tempfile.mkdtemp())
input_path = Path(video_file)
audio_path = tempdir / "audio.wav"
srt_path = tempdir / f"subtitles_{input_path.stem}.srt"
processed_video_path = None
extract_audio(str(input_path), str(audio_path))
segments_all = []
transcribe_options = {"beam_size": 5, "word_timestamps": False}
for segment in model.transcribe(str(audio_path), beam_size=5, vad_filter=True, **transcribe_options):
segments_all.append({"start": segment.start, "end": segment.end, "text": segment.text})
if target_lang and target_lang.lower() not in ["", "none"]:
translator = Translator()
translated_segments = []
for seg in segments_all:
src_text = seg["text"].strip()
try:
res = translator.translate(src_text, dest=target_lang)
translated_text = res.text
except Exception:
translated_text = src_text
translated_segments.append({"start": seg["start"], "end": seg["end"], "text": translated_text})
segments_used = translated_segments
else:
segments_used = segments_all
srt_text = segments_to_srt(segments_used)
srt_path.write_text(srt_text, encoding="utf-8")
if burn_subs:
out_video = tempdir / f"burned_{input_path.name}"
cmd = [
"ffmpeg", "-y", "-i", str(input_path),
"-vf", f"subtitles={str(srt_path)}:force_style='FontName=Arial,FontSize=24'",
"-c:a", "copy", str(out_video),
]
subprocess.run(cmd, check=True)
processed_video_path = str(out_video)
return str(srt_path), processed_video_path
# ---------------- UI (Streamlit) ----------------
st.set_page_config(page_title="Video Subtitle Editor + Translator", layout="wide")
st.title("🎬 Video Subtitle Editor + Translator (Streamlit)")
video_file = st.file_uploader("Upload your video (mp4, mov, mkv)", type=["mp4", "mov", "mkv"])
lang_choice = st.selectbox(
"Translate subtitles to:",
["None", "English (en)", "Urdu (ur)", "Hindi (hi)", "Spanish (es)", "French (fr)", "German (de)"]
)
burn_option = st.checkbox("Burn subtitles into video (hardcoded) - slow but permanent", value=False)
if st.button("Run"):
if video_file is None:
st.warning("Please upload a video file first.")
else:
with st.spinner("Processing... Please wait ⏳"):
temp_input = Path(tempfile.mkdtemp()) / video_file.name
with open(temp_input, "wb") as f:
f.write(video_file.read())
lang_code = lang_choice.split("(")[-1].replace(")", "").strip().lower() if "(" in lang_choice else "none"
try:
srt_path, processed_video = transcribe_and_translate(str(temp_input), lang_code, burn_option)
st.success("✅ Done! Subtitles generated successfully.")
st.download_button("📥 Download SRT", open(srt_path, "rb"), file_name=os.path.basename(srt_path))
if processed_video:
st.download_button("📽️ Download Video with Subtitles", open(processed_video, "rb"), file_name=os.path.basename(processed_video))
except subprocess.CalledProcessError as e:
st.error(f"ffmpeg error: {e}")
except Exception as e:
st.error(f"Error: {e}")
|