rd_average.py

#!/usr/bin/python3
# Copyright 2017-2020 Robert-André Mauchin
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#     and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from this
#     software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
#

import os
import sys
import glob
import numpy as np
import pandas as pd
import six
import pytablewriter
from multiprocessing import Pool


def get_lossless_average(path, reference_format):
    merged_data = {}
    columns = [
        "format",
        "avg_bpp",
        "avg_compression_ratio",
        "avg_space_saving",
        "wavg_encode_time",
        "wavg_decode_time",
    ]
    final_data = pd.DataFrame(columns=columns)
    final_data.set_index("format", drop=False, inplace=True)

    for format in next(os.walk(path))[1]:
        if not glob.glob(path + "/" + format + "/lossless/*.out"):
            print(
                "Lossless results files could not be found for format {}.".format(
                    format
                )
            )
            continue

        rawdata = []
        data_path = path + "/" + format + "/lossless/"

        for f in glob.glob(data_path + "/*.out"):
            rawdata.append(pd.read_csv(f, sep=":"))

        merged_data[format] = pd.concat(rawdata)
        sum_orig_file_size = np.sum(merged_data[format]["orig_file_size"])
        sum_compressed_file_size = np.sum(merged_data[format]["compressed_file_size"])
        sum_pixels = np.sum(merged_data[format]["pixels"])
        avg_bpp = sum_compressed_file_size * 8 / sum_pixels
        avg_compression_ratio = sum_orig_file_size / sum_compressed_file_size
        avg_space_saving = 1 - (1 / avg_compression_ratio)
        wavg_encode_time = np.average(
            merged_data[format]["encode_time"], weights=merged_data[format]["pixels"]
        )
        wavg_decode_time = np.average(
            merged_data[format]["decode_time"], weights=merged_data[format]["pixels"]
        )
        final_data.loc[format] = [
            format,
            avg_bpp,
            avg_compression_ratio,
            avg_space_saving,
            wavg_encode_time,
            wavg_decode_time,
        ]

    final_data = final_data.assign(
        weissman_score=lambda x: x.avg_compression_ratio
        / x.loc[reference_format, "avg_compression_ratio"]
        * np.log(x.loc[reference_format, "wavg_encode_time"] * 1000)
        / np.log(x.wavg_encode_time * 1000)
    )
    final_data.sort_values("weissman_score", ascending=False, inplace=True)
    results_file = path + "/" + os.path.basename(path) + ".lossless.out"

    final_data.to_csv(results_file, sep=":")

    file = open(path + "/" + os.path.basename(path) + ".lossless.md", "w")
    markdown_writer = pytablewriter.MarkdownTableWriter()
    markdown_writer.from_dataframe(final_data)
    markdown_writer.stream = six.StringIO()
    markdown_writer.write_table()
    file.write(markdown_writer.stream.getvalue())
    file.close()

    print("Lossless results file successfully saved to {}.".format(results_file))


