Add SSIM and MSE (using PSNR) scores

Sparks_in_Blender_encode.py

@@ -66,6 +66,8 @@ def write_line(
     enctime: int,
     vmafmean: float,
     vmafmin: float,
+    ssim: float,
+    mse: float,
 ):
     """
     Write line to data csv
@@ -79,6 +81,8 @@ def write_line(
     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)
@@ -92,6 +96,8 @@ def write_line(
                 enctime,
                 vmafmean,
                 vmafmin,
+                ssim,
+                mse,
             )
         )
 
@@ -179,25 +185,19 @@ def score_vmaf(outputfile: str, inputfile: str) -> dict[str, float]:
     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
+    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):
@@ -218,7 +218,51 @@ def score_ssim(outputfile: str, inputfile: str):
 
     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__":
@@ -239,6 +283,8 @@ if __name__ == "__main__":
                 "Encode time (s)",
                 "VMAF Score (mean)",
                 "VMAF Score (min)",
+                "SSIM Score",
+                "MSE Score",
             )
         )
 
@@ -287,6 +333,10 @@ if __name__ == "__main__":
 
                 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,
@@ -296,6 +346,8 @@ if __name__ == "__main__":
                     enctime=difftime,
                     vmafmean=vmaf["mean"],
                     vmafmin=vmaf["min"],
+                    ssim=ssim,
+                    mse=mse,
                 )
 
                 os.remove(outputfile)