autohighlights/efficientjson.py

#!/bin/python
"""jsonclips.py - generate highlights based on JSON

usage: jsonclips.py clips.json

JSON file can specify:
    - output filename
    - source filepath
    - start time
    - duration
    - transition type [default: luma]
    - transition duration [default: 30]
"""

TRANSITION_DURATION = 2  # seconds

FFMPEG_COMMON_ARGS = ["ffmpeg", "-loglevel", "quiet", "-hide_banner"]


def process_highlights(jsonfile):
    """Go through highlights, make clips, join to a video"""
    import tempfile

    def parse_time(timestring):
        from datetime import datetime
        """Custom time parsing function"""
        if timestring.count(":") == 0:
            # probably seconds
            formatstring = "%S"
        elif timestring.count(":") == 1:
            formatstring = "%M:%S"
        elif timestring.count(":") == 2:
            formatstring = "%H:%M:%S"

        if "." in timestring:
            formatstring += ".%f"

        return datetime.strptime(timestring, formatstring)

    def datetime_to_seconds(dt):
        """Take a datetime object and convert to seconds"""
        from datetime import datetime

        zero = datetime.strptime("0", "%S")
        return (dt - zero).seconds

    def make_clips(sources=None, outputdir=None):
        """Use ffmpeg to create output from processed JSON"""
        import subprocess
        import os

        clipnum = 0

        for source in sources:
            # decide on duration vs start+end
            if "duration" not in source:
                try:
                    start = parse_time(source["start"])
                    end = parse_time(source["end"])
                    duration = (end - start).seconds
                    start = datetime_to_seconds(start)
                except ValueError as e:
                    print("Error: {}".format(e))
            else:
                start = source["start"]
                duration = source["duration"]

            print("Making {}s clip from {} at {}".format(
                duration, source["filepath"], start))

            ffmpeg_file_args = ["-ss", str(start),
                                "-i", source["filepath"],
                                "-t", str(duration),
                                "-c", "copy",
                                os.path.join(outputdir,
                                             "clip{}.mkv".format(clipnum))]

            ffmpeg_args = FFMPEG_COMMON_ARGS + ffmpeg_file_args
            print("Running ffmpeg with args:\n\n{}".
                  format(ffmpeg_args))
            subprocess.run(ffmpeg_args)

            clipnum += 1

    def concat(workingdir=None, outputfile=None):
        """Use ffmpeg concat to join clips"""
        import subprocess
        import os

        joinfile = os.path.join(workingdir, "join.txt")
        ffmpeg_concat_args = ["-y", "-safe", "0",
                              "-f", "concat",
                              "-i", str(joinfile),
                              "-c", "copy",
                              str(outputfile)]

        print(FFMPEG_COMMON_ARGS + ffmpeg_concat_args)
        subprocess.run(FFMPEG_COMMON_ARGS + ffmpeg_concat_args)

    def parse_json(jsonfile=None):
        """Parse global / per-clip options from jsonfile

        Requires:
            - outputfile
            - list of sources
        """
        import json
        import sys

        if jsonfile:
            with open(jsonfile, "r") as fh:
                json_in = json.load(fh)

        if "outputfile" not in json_in:
            print("Please specify an outputfile to write to")
            sys.exit(1)

        if "sources" not in json_in:
            print("Please specify some sources to process")
            sys.exit(1)

        return json_in

    def split_first(videofile=None):
        """Split first videofile into videofileM + videofileO"""
        from basicclips import get_video_duration
        import subprocess

        video_length = get_video_duration(videofile)

        # main clip: full length minus half transition duration

        # cut on keyframe
        cut_time = nearest_keyframe(videofile, (video_length - 1))

        ffmpeg_video_args = ["-i", videofile,
                             "-t", str(cut_time),
                             "-c", "copy",
                             "{}".format(videofile.replace(".mkv", "M.mkv"))]

        subprocess.run(FFMPEG_COMMON_ARGS + ffmpeg_video_args)

        # omega clip: half transition duration
        ffmpeg_video_args = ["-i", videofile,
                             "-ss", str(cut_time),
                             "-c:v", "libx264", "-crf", "20", "-c:a", "copy",
                             "{}".format(videofile.replace(".mkv", "O.mkv"))]

        subprocess.run(FFMPEG_COMMON_ARGS + ffmpeg_video_args)

    def split_last(videofile=None):
        """Split last videofile into videofileA + videofileM"""
        import subprocess

        # cut on keyframe
        cut_time = nearest_keyframe(videofile, TRANSITION_DURATION/2)

        # alpha clip: half transition duration
        ffmpeg_video_args = ["-i", videofile,
                             "-t", str(cut_time),
                             "-c:v", "libx264", "-crf", "20", "-c:a", "copy",
                             "{}".format(videofile.replace(".mkv", "A.mkv"))]

        subprocess.run(FFMPEG_COMMON_ARGS + ffmpeg_video_args)

        # main clip: start at half transition duration, full length
        ffmpeg_video_args = ["-i", videofile,
                             "-ss", str(cut_time),
                             "-c", "copy",
                             "{}".format(videofile.replace(".mkv", "M.mkv"))]

        subprocess.run(FFMPEG_COMMON_ARGS + ffmpeg_video_args)

    def split_middle(videofile=None):
        """Split middle videofile into videofileA + videofileM + videofileO"""
        from basicclips import get_video_duration
        import subprocess

        video_length = get_video_duration(videofile)

        # cut on keyframes
        cut_time1 = nearest_keyframe(videofile, TRANSITION_DURATION/2)
        cut_time2 = nearest_keyframe(videofile, (
            video_length - (TRANSITION_DURATION/2)))

        # alpha clip: half transition duration
        ffmpeg_video_args = ["-i", videofile,
                             "-t", str(cut_time1),
                             "-c:v", "libx264", "-crf", "20", "-c:a", "copy",
                             "{}".format(videofile.replace(".mkv", "A.mkv"))]

        subprocess.run(FFMPEG_COMMON_ARGS + ffmpeg_video_args)

        # main clip: full length minus half transition duration
        ffmpeg_video_args = ["-i", videofile,
                             "-ss", str(cut_time1),
                             "-t", str(cut_time2),
                             "-c", "copy",
                             "{}".format(videofile.replace(".mkv", "M.mkv"))]

        subprocess.run(FFMPEG_COMMON_ARGS + ffmpeg_video_args)

        # omega clip: half transition duration
        ffmpeg_video_args = ["-i", videofile,
                             "-ss",
                             str(cut_time2),
                             "-c:v", "libx264", "-crf", "20", "-c:a", "copy",
                             "{}".format(videofile.replace(".mkv", "O.mkv"))]

        subprocess.run(FFMPEG_COMMON_ARGS + ffmpeg_video_args)

    def transition_join(video1=None, video2=None, workingdir=None, i=None):
        """Join videos 1 and 2 with a transition"""
        from moviepy.editor import VideoFileClip
        from moviepy.editor import CompositeVideoClip

        tr = TRANSITION_DURATION/2

        try:
            clip1 = VideoFileClip(os.path.join(workingdir, video1))
        except IOError as e:
            print("Error making joined clip: {}".format(e))
            import pdb
            pdb.set_trace()

        try:
            clip2 = VideoFileClip(os.path.join(workingdir, video2))
        except IOError as e:
            print("Error making joined clip: {}".format(e))
            import pdb
            pdb.set_trace()

        final = CompositeVideoClip([clip1,
                                    clip2.set_start(clip1.end-tr)
                                    .crossfadein(tr)])

        outputfile = os.path.join(workingdir, "clip{}J.mkv".format(i))
        final.write_videofile(outputfile, codec="libx264", audio_codec="aac")

        return outputfile

    def nearest_keyframe(videofile=None, checktime=0):
        """Find nearest keyframe to checktime using ffprobe

        see eg https://superuser.com/a/1426307"""
        import subprocess
        print("Finding nearest keyframe to {} for {}".format(
            checktime, videofile))

        if checktime > 15:
            skiptime = checktime - 15
        else:
            skiptime = checktime

        frames = []

        ffprobe_args = ["ffprobe", "-read_intervals", str(skiptime),
                        "-select_streams", "v",
                        "-skip_frame", "nokey",  # skip all non keyframes
                        "-show_frames", "-show_entries",
                        "frame=pkt_pts_time",
                        videofile]

        # complete = subprocess.run(ffprobe_args, capture_output=True)
        proc = subprocess.Popen(ffprobe_args,
                                stdout=subprocess.PIPE,
                                stderr=subprocess.DEVNULL,
                                universal_newlines=True)
        try:
            lines, error = proc.communicate(timeout=10)
        except subprocess.TimoutExpired:
            proc.kill()
            lines, error = proc.communicate()

        print(lines)
        for line in lines.split("\n"):
            if "pkt_pts_time" in line:
                frametime = float(line.split("=")[1])
                if round(frametime, 2) != 0.00:  # avoid 0s iframe
                    frames.append(frametime)
            if len(frames) > 20:
                break

        # if "stdout" in complete:
        #    for line in complete["stdout"].split("\n"):
        #        if "pkt_pts_time" in line:
        #            frames.extend(line.split("=")[1])

        # get closest value, see https://stackoverflow.com/a/12141207
        keyframetime = min(frames, key=lambda x: abs(x-checktime))

        return str(round(keyframetime, 2))
        # we return the time just before it
        # so that ffmpeg seeks to that when
        # cutting

    # Start
    import os

    if jsonfile:
        json_in = parse_json(jsonfile)

    # in temporary dir
    with tempfile.TemporaryDirectory() as tmpdir:
        # make main clips using codec copy
        # from clip0.mkv to clip[len(sources)-1]
        make_clips(sources=json_in["sources"], outputdir=tmpdir)
        # make subclips (main + transition subclips)
        # add to concat 'join.txt' as we go
        joinfile = os.path.join(tmpdir, "join.txt")
        with open(joinfile, "w") as fh:
            for i in range(len(json_in["sources"])):
                # First clip needs only main + end
                if i == 0:
                    split_first(os.path.join(tmpdir, "clip0.mkv"))
                    fh.write("file '{}'\n".format(
                        os.path.join(tmpdir, "clip0M.mkv")))
                # End clip needs only start + main
                elif i == len(json_in["sources"]):
                    split_last(os.path.join(tmpdir, "clip{}.mkv".format(i)))
                    # Join previous clip omega + this clip alpha
                    joined = transition_join("clip{}O.mkv".format(i-1),
                                             "clip{}A.mkv".format(i),
                                             tmpdir,
                                             i)
                    fh.write("file '{}'\n".format(joined))
                    fh.write("file '{}'\n".format(
                        os.path.join(tmpdir, "clip{}M.mkv".format(i))))
                # Others need start + main + end
                else:
                    split_middle(os.path.join(tmpdir, "clip{}.mkv".format(i)))
                    # Join previous clip omega + this clip alpha
                    joined = transition_join("clip{}O.mkv".format(i-1),
                                             "clip{}A.mkv".format(i),
                                             tmpdir,
                                             i)
                    fh.write("file '{}'\n".format(joined))
                    fh.write("file '{}'\n".format(
                        os.path.join(tmpdir, "clip{}M.mkv".format(i))))
        # Done, join all of the above
        print("Making concatenated video: {}".format(json_in["outputfile"]))
        with open(os.path.join(tmpdir, "join.txt"), "r") as fh:
            for line in fh:
                print(line)
        concat(workingdir=tmpdir, outputfile=json_in["outputfile"])


if __name__ == "__main__":
    import sys

    process_highlights(sys.argv[1])