parse_psnrs.py

import argparse
import re

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-i')
    args = parser.parse_args()
    file = args.i
    regexpsnr = re.compile(r"[\d]+--> .* (\d[\d]?.\d\d)dB ")
    regexpsnry = re.compile(r"[\d]+--> .* (\d[\d]?.\d\d)dBy ")
    sums_psnr = {}
    sums_psnry = {}
    counts = {}
    with open(file, 'r') as handle:
        for line in handle:
            if '-->' in line:
                if '_easy_' in line:
                    key = 'easy'
                elif '_medium_' in line:
                    key = 'medium'
                elif '_hard_' in line:
                    key = 'hard'
                else: 
                    print('Could not parse line')
                    print(line)
                    print('Continuing')
                psnr = regexpsnr.search(line)
                if psnr is None:
                    print(line)
                sums_psnr[key] = sums_psnr.get(key, 0) + float(psnr.group(1))
                psnry = regexpsnry.search(line).group(1)
                sums_psnry[key] = sums_psnry.get(key, 0) + float(psnry)
                counts[key] = counts.get(key, 0) + 1
    for key in ['easy', 'medium', 'hard']:
        print(key, 'psnr:', str(round(sums_psnr[key] / counts[key],2)) + 'dB', 'psnr y:', str(round(sums_psnry[key] / counts[key],2)) + 'dBy')

if __name__ == '__main__':
    main()