plots.py (6718B)
1 #!/usr/bin/env python3 2 # Copyright (c) the JPEG XL Project Authors. All rights reserved. 3 # 4 # Use of this source code is governed by a BSD-style 5 # license that can be found in the LICENSE file. 6 7 import csv 8 import sys 9 import math 10 import plotly.graph_objects as go 11 12 _, results, output_dir, *rest = sys.argv 13 OUTPUT = rest[0] if rest else 'svg' 14 # valid values: html, svg, png, webp, jpeg, pdf 15 16 with open(results, 'r') as f: 17 reader = csv.DictReader(f) 18 all_results = list(reader) 19 20 nonmetric_columns = set([ 21 "method", "image", "error", "size", "pixels", "enc_speed", "dec_speed", 22 "bpp", "bppp", "qabpp" 23 ]) 24 25 metrics = set(all_results[0].keys()) - nonmetric_columns 26 27 28 def codec(method): 29 sm = method.split(':') 30 ssm = set(sm) 31 speeds = set([ 32 'kitten', 'falcon', 'wombat', 'cheetah', 'tortoise', 'squirrel', 33 'hare', 'fast' 34 ]) 35 s = speeds.intersection(ssm) 36 if sm[0] == 'custom': 37 return sm[1] 38 if sm[0] == 'jxl' and s: 39 return 'jxl-' + list(s)[0] 40 return sm[0] 41 42 43 data = {(m, img): {c: [] 44 for c in {codec(x['method']) 45 for x in all_results}} 46 for m in metrics for img in {x['image'] 47 for x in all_results}} 48 49 for r in all_results: 50 c = codec(r['method']) 51 img = r['image'] 52 bpp = r['bpp'] 53 for m in metrics: 54 data[(m, img)][c].append((float(bpp), float(r[m]))) 55 56 57 def pos(codec): 58 if 'jxl-dis' in codec: 59 return 6, codec 60 elif 'jxl' in codec: 61 return 7, codec 62 elif 'avif' in codec: 63 return 5, codec 64 elif 'kdu' in codec: 65 return 4, codec 66 elif 'heif' in codec: 67 return 3, codec 68 elif 'fuif' in codec or 'pik' in codec: 69 return 2, codec 70 elif 'jpg' in codec or 'jpeg' in codec or 'web' in codec: 71 return 1, codec 72 else: 73 return 0, codec 74 75 76 def style(codec): 77 configs = { 78 'jxl-cheetah': { 79 'color': '#e41a1c', 80 'dash': '1px, 1px', 81 'width': 2 82 }, 83 'jxl-wombat': { 84 'color': '#e41a1c', 85 'dash': '2px, 2px', 86 'width': 2 87 }, 88 'jxl-squirrel': { 89 'color': '#e41a1c', 90 'dash': '5px, 5px', 91 'width': 2 92 }, 93 'jxl-kitten': { 94 'color': '#e41a1c', 95 'width': 2 96 }, 97 'jxl-dis-cheetah': { 98 'color': '#377eb8', 99 'dash': '1px, 1px', 100 'width': 2 101 }, 102 'jxl-dis-wombat': { 103 'color': '#377eb8', 104 'dash': '2px, 2px', 105 'width': 2 106 }, 107 'jxl-dis-squirrel': { 108 'color': '#377eb8', 109 'dash': '5px, 5px', 110 'width': 2 111 }, 112 'jxl-dis-kitten': { 113 'color': '#377eb8', 114 'width': 2 115 }, 116 'rav1e.avif': { 117 'color': '#4daf4a', 118 'dash': '3px, 3px', 119 'width': 2 120 }, 121 '420.rav1e.avif': { 122 'color': '#4daf4a', 123 'dash': '1px, 1px', 124 'width': 2 125 }, 126 '444.rav1e.avif': { 127 'color': '#4daf4a', 128 'dash': '3px, 3px', 129 'width': 2 130 }, 131 'psnr.420.aom.avif': { 132 'color': '#4daf4a', 133 'dash': '5px, 5px', 134 'width': 2 135 }, 136 'psnr.444.aom.avif': { 137 'color': '#4daf4a', 138 'dash': '7px, 7px', 139 'width': 2 140 }, 141 'ssim.420.aom.avif': { 142 'color': '#4daf4a', 143 'dash': '9px, 9px', 144 'width': 2 145 }, 146 'ssim.444.aom.avif': { 147 'color': '#4daf4a', 148 'width': 2 149 }, 150 'heif': { 151 'color': '#984ea3', 152 'width': 2 153 }, 154 'fuif': { 155 'color': '#ff7f00', 156 'dash': '2px, 2px', 157 'width': 2 158 }, 159 'pik-cfp': { 160 'color': '#ff7f00', 161 'width': 2 162 }, 163 'pik-cfp-fast': { 164 'color': '#ff7f00', 165 'dash': '4px, 4px', 166 'width': 2 167 }, 168 'webp': { 169 'color': '#000000', 170 'width': 2 171 }, 172 'jpeg': { 173 'color': '#a65628', 174 'width': 2 175 }, 176 'xt.jpg': { 177 'color': '#a65628', 178 'width': 2 179 }, 180 'perc1.kdu.j2k': { 181 'color': '#f781bf', 182 'dash': '1px, 1px', 183 'width': 2 184 }, 185 'perc2.kdu.j2k': { 186 'color': '#f781bf', 187 'dash': '3px, 3px', 188 'width': 2 189 }, 190 'perc3.kdu.j2k': { 191 'color': '#f781bf', 192 'dash': '5px, 5px', 193 'width': 2 194 }, 195 'perc4.kdu.j2k': { 196 'color': '#f781bf', 197 'dash': '7px, 7px', 198 'width': 2 199 }, 200 'default.kdu.j2k': { 201 'color': '#f781bf', 202 'width': 2 203 }, 204 } 205 return configs.get(codec, dict()) 206 207 208 visible_by_default = set([ 209 'jxl-kitten', 'ssim.444.aom.avif', 'heif', 'webp', 'jpeg', 'xt.jpg', 210 'default.kdu.j2k' 211 ]) 212 213 column_remap = { 214 'p': '6-Butteraugli', 215 'dist': 'Max-Butteraugli', 216 'psnr': "PSNR-YUV 6/8 Y", 217 'MS-SSIM-Y': '-log10(1 - MS-SSIM-Y)', 218 'puSSIM': '-log10(1 - puSSIM)', 219 'FSIM-Y': '-log10(1 - FSIM-Y)', 220 'FSIM-RGB': '-log10(1 - FSIM-RGB)', 221 'VMAF': '-log10(1 - VMAF / 100)', 222 } 223 224 225 def remap(metric): 226 funs = { 227 'MS-SSIM-Y': lambda x: -math.log10(1 - x), 228 'puSSIM': lambda x: -math.log10(1 - x), 229 'FSIM-Y': lambda x: -math.log10(1 - x), 230 'FSIM-RGB': lambda x: -math.log10(1 - x), 231 'VMAF': lambda x: -math.log10(1 + 1e-8 - x / 100), 232 } 233 return funs.get(metric, lambda x: x) 234 235 236 for (m, img) in data: 237 fname = "%s/%s_%s" % (output_dir, m, img) 238 fig = go.Figure() 239 for method in sorted(data[(m, img)].keys(), key=pos): 240 vals = data[(m, img)][method] 241 zvals = list(zip(*sorted(vals))) 242 if not zvals: 243 continue 244 fig.add_trace( 245 go.Scatter(x=zvals[0], 246 y=[remap(m)(x) for x in zvals[1]], 247 mode='lines', 248 name=method, 249 line=style(method), 250 visible=True 251 if method in visible_by_default else 'legendonly')) 252 fig.update_layout(title=img, 253 xaxis_title='bpp', 254 yaxis_title=column_remap.get(m, m)) 255 fig.update_xaxes(type='log') 256 if OUTPUT == 'html': 257 fig.write_html(fname + '.html', include_plotlyjs='directory') 258 else: 259 fig.write_image(fname + '.' + OUTPUT, scale=4)