| #!/usr/bin/env python | |
| ## Copyright (c) 2021, Alliance for Open Media. All rights reserved | |
| ## | |
| ## This source code is subject to the terms of the BSD 3-Clause Clear License and the | |
| ## Alliance for Open Media Patent License 1.0. If the BSD 3-Clause Clear License was | |
| ## not distributed with this source code in the LICENSE file, you can obtain it | |
| ## at aomedia.org/license/software-license/bsd-3-c-c/. If the Alliance for Open Media Patent | |
| ## License 1.0 was not distributed with this source code in the PATENTS file, you | |
| ## can obtain it at aomedia.org/license/patent-license/. | |
| ## | |
| __author__ = "maggie.sun@intel.com, ryanlei@meta.com" | |
| import re | |
| import openpyxl | |
| import xlsxwriter | |
| import shutil | |
| import Config | |
| from Config import QPs, DnScaleRatio, CTC_ASXLSTemplate, CTC_RegularXLSTemplate, InterpolatePieces, \ | |
| UsePCHIPInterpolation | |
| import Utils | |
| from Utils import ParseCSVFile, plot_rd_curve, Interpolate_Bilinear, Interpolate_PCHIP, convex_hull | |
| import matplotlib.pyplot as plt | |
| from matplotlib.backends.backend_pdf import PdfPages | |
| from CalcBDRate import BD_RATE | |
| from itertools import cycle | |
| qtys = ["psnr_y", "psnr_u", "psnr_v", "overall_psnr", "ssim_y", "ms_ssim_y", | |
| "vmaf", "vmaf_neg", "psnr_hvs","ciede2k", "apsnr_y", "apsnr_u", | |
| "apsnr_v", "overall_apsnr"] | |
| ''' | |
| csv_files = { | |
| "HM_CloseGOP": | |
| { | |
| "LD": "D:\\HEVC-AV1-Study\\HEVC-RA-CloseGop\\analysis\\rdresult\\RDResults_hm_hevc_RA_Preset_0.csv", | |
| }, | |
| "HM_OpenGOP": | |
| { | |
| "LD": "D:\\HEVC-AV1-Study\\HEVC-RA-OpenGop\\analysis\\rdresult\\RDResults_hm_hevc_RA_Preset_0.csv", | |
| }, | |
| "AV1_CloseGOP": | |
| { | |
| "LD": "D:\\HEVC-AV1-Study\\AV1-RA-CloseGop\\analysis\\rdresult\\RDResults_aom_av1_RA_Preset_0.csv", | |
| }, | |
| "AV1_OpenGOP": | |
| { | |
| "LD": "D:\\HEVC-AV1-Study\\AV1-RA-OpenGop\\analysis\\rdresult\\RDResults_aom_av1_RA_Preset_0.csv", | |
| }, | |
| } | |
| ''' | |
| csv_files = { | |
| "v1.0.0": | |
| { | |
| "AI": "F:\\Av2-CTC-v4-ToolOffTest\\v1.0-alt\\analysis\\rdresult\\RDResults_aom_av2_AI_Preset_0.csv", | |
| "LD": "F:\\Av2-CTC-v4-ToolOffTest\\v1.0-alt\\analysis\\rdresult\\RDResults_aom_av2_LD_Preset_0.csv", | |
| "RA": "F:\\Av2-CTC-v4-ToolOffTest\\v1.0-alt\\analysis\\rdresult\\RDResults_aom_av2_RA_Preset_0.csv", | |
| "Still": "F:\\Av2-CTC-v4-ToolOffTest\\v1.0-alt\\analysis\\rdresult\\RDResults_aom_av2_STILL_Preset_0.csv", | |
| "AS": "F:\\Av2-CTC-v4-ToolOffTest\\v1.0-alt\\analysis\\rdresult\\RDResults_aom_av2_AS_Preset_0.csv", | |
| }, | |
| "v4.0.0": | |
| { | |
| "AI": "F:\\Av2-CTC-v4-ToolOffTest\\v4.0\\analysis\\rdresult\\RDResults_aom_av2_AI_Preset_0.csv", | |
| "LD": "F:\\Av2-CTC-v4-ToolOffTest\\v4.0\\analysis\\rdresult\\RDResults_aom_av2_LD_Preset_0.csv", | |
| "RA": "F:\\Av2-CTC-v4-ToolOffTest\\v4.0\\analysis\\rdresult\\RDResults_aom_av2_RA_Preset_0.csv", | |
| "Still": "F:\\Av2-CTC-v4-ToolOffTest\\v4.0\\analysis\\rdresult\\RDResults_aom_av2_STILL_Preset_0.csv", | |
| "AS": "F:\\Av2-CTC-v4-ToolOffTest\\v4.0\\analysis\\rdresult\\RDResults_aom_av2_AS_Preset_0.csv", | |
| }, | |
| } | |
| start_row = { | |
| "AI": 2, | |
| "AS": 2, | |
| "RA": 2, | |
| "Still": 2, | |
| "LD": 50, | |
| } | |
| formats = { | |
| "v1.0.0": ['r', '-', 'o'], | |
| "v4.0.0": ['g', '-', '*'], | |
| } | |
| AS_formats = { | |
| "3840x2160": ['r', '-.', 'o'], | |
| "2560x1440": ['g', '-.', '*'], | |
| "1920x1080": ['b', '-.', '^'], | |
| "1280x720": ['y', '-.', '+'], | |
| "960x540": ['c', '-.', 'x'], | |
| "640x360": ['k', '-.', '<'], | |
| } | |
| anchor = "v1.0.0" | |
| rd_curve_pdf = "rdcurve.pdf" | |
| colors = cycle('bgrycmk') | |
| markers = cycle('o*^+<x') | |
| def WriteSheet(csv_file, sht, start_row): | |
| csv = open(csv_file, 'rt') | |
| row = start_row | |
| for line in csv: | |
| if not line.startswith('TestCfg'): | |
| words = re.split(',', line.strip()) | |
| col = 1 | |
| for word in words: | |
| mycell = sht.cell(row=row, column=col) | |
| if col >= 12 and col <= 30 and word != "": | |
| mycell.value = float(word) | |
| else: | |
| mycell.value = word | |
| col += 1 | |
| row += 1 | |
| csv.close() | |
| def FillXlsFile(): | |
| for tag in csv_files.keys(): | |
| if tag == anchor: | |
| continue | |
| else: | |
| for cfg in csv_files[anchor].keys(): | |
| anchor_sht_name = "Anchor-%s" % cfg | |
| test_sht_name = "Test-%s" % cfg | |
| if cfg == "AS": | |
| xls_template = CTC_ASXLSTemplate | |
| xls_file = "CTC_AS_%s-%s.xlsm" % (anchor, tag) | |
| shutil.copyfile(xls_template, xls_file) | |
| anchor_sht_name = "Anchor" | |
| test_sht_name = "Test" | |
| elif cfg == "AI": | |
| xls_template = CTC_RegularXLSTemplate | |
| xls_file = "CTC_Regular_%s-%s.xlsm" % (anchor, tag) | |
| shutil.copyfile(xls_template, xls_file) | |
| wb = openpyxl.load_workbook(filename=xls_file, read_only=False, keep_vba=True) | |
| anchor_sht = wb[anchor_sht_name] | |
| anchor_csv = csv_files[anchor][cfg] | |
| WriteSheet(anchor_csv, anchor_sht, start_row[cfg]) | |
| test_sht = wb[test_sht_name] | |
| test_csv = csv_files[tag][cfg] | |
| WriteSheet(test_csv, test_sht, start_row[cfg]) | |
| wb.save(xls_file) | |
| def DrawIndividualRDCurve(records, anchor, pdf): | |
| with PdfPages(pdf) as export_pdf: | |
| for cfg in records[anchor].keys(): | |
| videos = records[anchor][cfg].keys() | |
| for video in videos: | |
| if cfg == "AS": | |
| DnScaledRes = [(int(3840/ratio), int(2160/ratio)) for ratio in DnScaleRatio] | |
| Int_RDPoints = {} | |
| # draw individual rd curves | |
| for tag in records.keys(): | |
| Int_RDPoints[tag] = [] | |
| record = records[tag][cfg][video] | |
| plt.figure(figsize=(15, 10)) | |
| plt.suptitle("%s : %s: %s" % (cfg, video, tag)) | |
| br = {}; apsnr = {} | |
| for key in record.keys(): | |
| res = re.split('_', key)[0] | |
| if res not in br.keys(): | |
| br[res] = [] | |
| apsnr[res] = [] | |
| br[res].append(record[key].bitrate) | |
| apsnr[res].append(record[key].overall_apsnr) | |
| for res in br.keys(): | |
| rdpnts = [(brt, qty) for brt, qty in zip(br[res], apsnr[res])] | |
| if UsePCHIPInterpolation: | |
| int_rdpnts = Interpolate_PCHIP(rdpnts, QPs['AS'][:], InterpolatePieces, True) | |
| else: | |
| int_rdpnts = Interpolate_Bilinear(rdpnts, QPs['AS'][:], InterpolatePieces, True) | |
| Int_RDPoints[tag] += int_rdpnts | |
| plot_rd_curve(br[res], apsnr[res], "overall_apsnr", res, "bitrate(Kbps)", | |
| AS_formats[res][0], AS_formats[res][1], AS_formats[res][2]) | |
| plt.legend(loc='lower right') | |
| plt.grid(True) | |
| export_pdf.savefig() | |
| plt.close() | |
| #draw convex hull | |
| plt.figure(figsize=(15, 10)) | |
| plt.suptitle("%s : %s: convex hull" % (cfg, video)) | |
| for tag in records.keys(): | |
| lower, upper = convex_hull(Int_RDPoints[tag]) | |
| br = [h[0] for h in upper] | |
| apsnr = [h[1] for h in upper] | |
| plot_rd_curve(br, apsnr, "overall_apsnr(dB)", tag, "bitrate(kbps)", | |
| formats[tag][0], formats[tag][1], formats[tag][2]) | |
| plt.legend(loc='lower right') | |
| plt.grid(True) | |
| export_pdf.savefig() | |
| plt.close() | |
| else: | |
| # bit rate to quality | |
| plt.figure(figsize=(15, 10)) | |
| plt.suptitle("%s : %s" % (cfg, video)) | |
| for tag in records.keys(): | |
| record = records[tag][cfg][video] | |
| br = [record[key].bitrate for key in record.keys()] | |
| apsnr = [record[key].overall_apsnr for key in record.keys()] | |
| plot_rd_curve(br, apsnr, "overall_apsnr(dB)", tag, "bitrate(kbps)", | |
| formats[tag][0], formats[tag][1], formats[tag][2]) | |
| plt.legend(loc='lower right') | |
| plt.grid(True) | |
| export_pdf.savefig() | |
| plt.close() | |
| def DrawCombinedRDCurve(records): | |
| pdf = "combined_rdcurve.pdf" | |
| with PdfPages(pdf) as export_pdf: | |
| for tag in csv_files.keys(): | |
| for cfg in csv_files[tag].keys(): | |
| videos = records[tag][cfg].keys() | |
| plt.figure(figsize=(30, 30)) | |
| plt.suptitle("%s : %s" % (tag, cfg)) | |
| for video in videos: | |
| short_name = video.split('_')[0] | |
| if cfg == "AS": | |
| Int_RDPoints = [] | |
| record = records[tag][cfg][video] | |
| br = {}; | |
| apsnr = {} | |
| for key in record.keys(): | |
| res = re.split('_', key)[0] | |
| if res not in br.keys(): | |
| br[res] = [] | |
| apsnr[res] = [] | |
| br[res].append(record[key].bitrate) | |
| apsnr[res].append(record[key].overall_apsnr) | |
| for res in br.keys(): | |
| rdpnts = [(brt, qty) for brt, qty in zip(br[res], apsnr[res])] | |
| if UsePCHIPInterpolation: | |
| int_rdpnts = Interpolate_PCHIP(rdpnts, QPs['AS'][:], InterpolatePieces, True) | |
| else: | |
| int_rdpnts = Interpolate_Bilinear(rdpnts, QPs['AS'][:], InterpolatePieces, True) | |
| Int_RDPoints += int_rdpnts | |
| # draw convex hull | |
| lower, upper = convex_hull(Int_RDPoints) | |
| br = [h[0] for h in upper] | |
| apsnr = [h[1] for h in upper] | |
| plot_rd_curve(br, apsnr, "overall_apsnr(dB)", short_name, "bitrate(kbps)", | |
| next(colors), '-', next(markers)) | |
| else: | |
| record = records[tag][cfg][video] | |
| br = [record[key].bitrate for key in record.keys()] | |
| apsnr = [record[key].overall_apsnr for key in record.keys()] | |
| plot_rd_curve(br, apsnr, "overall_apsnr(dB)", short_name, "bitrate(kbps)", | |
| next(colors), '-', next(markers)) | |
| plt.legend(loc='lower right') | |
| plt.grid(True) | |
| export_pdf.savefig() | |
| plt.close() | |
| def DrawCombinedRuntime(records): | |
| pdf = "combined_runtime.pdf" | |
| with PdfPages(pdf) as export_pdf: | |
| for tag in csv_files.keys(): | |
| for cfg in csv_files[tag].keys(): | |
| videos = records[tag][cfg].keys() | |
| plt.figure(figsize=(30, 30)) | |
| plt.suptitle("%s : %s" % (tag, cfg)) | |
| for video in videos: | |
| short_name = video.split('_')[0] | |
| if cfg == "AS": | |
| record = records[tag][cfg][video] | |
| br = {}; | |
| enc_time = {} | |
| for key in record.keys(): | |
| res = re.split('_', key)[0] | |
| if res not in br.keys(): | |
| br[res] = [] | |
| enc_time[res] = [] | |
| br[res].append(record[key].bitrate) | |
| enc_time[res].append(record[key].enc_time) | |
| for res in br.keys(): | |
| plot_rd_curve(br[res], enc_time[res], "enc_time(s)", short_name+'_'+res, "bitrate(kbps)", | |
| next(colors), '-', next(markers)) | |
| else: | |
| record = records[tag][cfg][video] | |
| br = [record[key].bitrate for key in record.keys()] | |
| enc_time = [record[key].enc_time for key in record.keys()] | |
| plot_rd_curve(br, enc_time, "enc_time(s)", short_name, "bitrate(kbps)", | |
| next(colors), '-', next(markers)) | |
| plt.legend(loc='lower right') | |
| plt.grid(True) | |
| export_pdf.savefig() | |
| plt.close() | |
| def GetQty(record, qty): | |
| qtys = [] | |
| for key in record.keys(): | |
| if qty == 'psnr_y': | |
| qtys.append(record[key].psnr_y) | |
| elif qty == 'psnr_u': | |
| qtys.append(record[key].psnr_u) | |
| elif qty == 'psnr_v': | |
| qtys.append(record[key].psnr_v) | |
| elif qty == 'overall_psnr': | |
| qtys.append(record[key].overall_psnr) | |
| elif qty == 'ssim_y': | |
| qtys.append(record[key].ssim_y) | |
| elif qty == 'ms_ssim_y': | |
| qtys.append(record[key].ms_ssim_y) | |
| elif qty == 'vmaf': | |
| qtys.append(record[key].vmaf_y) | |
| elif qty == 'vmaf_neg': | |
| qtys.append(record[key].vmaf_y_neg) | |
| elif qty == 'psnr_hvs': | |
| qtys.append(record[key].psnr_hvs) | |
| elif qty == 'ciede2k': | |
| qtys.append(record[key].ciede2k) | |
| elif qty == 'apsnr_y': | |
| qtys.append(record[key].apsnr_y) | |
| elif qty == 'apsnr_u': | |
| qtys.append(record[key].apsnr_u) | |
| elif qty == 'apsnr_v': | |
| qtys.append(record[key].apsnr_v) | |
| elif qty == 'overall_apsnr': | |
| qtys.append(record[key].overall_apsnr) | |
| else: | |
| assert(0) | |
| return qtys | |
| def CalcBDRate(anchor, test): | |
| anchor_qty = {}; test_qty = {} | |
| br_anchor = []; br_test = [] | |
| for key in anchor.keys(): | |
| br_anchor.append(anchor[key].bitrate) | |
| for key in test.keys(): | |
| br_test.append(test[key].bitrate) | |
| for qty in qtys: | |
| anchor_qty[qty] = GetQty(anchor, qty) | |
| test_qty[qty] = GetQty(test, qty) | |
| bdrate = {}; err = 0 | |
| for qty in qtys: | |
| (err, bdrate[qty]) = BD_RATE(qty, br_anchor, anchor_qty[qty], | |
| br_test, test_qty[qty]) | |
| return bdrate | |
| def CalcFullBDRate(cfg): | |
| bdrate = {}; seq_time = {}; seq_instr = {} | |
| for video in records[anchor][cfg].keys(): | |
| bdrate[video] = {}; seq_time[video] = {} | |
| seq_instr[video] = {} | |
| for mode in records.keys(): | |
| record = records[mode][cfg][video] | |
| if mode not in seq_time[video].keys(): | |
| seq_time[video][mode] = 0 | |
| if mode not in seq_instr[video].keys(): | |
| seq_instr[video][mode] = 0 | |
| for key in record.keys(): | |
| seq_time[video][mode] += record[key].enc_time | |
| seq_instr[video][mode] += record[key].enc_instr | |
| if mode == anchor: | |
| continue | |
| bdrate[video][mode] = CalcBDRate(records[anchor][cfg][video], records[mode][cfg][video]) | |
| return (bdrate, seq_time, seq_instr) | |
| def WriteSummaryXlsFile(bdrate, seq_time, seq_instr, summary): | |
| csv = open(summary+".csv", "wt") | |
| csv.write('Video,mode') | |
| for qty in qtys: | |
| csv.write(',%s'%qty) | |
| csv.write(",EncTime(s),EncInstr\n") | |
| wb = xlsxwriter.Workbook(summary + ".xlsx") | |
| shts = [] | |
| for mode in csv_files.keys(): | |
| total = {}; avg_bdrate = {} | |
| sht = wb.add_worksheet(mode) | |
| shts.append(sht) | |
| row = 0; col = 0 | |
| sht.write(row, 0, 'Video') | |
| for qty in qtys: | |
| col = qtys.index(qty) + 1 | |
| sht.write(row, col, qty) | |
| sht.write(row, col + 1, "EncTime(s)") | |
| sht.write(row, col + 2, "EncInstr") | |
| row = 1 | |
| for video in bdrate.keys(): | |
| csv.write("%s,%s"%(video, mode)) | |
| for qty in qtys: | |
| if mode == anchor: | |
| csv.write(',0.0') | |
| elif not str(bdrate[video][mode][qty]).startswith('Error'): | |
| csv.write(',%f' % bdrate[video][mode][qty]) | |
| else: | |
| csv.write(',%s' % bdrate[video][mode][qty]) | |
| csv.write(",%f,%f\n" % (seq_time[video][mode], seq_instr[video][mode])) | |
| sht.write(row, 0, video) | |
| for qty in qtys: | |
| if not qty in total.keys(): | |
| total[qty] = 0 | |
| avg_bdrate[qty] = 0 | |
| if mode != anchor and not str(bdrate[video][mode][qty]).startswith('Error'): | |
| total[qty] += 1 | |
| avg_bdrate[qty] += bdrate[video][mode][qty] | |
| col = qtys.index(qty) + 1 | |
| if mode == anchor: | |
| sht.write(row, col, 0.00) | |
| elif (mode in bdrate[video].keys()): | |
| sht.write(row, col, bdrate[video][mode][qty]) | |
| sht.write(row, col + 1, seq_time[video][mode]) | |
| sht.write(row, col + 2, seq_instr[video][mode]) | |
| row += 1 | |
| row += 1 | |
| for qty in qtys: | |
| if total[qty] != 0: | |
| avg_bdrate[qty] /= total[qty] | |
| else: | |
| avg_bdrate[qty] = 0.0 | |
| sht.write(row, 0, "Average") | |
| for qty in qtys: | |
| col = qtys.index(qty) + 1 | |
| sht.write(row, col, avg_bdrate[qty]) | |
| row += 1 | |
| wb.close() | |
| csv.close() | |
| ###################################### | |
| # main | |
| ###################################### | |
| if __name__ == "__main__": | |
| records = {} | |
| for tag in csv_files.keys(): | |
| records[tag] = {} | |
| for test_cfg in csv_files[tag].keys(): | |
| records[tag][test_cfg] = ParseCSVFile(csv_files[tag][test_cfg]) | |
| FillXlsFile() | |
| DrawCombinedRDCurve(records) | |
| DrawCombinedRuntime(records) | |
| DrawIndividualRDCurve(records, anchor, rd_curve_pdf) | |
| #Calculate BDRate and collect total time | |
| for test_cfg in csv_files[anchor].keys(): | |
| (bdrate, seq_time, seq_instr) = CalcFullBDRate(test_cfg) | |
| #Write output summary xls file | |
| filename = "Summary-AV1-vs-AV2_v4.0_%s"%test_cfg | |
| WriteSummaryXlsFile(bdrate, seq_time, seq_instr, filename) |