def get_lossy_average(args):
    [path, format, reference_format] = args

    if not glob.glob(path + "/" + format + "/lossy/*.out"):
        print("Lossy results files could not be found for format {}.".format(format))
        return

    rawdata = []
    merged_data = []
    columns = [
        "file_name",
        "quality",
        "orig_file_size",
        "compressed_file_size",
        "pixels",
        "bpp",
        "compression_ratio",
        "encode_time",
        "decode_time",
        "ssim_score",
        "msssim_score",
        "ciede2000_score",
        "psnrhvs_score",
        "vmaf_score",
        "butteraugli_score",
        "dssim_score",
        "ssimulacra_score",
    ]
    final_columns = [
        "quality",
        "avg_bpp",
        "avg_compression_ratio",
        "avg_space_saving",
        "wavg_encode_time",
        "wavg_decode_time",
        "wavg_ssim_score",
        "wavg_msssim_score",
        "wavg_ciede2000_score",
        "wavg_psnrhvs_score",
        "wavg_vmaf_score",
        "wavg_butteraugli_score",
        "wavg_dssim_score",
        "wavg_ssimulacra_score",
    ]
    final_data = pd.DataFrame(columns=final_columns)

    data_path = path + "/" + format + "/lossy/"

    for f in glob.glob(data_path + "*.out"):
        rawdata.append(pd.read_csv(f, sep=":"))

    quality_length = len(rawdata[0].index)

    for i in range(quality_length):
        merged_data.insert(i, pd.DataFrame(columns=columns))

        for data in rawdata:
            merged_data[i] = merged_data[i].append(data.iloc[[i]])

        merged_data[i].sort_values("file_name", ascending=True, inplace=True)

        quality = np.mean(merged_data[i]["quality"])
        sum_orig_file_size = np.sum(merged_data[i]["orig_file_size"])
        sum_compressed_file_size = np.sum(merged_data[i]["compressed_file_size"])
        sum_pixels = np.sum(merged_data[i]["pixels"])
        avg_bpp = sum_compressed_file_size * 8 / sum_pixels
        avg_compression_ratio = sum_orig_file_size / sum_compressed_file_size
        avg_space_saving = 1 - (1 / avg_compression_ratio)
        wavg_encode_time = np.average(
            merged_data[i]["encode_time"], weights=merged_data[i]["pixels"]
        )
        wavg_decode_time = np.average(
            merged_data[i]["decode_time"], weights=merged_data[i]["pixels"]
        )
        wavg_ssim_score = np.average(
            merged_data[i]["ssim_score"], weights=merged_data[i]["pixels"]
        )
        wavg_ciede2000_score = np.average(
            merged_data[i]["ciede2000_score"], weights=merged_data[i]["pixels"]
        )
        wavg_msssim_score = np.average(
            merged_data[i]["msssim_score"], weights=merged_data[i]["pixels"]
        )
        wavg_psnrhvs_score = np.average(
            merged_data[i]["psnrhvs_score"], weights=merged_data[i]["pixels"]
        )
        wavg_vmaf_score = np.average(
            merged_data[i]["vmaf_score"], weights=merged_data[i]["pixels"]
        )
        wavg_butteraugli_score = np.average(
            merged_data[i]["butteraugli_score"], weights=merged_data[i]["pixels"]
        )
        wavg_dssim_score = np.average(
            merged_data[i]["dssim_score"], weights=merged_data[i]["pixels"]
        )
        wavg_ssimulacra_score = np.average(
            merged_data[i]["ssimulacra_score"], weights=merged_data[i]["pixels"]
        )

        final_data.loc[i] = [
            quality,
            avg_bpp,
            avg_compression_ratio,
            avg_space_saving,
            wavg_encode_time,
            wavg_decode_time,
            wavg_ssim_score,
            wavg_msssim_score,
            wavg_ciede2000_score,
            wavg_psnrhvs_score,
            wavg_vmaf_score,
            wavg_butteraugli_score,
            wavg_dssim_score,
            wavg_ssimulacra_score,
        ]
    results_file = path + "/" + os.path.basename(path) + "." + format + ".lossy.out"
    final_data.to_csv(results_file, sep=":", index=False)
    print(
        "Lossy results file for format {} successfully saved to {}.".format(
            format, results_file
        )
    )


def main(argv):
    if sys.version_info[0] < 3 and sys.version_info[1] < 5:
        raise Exception("Python 3.5 or a more recent version is required.")

    if len(argv) < 2 or len(argv) > 3:
        print(
            "rd_average.py: Calculate a per format weighted averages of the results files generated by rd_collect.py"
        )
        print("Arg 1: Path to the results of a subset generated by rd_collect.py")
        print('       For ex: rd_average.py "results/subset1"')
        print("Arg 2: Reference format with which to compare other formats.")
        print("       Default to mozjpeg")
        return

    results_folder = os.path.normpath(argv[1])
    available_formats = next(os.walk(results_folder))[1]

    # Check is there is actually results files in the path provided
    if (
        not os.path.isdir(results_folder)
        or not available_formats
        or not glob.glob(results_folder + "/**/*.out", recursive=True)
    ):
        print(
            "Could not find all results file. Please make sure the path provided is correct."
        )
        return

    try:
        reference_format = argv[2]
    except IndexError:
        reference_format = "mozjpeg"

    if (
        reference_format not in available_formats
        or not glob.glob(results_folder + "/" + reference_format + "/lossless/*.out")
        or not glob.glob(results_folder + "/" + reference_format + "/lossy/*.out")
    ):
        print(
            "Could not find reference format results files. Please choose a format among {} or check if the reference format results files are present.".format(
                available_formats
            )
        )
        return

    get_lossless_average(results_folder, reference_format)

    Pool().map(
        get_lossy_average,
        [
            (results_folder, format, reference_format)
            for format in next(os.walk(results_folder))[1]
        ],
    )


if __name__ == "__main__":
    main(sys.argv)