Add SSIM and MSE (using PSNR) scores

This commit is contained in:
RealStickman 2022-10-24 21:26:45 +02:00
parent 9db7bc6adb
commit 5b3fc22f55

View File

@ -66,6 +66,8 @@ def write_line(
enctime: int, enctime: int,
vmafmean: float, vmafmean: float,
vmafmin: float, vmafmin: float,
ssim: float,
mse: float,
): ):
""" """
Write line to data csv Write line to data csv
@ -79,6 +81,8 @@ def write_line(
enctime (int): Time to encode enctime (int): Time to encode
vmafmean (float): Mean VMAF score vmafmean (float): Mean VMAF score
vmafmin (float): Min VMAF score vmafmin (float): Min VMAF score
ssim (float): SSIM Score
mse (float): MSE Score
""" """
with open(datafile, "a", newline="") as file: with open(datafile, "a", newline="") as file:
write = csv.writer(file) write = csv.writer(file)
@ -92,6 +96,8 @@ def write_line(
enctime, enctime,
vmafmean, vmafmean,
vmafmin, vmafmin,
ssim,
mse,
) )
) )
@ -179,25 +185,19 @@ def score_vmaf(outputfile: str, inputfile: str) -> dict[str, float]:
return vmaf return vmaf
def score_psnr(outputfile: str, inputfile: str): def parse_kv_files(inputfile: str) -> list[dict[str, Any]]:
""" # create list of dicts. Each dict is one line in the file
Calculate a file's MSE (mean-square error) using PSNR. A lower value is better lines: list[dict[str, Any]] = []
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"},
)
ff.run()
# Steps to get mse value # 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): def score_ssim(outputfile: str, inputfile: str):
@ -218,7 +218,51 @@ def score_ssim(outputfile: str, inputfile: str):
ff.run() ff.run()
# Steps to get ssim value # 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 __name__ == "__main__":
@ -239,6 +283,8 @@ if __name__ == "__main__":
"Encode time (s)", "Encode time (s)",
"VMAF Score (mean)", "VMAF Score (mean)",
"VMAF Score (min)", "VMAF Score (min)",
"SSIM Score",
"MSE Score",
) )
) )
@ -287,6 +333,10 @@ if __name__ == "__main__":
vmaf = score_vmaf(outputfile=outputfile, inputfile=inputfile) vmaf = score_vmaf(outputfile=outputfile, inputfile=inputfile)
ssim = score_ssim(outputfile=outputfile, inputfile=inputfile)
mse = score_psnr(outputfile=outputfile, inputfile=inputfile)
write_line( write_line(
codec=codec, codec=codec,
crf=crf, crf=crf,
@ -296,6 +346,8 @@ if __name__ == "__main__":
enctime=difftime, enctime=difftime,
vmafmean=vmaf["mean"], vmafmean=vmaf["mean"],
vmafmin=vmaf["min"], vmafmin=vmaf["min"],
ssim=ssim,
mse=mse,
) )
os.remove(outputfile) os.remove(outputfile)