| |
| """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): |
| 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]) |
|
|
