create_captions.py

import ffmpeg
import sys
import re
import subprocess
import os



def time_to_float(time_str):
    # Split the time string into minutes, seconds, and milliseconds
    minutes, seconds = time_str.split(':')
    seconds, milliseconds = seconds.split('.')
    
    # Convert to float (total seconds)
    return float(minutes) * 60 + float(seconds) + float(milliseconds) / 1000


def create_caps(steps):
    captions=[]
    video_start=0
    for step in steps:
        print(step)
        start=time_to_float(step['start'])
        end= time_to_float(step['end'])
        description_sentences=step['description'].split(".")

        print(description_sentences)

        interval=end/len(description_sentences)
        start_interval=video_start

        for description_sentence in description_sentences:
            caption=(start_interval,start_interval+interval,description_sentence)
            start_interval+=interval
            captions.append(caption)
        video_start+=end
    
    return captions




def create_srt(subtitles, output_file):
    with open(output_file, 'w', encoding='utf-8') as f:
        for i, (start, end, text) in enumerate(subtitles, 1):
            start_time = f"{int(start//3600):02d}:{int((start%3600)//60):02d}:{int(start%60):02d},{int((start%1)*1000):03d}"
            end_time = f"{int(end//3600):02d}:{int((end%3600)//60):02d}:{int(end%60):02d},{int((end%1)*1000):03d}"
            f.write(f"{i}\n{start_time} --> {end_time}\n{text}\n\n")



def add_captions_to_video(input_video,soft_subtitle, subtitles,output_video="output_video_subtitles.mp4",subtitle_language="en"):
    subtitle_file="subtitles.srt"
    create_srt(subtitles,subtitle_file)

    video_input_stream = ffmpeg.input(input_video)
    subtitle_input_stream = ffmpeg.input(subtitle_file)
    output_video = output_video
    subtitle_track_title = subtitle_file.replace(".srt", "")

    if soft_subtitle:
        stream = ffmpeg.output(
            video_input_stream, subtitle_input_stream, output_video, **{"c": "copy", "c:s": "mov_text"},
            **{"metadata:s:s:0": f"language={subtitle_language}",
            "metadata:s:s:0": f"title={subtitle_track_title}"}
        )
        ffmpeg.run(stream, overwrite_output=True)
    else:
        subtitle_style = (
            "force_style='FontName=Arial,FontSize=15,PrimaryColour=&H00FFFFFF,"
            "Alignment=2,MarginV=25'"
        )

        stream = ffmpeg.output(
                video_input_stream,
                output_video,
                **{
                    "c:v": "libx264", 
                    # "q:v": 50, # Specify video codec (H.264)
                    "c:a": "copy",     # Copy audio stream to avoid re-encoding
                    "vf": f"subtitles={subtitle_file}:{subtitle_style}",  # Burn subtitles
                    "threads": 0
                }
            )
            

    print("here")
    ffmpeg.run(stream, overwrite_output=True)

    return output_video







if __name__ == "__main__":
    # Example subtitles
    subtitles = [
        (0, 4, "Hello, this is the first subtitle."),
        (5, 10, "This is the second subtitle.")
    ]
    # create_srt(subtitles, "subtitles.srt")

    # Input video and subtitle files
    input_video = "/Users/georgia.bucea/products/ShortsAI/21_aug/IMG_8280.MOV"
    subtitle_file = "subtitles.srt"
    output_video = "output_video_subtitles.mp4"


    add_captions_to_video(input_video,False,subtitles,output_video,"en")