129 lines
3.9 KiB
Python
129 lines
3.9 KiB
Python
#!/usr/bin/python3
|
|
import subprocess
|
|
import pandas as pd
|
|
import logging
|
|
|
|
|
|
def run_psnr_check(original, encoded, video_info):
|
|
out = ""
|
|
# bad practice, but idgaf
|
|
# -f rawvideo {video_info}
|
|
options = f"-f rawvideo {video_info} -i {original} -i {encoded} -filter_complex psnr -f null /dev/null"
|
|
with open("ffmpeg-log.txt", "w") as f:
|
|
proc = subprocess.run(["ffmpeg", *options.split()],
|
|
stdout=f, stderr=subprocess.STDOUT, text=True)
|
|
logging.info(f"Return code: {proc.returncode}")
|
|
with open("ffmpeg-log.txt", "r") as f:
|
|
out = f.read()
|
|
return out
|
|
|
|
|
|
def run_ssim_check(original, encoded, video_info):
|
|
# bad practice, but idgaf
|
|
# -f rawvideo {video_info}
|
|
# we don't need additional information with h264 encoded files
|
|
options = f"-f rawvideo {video_info} -i {original} -i {encoded} -filter_complex ssim -f null /dev/null"
|
|
with open("ffmpeg-log.txt", "w") as f:
|
|
proc = subprocess.run(["ffmpeg", *options.split()],
|
|
stdout=f, stderr=subprocess.STDOUT, text=True)
|
|
logging.info(f"Return code: {proc.returncode}")
|
|
with open("ffmpeg-log.txt", "r") as f:
|
|
out = f.read()
|
|
return out
|
|
|
|
|
|
def parse_psnr_output(output):
|
|
for line in output.splitlines():
|
|
if "[Parsed_psnr" in line and "PSNR" in line:
|
|
parts = line.split()
|
|
y = parts[4].split(":")[1]
|
|
u = parts[5].split(":")[1]
|
|
v = parts[6].split(":")[1]
|
|
avg = parts[7].split(":")[1]
|
|
minYUV = parts[8].split(":")[1]
|
|
maxYUV = parts[9].split(":")[1]
|
|
return {
|
|
"Y": y,
|
|
"U": u,
|
|
"V": v,
|
|
"Average": avg,
|
|
"MinYUV": minYUV,
|
|
"MaxYUV": maxYUV
|
|
}
|
|
return {}
|
|
|
|
|
|
def parse_ssim_output(output):
|
|
for line in output.splitlines():
|
|
if "[Parsed_ssim" in line and "SSIM" in line:
|
|
parts = line.split()
|
|
all_value = parts[10].split(":")[1]
|
|
y = parts[4].split(":")[1]
|
|
u = parts[6].split(":")[1]
|
|
v = parts[8].split(":")[1]
|
|
return {
|
|
"Y": y,
|
|
"U": u,
|
|
"V": v,
|
|
"Average": all_value
|
|
}
|
|
return {}
|
|
|
|
|
|
def run_quality_check(original, encoded, option):
|
|
psnr_result = run_psnr_check(original, encoded, option)
|
|
ssim_result = run_ssim_check(original, encoded, option)
|
|
psnr_metrics = parse_psnr_output(psnr_result)
|
|
ssim_metrics = parse_ssim_output(ssim_result)
|
|
logging.info("PSNR Metrics:", psnr_metrics)
|
|
logging.info("SSIM Metrics:", ssim_metrics)
|
|
return psnr_metrics, ssim_metrics
|
|
|
|
|
|
def parse_quality_report(psnr_metrics, ssim_metrics):
|
|
psnrSeries = pd.Series(psnr_metrics)
|
|
ssimSeries = pd.Series(ssim_metrics)
|
|
combined = pd.concat([psnrSeries, ssimSeries], axis=1)
|
|
combined.columns = ["PSNR", "SSIM"]
|
|
combined = combined.fillna(0)
|
|
return combined
|
|
|
|
|
|
if __name__ == "__main__":
|
|
psnr, ssim = run_quality_check(
|
|
"base-x264enc-kpop-test-10.yuv",
|
|
"encoded-x264enc-kpop-test-10.mp4",
|
|
"-pixel_format yuv420p -color_range tv -video_size 1920x1080 -framerate 23.98 "
|
|
)
|
|
|
|
combined = parse_quality_report(
|
|
psnr,
|
|
ssim
|
|
)
|
|
|
|
encoder = "x264enc"
|
|
profile = "main"
|
|
params = "bitrate=5000"
|
|
|
|
columns = pd.MultiIndex.from_tuples(
|
|
[(encoder, profile, params, col) for col in combined.columns]
|
|
)
|
|
|
|
combined.columns = columns
|
|
|
|
main_df = combined
|
|
profile = "baseline"
|
|
|
|
combined2 = parse_quality_report(
|
|
psnr,
|
|
ssim
|
|
)
|
|
columns = pd.MultiIndex.from_tuples(
|
|
[(encoder, profile, params, col) for col in combined2.columns]
|
|
)
|
|
combined2.columns = columns
|
|
main_df = pd.concat([main_df, combined2], axis=1)
|
|
logging.info(main_df)
|
|
|
|
main_df.to_csv("quality_report.csv")
|