extend

#!/bin/python3
import argparse
import cv2
import os
from subprocess import Popen

parser = argparse.ArgumentParser(
    description="A script for extending still solid frames in a video."
)

parser.add_argument("-i", "--input", help="the input video name", required=True)
parser.add_argument("-o", "--output", help="the output video name", required=True)
parser.add_argument("-c", "--color", help="the color to extend (defaults to green, i.e. 0,255,0)", default="0,255,0")
parser.add_argument("-t", "--threshold", help="how close the frames have to be to the color (defaults to 30)", default=30)

arguments = parser.parse_args()

video = cv2.VideoCapture(arguments.input)

rd, gd, bd = list(map(int, arguments.color.strip().split(",")))

def range_to_filter(r):
    return f"between(n\,{min(r)}\,{max(r)})"

def is_valid_color(frame):
    r, g, b = frame.mean(axis=0).mean(axis=0)

    for a, b in [[r, rd], [g, gd], [b, bd]]:
        if abs(a - b) > arguments.threshold:
            return False
    return True

print("Filtering frames...")

frame_number = 0
frames_to_filter = []
current_range = []


while(True):
    ret, frame = video.read()

    if ret:
        if not is_valid_color(frame):
            current_range.append(frame_number)
        else:
            if len(current_range) != 0:
                frames_to_filter.append(current_range)
                current_range = []

        frame_number += 1
    else:
        break

if len(current_range) != 0:
    frames_to_filter.append(current_range)

select_command = "+".join([range_to_filter(n) for n in frames_to_filter])

command = [
    "ffmpeg", "-i", arguments.input,
    "-vf", f"select='{select_command}'",
    "-crf", "10",  # re-encoding bad
    "-vsync", "cfr",
    "-c:a", "copy", arguments.output
]

process = Popen(command).communicate()

video.release()
cv2.destroyAllWindows()