compute_mean.py

"""Append mean values to results txt file."""
import argparse
import json
import multiprocessing as mp
import sys, os

from collections import defaultdict
from itertools import starmap
from pathlib import Path
from typing import List
from math import nan
import numpy as np

from .codecs import AV1, BPG, HM, JPEG, JPEG2000, TFCI, VTM, Codec, WebP

codecs = [JPEG, WebP, JPEG2000, BPG, TFCI, VTM, HM, AV1]


def append_mean(out_dir:str, codec_name:str, qps:List[int]):
    out_dir_codec = Path(out_dir).joinpath(codec_name)
    for q in qps:
        fp = out_dir_codec.joinpath(f"rate_metrics_rec_q_{q:02d}.txt")  
        results = np.loadtxt(fp, usecols=range(1,6))
        mean = results.mean(axis=0, where=~np.isnan(results))
        
        with open(fp, "r") as f:
            lines = f.readlines()
        
        nameWidth = len(lines[-1].split()[0])
        metricWidth, rateWidth = 22, 30
        fmt_psnr = '{:^' + str(metricWidth) + '.2f}'
        fmt_ssim = '{:^' + str(metricWidth) + '.4f}'
        fmt_rate = '{:<' + str(rateWidth) + '}'
        
        line = "Average".center(nameWidth) + (fmt_psnr * 3).format(*mean[:3]) + \
            fmt_ssim.format(mean[3]) + fmt_rate.format(mean[-1]) + "\n"
        
        with open(fp, "a") as f:
            f.write("-"*len(line)+"\n") # before that print a line of "-"
            f.write(line)


def setup_args(parser):
    parser.add_argument(
        "-c",
        "--codec",
        type=str,
        help="Select codec from [jpeg, webp, jpeg2000, bpg, tfci, vtm, hm, av1]",
    )
    parser.add_argument(
        "-o",
        "--out-dir",
        type=str,
        help="Path to directory where to store rates and metrics for each image",
    )
    parser.add_argument(
        "-q",
        "--qps",
        type=str,
        default="75",
        help="list of quality/quantization parameter seperated by comma (default: %(default)s)",
    )
    return parser

def main(argv):
    parser = argparse.ArgumentParser()
    parser = setup_args(parser)
    args = parser.parse_args(argv)

    codec_name = next(c.__name__ for c in codecs if c.__name__.lower() == args.codec)
    qps = [int(q) for q in args.qps.split(",") if q]
    append_mean(args.out_dir, codec_name, qps)


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