app.py

import streamlit as st
import whisper
import os
import tempfile
import subprocess
import srt
import datetime

# Load Whisper model (using cache for performance)
@st.cache_resource
def load_model():
    return whisper.load_model("small")  # Use the small model for free usage

# Transcribe video using Whisper
def transcribe_video(video_path):
    model = load_model()
    result = model.transcribe(video_path)
    return result["segments"]

# Create an SRT file from the transcription segments
def create_srt(subtitles, srt_path):
    subs = []
    for i, segment in enumerate(subtitles):
        # Convert float seconds to timedelta objects
        start_time = datetime.timedelta(seconds=segment["start"])
        end_time = datetime.timedelta(seconds=segment["end"])
        subs.append(srt.Subtitle(
            index=i + 1,
            start=start_time,
            end=end_time,
            content=segment["text"]
        ))
    with open(srt_path, "w", encoding="utf-8") as f:
        f.write(srt.compose(subs))

# Burn subtitles into the video using FFmpeg
def burn_captions(video_path, srt_path, output_path):
    command = [
        "ffmpeg",
        "-i", video_path,
        "-vf", f"subtitles={srt_path}",
        "-c:a", "copy",
        output_path
    ]
    subprocess.run(command, check=True)

# Streamlit app UI
st.title("🎥 AI Video Captioning App")

uploaded_file = st.file_uploader("Upload a video file", type=["mp4", "mkv", "avi", "mov"])

if uploaded_file:
    # Save the uploaded video to a temporary file
    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as temp_video:
        temp_video.write(uploaded_file.read())
        video_path = temp_video.name

    st.video(video_path)

    if st.button("Generate Captions & Burn into Video"):
        with st.spinner("Generating captions..."):
            captions = transcribe_video(video_path)

        srt_path = video_path.replace(".mp4", ".srt")
        create_srt(captions, srt_path)

        output_video_path = video_path.replace(".mp4", "_captioned.mp4")

        with st.spinner("Burning captions into video..."):
            burn_captions(video_path, srt_path, output_video_path)

        st.success("Processing complete! Download your video below.")
        with open(output_video_path, "rb") as file:
            st.download_button("📥 Download Captioned Video", file, file_name="captioned_video.mp4", mime="video/mp4")