File size: 2,017 Bytes
1004324 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 | #!/usr/bin/env python3
"""Compare two DDColor ONNX files on photos: ab max/mean difference plus side-by-side colourised JPGs.
Usage: compare_ddcolor.py REF.onnx TEST.onnx OUTDIR IMG [IMG ...]"""
import sys, time
from pathlib import Path
import cv2, numpy as np, onnxruntime as ort
def preprocess(img_bgr): # DDColor ColorizationPipeline preprocessing
img = (img_bgr / 255.0).astype(np.float32)
l = cv2.cvtColor(cv2.resize(img, (512, 512)), cv2.COLOR_BGR2Lab)[:, :, :1]
rgb = cv2.cvtColor(np.concatenate((l, 0 * l, 0 * l), -1), cv2.COLOR_LAB2RGB)
return rgb.transpose(2, 0, 1)[None].astype(np.float32)
def colourise(img_bgr, ab):
h, w = img_bgr.shape[:2]
l = cv2.cvtColor((img_bgr / 255.0).astype(np.float32), cv2.COLOR_BGR2Lab)[:, :, :1]
ab = cv2.resize(ab[0].transpose(1, 2, 0), (w, h), interpolation=cv2.INTER_NEAREST)
return (cv2.cvtColor(np.concatenate((l, ab), -1), cv2.COLOR_LAB2BGR) * 255).round().clip(0, 255).astype(np.uint8)
ref_s = ort.InferenceSession(sys.argv[1], providers=["CPUExecutionProvider"])
tst_s = ort.InferenceSession(sys.argv[2], providers=["CPUExecutionProvider"])
out = Path(sys.argv[3]); out.mkdir(parents=True, exist_ok=True)
for p in sys.argv[4:]:
img = cv2.imread(p); x = preprocess(img)
a = ref_s.run(["output"], {"input": x})[0]
t = time.time(); b = tst_s.run(["output"], {"input": x})[0]; dt = time.time() - t
d = np.abs(a - b)
ca, cb = colourise(img, a), colourise(img, b)
pix = np.abs(ca.astype(int) - cb.astype(int))
print(f"{Path(p).name}: ab max {d.max():.3f} mean {d.mean():.4f} p99 {np.percentile(d, 99):.3f} | "
f"8-bit RGB max {pix.max()} mean {pix.mean():.3f} | test run {dt:.2f}s")
stem = Path(p).stem.split()[0]
gap = np.full((ca.shape[0], 16, 3), 255, np.uint8)
side = np.concatenate((ca, gap, cb), 1)
s = 1400 / side.shape[1]
cv2.imwrite(str(out / f"{stem}_fp32_vs_test.jpg"), cv2.resize(side, None, fx=s, fy=s, interpolation=cv2.INTER_AREA), [cv2.IMWRITE_JPEG_QUALITY, 92])
|