85 lines
2.8 KiB
Python
85 lines
2.8 KiB
Python
#!/usr/bin/env python3
|
|
"""compare_examples.py - every 12-bit example scored on one footing.
|
|
|
|
~/.venvs/fgc-plots/bin/python compare_examples.py <examples-dir> [<out.json>]
|
|
|
|
Decodes each 12-bit example .jxl from `raw12_bench examples` (needs imagecodecs
|
|
and OpenCV), turns the mosaic kinds into colour with the same edge-aware
|
|
de-Bayer, and compares every result with the lossless 12-bit colour image of
|
|
the same frame (c12_ea_lossless.jxl, exact by construction):
|
|
PSNR on the 12-bit scale, all three channels
|
|
SSIM on luma, 12-bit, 11x11 Gaussian window (standard constants)
|
|
So compression is judged on what the downstream algorithms would receive, not
|
|
against each kind's own reference.
|
|
"""
|
|
import glob
|
|
import json
|
|
import os
|
|
import sys
|
|
|
|
import cv2
|
|
import imagecodecs
|
|
import numpy as np
|
|
|
|
d = sys.argv[1]
|
|
ref_path = os.path.join(d, "c12_ea_lossless.jxl")
|
|
|
|
|
|
def decode(path):
|
|
a = imagecodecs.jpegxl_decode(open(path, "rb").read())
|
|
return a
|
|
|
|
|
|
def to_rgb12(name, a):
|
|
"""Decoded array -> HxWx3 float RGB on the 12-bit scale."""
|
|
if a.ndim == 3:
|
|
return a.astype(np.float64)
|
|
m = a.astype(np.uint16)
|
|
h, w = m.shape
|
|
if "planes" in name: # undo the 2x2 tiling of R, G1 / G2, B
|
|
hh, hw = h // 2, w // 2
|
|
mo = np.empty_like(m)
|
|
mo[0::2, 0::2] = m[:hh, :hw]
|
|
mo[0::2, 1::2] = m[:hh, hw:]
|
|
mo[1::2, 0::2] = m[hh:, :hw]
|
|
mo[1::2, 1::2] = m[hh:, hw:]
|
|
m = mo
|
|
return cv2.cvtColor(m, cv2.COLOR_BayerBG2RGB_EA).astype(np.float64)
|
|
|
|
|
|
def luma(rgb):
|
|
return 0.2126 * rgb[..., 0] + 0.7152 * rgb[..., 1] + 0.0722 * rgb[..., 2]
|
|
|
|
|
|
def ssim(a, b, peak=4095.0):
|
|
c1, c2 = (0.01 * peak) ** 2, (0.03 * peak) ** 2
|
|
g = lambda x: cv2.GaussianBlur(x, (11, 11), 1.5)
|
|
mu_a, mu_b = g(a), g(b)
|
|
va, vb, cov = g(a * a) - mu_a ** 2, g(b * b) - mu_b ** 2, g(a * b) - mu_a * mu_b
|
|
s = ((2 * mu_a * mu_b + c1) * (2 * cov + c2)) / ((mu_a ** 2 + mu_b ** 2 + c1) * (va + vb + c2))
|
|
return float(s.mean())
|
|
|
|
|
|
ref = decode(ref_path)
|
|
if ref.max() > 4095: # decoded to 16-bit range: bring back to 12-bit
|
|
ref = ref / 16.0
|
|
ref = ref.astype(np.float64)
|
|
ref_l = luma(ref)
|
|
res = []
|
|
for path in sorted(glob.glob(os.path.join(d, "*.jxl"))):
|
|
name = os.path.splitext(os.path.basename(path))[0]
|
|
if not any(k in name for k in ("c12", "raw12")):
|
|
continue # 8-bit examples are other captures, not comparable pixel by pixel
|
|
a = decode(path)
|
|
if a.max() > 4095:
|
|
a = a / 16.0
|
|
rgb = to_rgb12(name, a)
|
|
mse = float(np.mean((rgb - ref) ** 2))
|
|
psnr = 99.0 if mse == 0 else 10 * np.log10(4095.0 ** 2 / mse)
|
|
r = dict(name=name, bytes=os.path.getsize(path), psnr_db=round(psnr, 2),
|
|
ssim=round(ssim(luma(rgb), ref_l), 5))
|
|
res.append(r)
|
|
print(json.dumps(r))
|
|
if len(sys.argv) > 2:
|
|
json.dump(res, open(sys.argv[2], "w"), indent=1)
|