Add subtitles to videos use ffmpeg

add_subtitles.py

"""
Add subtitles to video.

Subtitles files should be in same directory with videos,
and with same file name.

usage:
    $ python add_subtitles.py INPUT_DIR OUTPUT_DIR

"""
import os
from os.path import join, splitext, split
from os.path import exists
import sys
import subprocess

SUBTITLE_FORMATS = [".en.srt"]
VIDEO_FORMATS = [".mp4", ".webm"]


def gen_videos_and_subtitle(in_dir):
    files = os.listdir(in_dir)
    for f in files:
        f_ = join(in_dir, f)
        prefix, extension = splitext(f_)
        if extension in VIDEO_FORMATS:
            for fmt in SUBTITLE_FORMATS:
                if prefix.endswith("."):
                    prefix = prefix[:-1]
                s_ = prefix + fmt
                if split(s_)[1] in files:
                    yield f_, s_

def make_tmp_copy(path):
    dir_name = split(path)[0]
    tmp_name = lambda id_: join(dir_name, 'tmp.{}'.format(id_))
    tmp_id = 0
    while exists(tmp_name(tmp_id)):
        tmp_id += 1
    result = tmp_name(tmp_id)
    subprocess.call(["cp", path, result])
    return result

def delete_tmp_copy(tmp_path):
    subprocess.call(["rm", tmp_path])

def add_subtitle(video_path, subtitle_path, output_dir):
    subtitle_tmp = make_tmp_copy(subtitle_path)

    output = join(out_dir, split(video_path)[1])

    if exists(output):
        print(output, "aleady exists.")
        return

    cmd = ["ffmpeg", "-i", video_path, "-strict", "-2",
           "-vf", "subtitles={}".format(subtitle_tmp),
           output]

    print(" ".join(cmd))
    status = subprocess.call(cmd)

    delete_tmp_copy(subtitle_tmp)
    return status


if __name__ == "__main__":
    in_dir = sys.argv[1]
    out_dir = sys.argv[2]
    video_and_subtitles = gen_videos_and_subtitle(in_dir) 
    for video, subtitle in video_and_subtitles:
        add_subtitle(video, subtitle, out_dir)