autoditor/autoditor.py

import moviepy.editor as mp
import numpy as np
from matplotlib import pyplot as plt
import pandas as pd
import scipy.io.wavfile as wav
from typing import Tuple
from itertools import zip_longest
import argparse
import os
import tempfile


class Moment:
    def __init__(self, start, stop):
        self.start = start
        self.stop = stop
        self.duration = 0
        if stop is not None:
            self.duration = stop - start

    def __str__(self):
        return f"Start {self.start} \t\t Stop {self.stop} \t\t Duration {self.duration}"

    def __repr__(self):
        return f"Start {self.start} \t\t Stop {self.stop} \t\t Duration {self.duration}"


def process_audio(source_audio_path: str) -> Tuple[np.ndarray, int, int]:
    rate, data_raw = wav.read(source_audio_path)
    data_raw = data_raw.astype(np.int32)
    mono = (data_raw[:, 0] + data_raw[:, 1]) / 2
    duration = len(mono) / rate
    return mono, duration, rate


def convert_video_to_audio(
    source_video_path: str, destination_audio_location=None
) -> str:
    tdir = tempfile.gettempdir()
    dest_location = f"{tdir}/{source_video_path}.wav"
    print(f"checking to see if {dest_location} exists")
    if destination_audio_location is not None:
        dest_location = destination_audio_location
    if os.path.isfile(dest_location):
        print(f"{dest_location} exists, using cached")
        return dest_location
    vid = mp.VideoFileClip(source_video_path)
    vid.audio.write_audiofile(dest_location)
    vid.close()
    return dest_location


def get_subclips(source_video_path: str, moments):
    vid = mp.VideoFileClip(source_video_path)
    clips = []

    for m in moments:
        if m.duration > 30:
            clips.append(vid.subclip(m.start, m.stop))
    return clips


def sub_resample(data: np.ndarray, factor: int):
    return data[::factor].copy()


def moving_average(x, w):
    return np.convolve(x, np.ones(w), "valid") / w


def find_highlights(data, threshold, rate, factor):
    chunks = []
    for i in range(len(data) - 1):
        if data[i] < threshold < data[i + 1]:
            chunks.append(i * factor / rate)
    return chunks


def find_moving_average_highlights(short_ma, long_ma, bitrate, resample_factor):
    in_a_clip = False
    timestamps = []
    for t in range(1, len(long_ma)):
        if (
            not in_a_clip
            and (short_ma[t - 1] < long_ma[t - 1])
            and (short_ma[t] > long_ma[t])
        ):
            in_a_clip = True
            timestamps.append(t * resample_factor / bitrate)
        elif (
            in_a_clip
            and (short_ma[t - 1] > long_ma[t - 1])
            and (short_ma[t] < long_ma[t])
        ):
            in_a_clip = False
            timestamps.append(t * resample_factor / bitrate)

    ret_list = []
    raw_moments = list(blockwise(timestamps))
    for rm in raw_moments:
        ret_list.append(Moment(rm[0], rm[1]))

    return ret_list


def blockwise(t, size=2, fillvalue=None):
    it = iter(t)
    return zip_longest(*[it] * size, fillvalue=fillvalue)


def plot_audio(data):
    plt.plot(list(range(len(data))), data)
    plt.show()


def main(vidfilepath, outfile, res_factor, lw, sw, dry_run, minduration, maxduration):
    try:
        videofile = vidfilepath

        audiofile = convert_video_to_audio(videofile)
        data, duration, bitrate = process_audio(audiofile)
        RESAMPLE_FACTOR = res_factor
        subsampled_data = sub_resample(data, RESAMPLE_FACTOR)
        squared_subsample = np.square(subsampled_data)
        LONG_WINDOW = lw
        SHORT_WINDOW = sw

        assert LONG_WINDOW > SHORT_WINDOW

        long_ma = moving_average(squared_subsample, LONG_WINDOW)
        short_ma = moving_average(squared_subsample, SHORT_WINDOW)
        moments = find_moving_average_highlights(
            short_ma, long_ma, bitrate, RESAMPLE_FACTOR
        )
        total_time = 0
        for m in moments:
            if m.duration > minduration and m.duration < maxduration:
                print(
                    f"Start {round(m.start/60, 2)} \t\t Stop {round(m.stop/60, 2)} \t\t Duration {round(m.duration, 2)}"
                )
                total_time = total_time + m.duration

        print(total_time / 60)

        if not dry_run:
            clips = get_subclips(videofile, moments)
            clips
            to_render = mp.concatenate_videoclips(clips)
            to_render.write_videofile(outfile)
        os.remove(audiofile)
    except:
        os.remove(audiofile)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        prog="autoditor", description="autoditor is an automatic video editor."
    )
    parser.add_argument(
        "-v", "--video", required=True, metavar="Video file path", dest="vpath"
    )
    parser.add_argument(
        "-f",
        "--factor",
        default=16000,
        metavar="Subsampling factor",
        dest="factor",
        type=int,
    )
    parser.add_argument(
        "-l",
        "--longwindow",
        default=128,
        metavar="Long moving average time",
        dest="lwindow",
        type=int,
    )
    parser.add_argument(
        "-s",
        "--shortwindow",
        default=64,
        metavar="Short moving average time",
        dest="swindow",
        type=int,
    )
    parser.add_argument("-d", "--dryrun", dest="drun", action="store_true")
    parser.add_argument(
        "-o", "--output", required=True, metavar="Output file location", dest="opath"
    )
    parser.add_argument(
        "-i",
        "--minduration",
        default=30,
        metavar="Minimum clip duration",
        dest="mindur",
        type=int,
    )
    parser.add_argument(
        "-m",
        "--maxduration",
        default=100,
        metavar="Maximum clip duration",
        dest="maxdur",
        type=int,
    )

    args = parser.parse_args()
    # def main(vidfilepath, outfile, res_factor, lw, sw, dry_run):
    main(
        args.vpath,
        args.opath,
        args.factor,
        args.lwindow,
        args.swindow,
        args.drun,
        args.mindur,
        args.maxdur,
    )