From 394c6d501cac969227d22a524e68ed7bb5132cd5 Mon Sep 17 00:00:00 2001 From: RealStickman Date: Wed, 9 Nov 2022 20:30:29 +0100 Subject: [PATCH] TODO encoding single video file --- encode_single_video.py | 360 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 360 insertions(+) create mode 100755 encode_single_video.py diff --git a/encode_single_video.py b/encode_single_video.py new file mode 100755 index 0000000..e41462b --- /dev/null +++ b/encode_single_video.py @@ -0,0 +1,360 @@ +#!/usr/bin/env python3 + +import os + +import time + +import csv + +import ffmpy + +from multiprocessing import cpu_count + +from collections import OrderedDict + +import json + +from typing import Union, Any + +# encoding options used +encoding: dict[str, Any] = { + "libx264": { + "crf": [10, 14, 16, 18, 20, 22, 25], + "presets": [ + "superfast", + "veryfast", + "faster", + "fast", + "medium", + "slow", + "slower", + "veryslow", + ], + }, + "libx265": { + "crf": [10, 14, 16, 18, 20, 22, 25], + "presets": [ + "superfast", + "veryfast", + "faster", + "fast", + "medium", + "slow", + "slower", + "veryslow", + ], + }, + "libaom-av1": {"crf": [20, 25, 30, 35, 40], "presets": [0, 2, 4, 6]}, + "libsvtav1": {"crf": [20, 25, 30, 35, 40], "presets": [0, 4, 8, 12]}, +} + +# program version +# make tests reproducible by tag +version = "v0.0.1a3" + +# always round timestamp to integer +def now(): + return int(time.time()) + + +def write_line( + codec: str, + crf: int, + preset: Union[str, int], + infile: str, + outfilesize: float, + enctime: int, + vmafmean: float, + vmafmin: float, + ssim: float, + mse: float, +): + """ + Write line to data csv + + Parameters: + codec (str): Codec used + crf (int): CRF used + preset (str/int): Preset used + infile (str): Input file name + outfilesize (float): Size of output file + enctime (int): Time to encode + vmafmean (float): Mean VMAF score + vmafmin (float): Min VMAF score + ssim (float): SSIM Score + mse (float): MSE Score + """ + with open(datafile, "a", newline="") as file: + write = csv.writer(file) + write.writerow( + ( + codec, + crf, + preset, + infile, + outfilesize, + enctime, + vmafmean, + vmafmin, + ssim, + mse, + ) + ) + + +def encode_general( + inputfile: str, outputfile: str, codec: str, crf: int, preset: Union[str, int] +): + """ + General encoding function + + Parameters: + inputfile (str): Path to input file + outputfile (str): Path to output file + codec (str): Codec used + crf (int): CRF value + preset (str/int): Choosen preset + """ + ff = ffmpy.FFmpeg( + inputs={inputfile: None}, + outputs={ + outputfile: "-c:v {videocodec} -crf {crf} -preset {preset} -g 240 -map 0:v:0 ".format( + videocodec=codec, + crf=crf, + preset=preset, + ) + }, + ) + + return ff + + +def encode_libaom(inputfile: str, outputfile: str, crf: int, preset: Union[str, int]): + """ + Encoding with libaom + + Parameters: + inputfile (str): Path to input file + outputfile (str): Path to output file + crf (int): CRF value + preset (str/int): Choosen preset + """ + ff = ffmpy.FFmpeg( + inputs={inputfile: None}, + outputs={ + outputfile: "-c:v libaom-av1 -crf {crf} -b:v 0 -cpu-used {preset} -row-mt 1 -tiles 2x2 -g 240 -map 0:v:0 ".format( + crf=crf, + preset=preset, + ) + }, + ) + + return ff + + +def score_vmaf(outputfile: str, inputfile: str) -> dict[str, float]: + """ + Calculate a file's VMAF score. Higher is better + + Parameters: + outputfile (str): Path to output file + inputfile (str): Path to input file + + Return: + dict[str, float]: VMAF mean and min value + """ + ff = ffmpy.FFmpeg( + inputs=OrderedDict([(outputfile, None), (inputfile, None)]), + outputs={ + "-": "-filter_complex libvmaf=log_fmt=json:n_threads={cputhreads}:log_path=vmaf.json -f null".format( + cputhreads=cpu_count() + ) + }, + ) + + ff.run() + + with open("vmaf.json", "r") as file: + vmafall = json.load(file) + + vmaf: dict[str, float] = { + "mean": vmafall["pooled_metrics"]["vmaf"]["mean"], + "min": vmafall["pooled_metrics"]["vmaf"]["min"], + } + + return vmaf + + +def parse_kv_files(inputfile: str) -> list[dict[str, Any]]: + # create list of dicts. Each dict is one line in the file + lines: list[dict[str, Any]] = [] + # Steps to get mse value + with open(inputfile) as file: + for line in file: + linelist = line.rstrip().split() + # NOTE stripping the last item because ssim has values like this in the location: (26.088579) + # For PSNR, psnr_v:87.80 is affected. (YUV color space). Not really a priority to fix atm + linedict = dict(kv.split(":") for kv in linelist[:-1]) + lines.append(linedict) + + return lines + + +def score_ssim(outputfile: str, inputfile: str): + """ + Calculate a file's SSIM rating. TBD + + Parameters: + outputfile (str): Path to output file + inputfile (str): Path to input file + + Return: + TBD + """ + ff = ffmpy.FFmpeg( + inputs=OrderedDict([(outputfile, None), (inputfile, None)]), + outputs={"-": "-lavfi ssim=stats_file=ssim.log -f null"}, + ) + + ff.run() + + # get average ssim value + ssim: list[dict[str, Any]] = parse_kv_files("ssim.log") + + # add all ssim_avg values together + ssim_tot: float = 0.0 + for i in ssim: + ssim_tot += float(i["All"]) + + # get mse average + ssim_avg: float = ssim_tot / len(ssim) + + return ssim_avg + + +def score_psnr(outputfile: str, inputfile: str) -> float: + """ + Calculate a file's MSE (mean-square error) using PSNR. A lower value is better + + Parameters: + outputfile (str): Path to output file + inputfile (str): Path to input file + + Return: + TBD + """ + ff = ffmpy.FFmpeg( + inputs=OrderedDict([(outputfile, None), (inputfile, None)]), + outputs={"-": "-lavfi psnr=stats_file=psnr.log -f null"}, + # outputs={"-": "-filter_complex psnr=stats_file=psnr.log -f null"}, + ) + + ff.run() + + # get average mse value + psnr: list[dict[str, Any]] = parse_kv_files("psnr.log") + + # add all mse_avg values together + mse_tot: float = 0.0 + for i in psnr: + mse_tot += float(i["mse_avg"]) + + # get mse average + mse_avg: float = mse_tot / len(psnr) + + return mse_avg + + +if __name__ == "__main__": + if not os.path.isdir("encodes"): + os.mkdir("encodes") + + datafile = version + "-data-" + str(now()) + ".csv" + + with open(datafile, "w", newline="") as file: + write = csv.writer(file) + write.writerow( + ( + "Codec", + "CRF", + "Preset", + "Input file", + "Output file size (MiB)", + "Encode time (s)", + "VMAF Score (mean)", + "VMAF Score (min)", + "SSIM Score", + "MSE Score", + ) + ) + + # TODO expose this as variable/argument? + inputfile = "source/Sparks_in_Blender.webm" + + for codec in encoding: + for crf in encoding[codec]["crf"]: + for preset in encoding[codec]["presets"]: + # TODO selection of output location with arguments? + outputfile = os.path.join( + "encodes", + ( + "Sparks_in_Blender-codec_" + + codec + + "-crf_" + + str(crf) + + "-preset_" + + str(preset) + + ".mkv" + ), + ) + + # libaom needs additional options + if codec == "libaom-av1": + ff = encode_libaom( + inputfile=inputfile, + outputfile=outputfile, + crf=crf, + preset=preset, + ) + else: + ff = encode_general( + inputfile=inputfile, + outputfile=outputfile, + codec=codec, + crf=crf, + preset=preset, + ) + + # execute previously defined encoding settings + starttime = now() + ff.run() + endtime = now() + difftime = int(endtime - starttime) + + outputfilesize = os.path.getsize(outputfile) / 1024 / 1024 + + vmaf = score_vmaf(outputfile=outputfile, inputfile=inputfile) + + ssim = score_ssim(outputfile=outputfile, inputfile=inputfile) + + mse = score_psnr(outputfile=outputfile, inputfile=inputfile) + + write_line( + codec=codec, + crf=crf, + preset=preset, + infile="Sparks_in_Blender.webm", + outfilesize=outputfilesize, + enctime=difftime, + vmafmean=vmaf["mean"], + vmafmin=vmaf["min"], + ssim=ssim, + mse=mse, + ) + + # TODO make removal optional/togglable with cli switch + os.remove(outputfile) + + os.remove("vmaf.json") + os.remove("ssim.log") + os.remove("psnr.log")