| |
| """Benchmark the OpenVINO Real-ESRGAN variants (quality + speed), no PyTorch needed. |
| |
| Every model is compared against the FP16 model, which reproduces the original |
| PyTorch `RealESRGAN_x4plus` result to ~48 dB. |
| |
| python benchmark_models.py # default set, synthetic input |
| python benchmark_models.py --images a.png b.png # real images |
| python benchmark_models.py --models int4 --models int8 |
| """ |
|
|
| import argparse |
| import os |
| import pathlib |
| import time |
|
|
| import cv2 |
| import numpy as np |
|
|
| from realesrgan_ov import MODELS, RealESRGAN, _download_model |
|
|
| HERE = pathlib.Path(__file__).resolve().parent |
|
|
|
|
| def psnr(a, b, peak=255.0): |
| mse = np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2) |
| return 99.0 if mse == 0 else 10 * np.log10(peak ** 2 / mse) |
|
|
|
|
| def ssim(a, b): |
| c1, c2 = 0.01 ** 2, 0.03 ** 2 |
| a = cv2.cvtColor(a, cv2.COLOR_BGR2GRAY).astype(np.float64) / 255. |
| b = cv2.cvtColor(b, cv2.COLOR_BGR2GRAY).astype(np.float64) / 255. |
| mu_a, mu_b = cv2.GaussianBlur(a, (11, 11), 1.5), cv2.GaussianBlur(b, (11, 11), 1.5) |
| sa = cv2.GaussianBlur(a * a, (11, 11), 1.5) - mu_a ** 2 |
| sb = cv2.GaussianBlur(b * b, (11, 11), 1.5) - mu_b ** 2 |
| sab = cv2.GaussianBlur(a * b, (11, 11), 1.5) - mu_a * mu_b |
| return (((2 * mu_a * mu_b + c1) * (2 * sab + c2)) / ((mu_a ** 2 + mu_b ** 2 + c1) * (sa + sb + c2))).mean() |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument('--models', nargs='+', default=['fp16', 'int8', 'int4', 'int4-mixed', 'int4-w4a8']) |
| parser.add_argument('--reference', default='fp16') |
| parser.add_argument('--images', nargs='*', default=[]) |
| parser.add_argument('--tile', type=int, default=128) |
| parser.add_argument('--tile-pad', type=int, default=16) |
| parser.add_argument('--runs', type=int, default=2) |
| parser.add_argument('--device', default='CPU') |
| parser.add_argument('--size', type=int, nargs=2, default=[512, 256], help='synthetic input WxH') |
| parser.add_argument('--token', default=os.environ.get('HF_TOKEN'), help='Hugging Face token') |
| args = parser.parse_args() |
|
|
| if args.images: |
| inputs = [cv2.imread(p) for p in args.images] |
| inputs = [im for im in inputs if im is not None] |
| else: |
| rng = np.random.default_rng(0) |
| yy, xx = np.mgrid[0:args.size[1], 0:args.size[0]] |
| base = (127 + 100 * np.sin(xx / 7.0) * np.cos(yy / 11.0))[..., None] |
| inputs = [np.clip(base + rng.normal(0, 25, base.shape), 0, 255).astype(np.uint8).repeat(3, axis=2)] |
|
|
| def load(key): |
| path = HERE / 'models' / f'{MODELS[key][0]}.xml' |
| if not path.exists(): |
| path = pathlib.Path(_download_model(key, args.token)) |
| return RealESRGAN(model=str(path), device=args.device) |
|
|
| ref_up = load(args.reference) if args.reference in args.models else load(args.reference) |
| refs = [ref_up.enhance(im, tile=args.tile, tile_pad=args.tile_pad)[0] for im in inputs] |
|
|
| size_mb = {} |
| for key in args.models: |
| up = load(key) |
| xml = pathlib.Path(up.model_path) |
| size_mb[key] = (xml.stat().st_size + xml.with_suffix('.bin').stat().st_size) / 1e6 |
| ps, ss, ts, mp = [], [], [], [] |
| for im, ref in zip(inputs, refs): |
| up.enhance(im, tile=args.tile, tile_pad=args.tile_pad) |
| best = min([_time(up, im, args) for _ in range(args.runs)]) |
| out = up.enhance(im, tile=args.tile, tile_pad=args.tile_pad)[0] |
| ts.append(best) |
| mp.append(im.shape[0] * im.shape[1] / best / 1e6) |
| ps.append(psnr(ref, out)) |
| ss.append(ssim(ref, out)) |
| print(f'{key:12s} {size_mb[key]:6.1f} MB PSNR {np.mean(ps):6.2f} dB SSIM {np.mean(ss):.4f} ' |
| f'{np.mean(ts):6.3f} s/img {np.mean(mp):5.2f} MPix/s') |
|
|
|
|
| def _time(up, im, args): |
| t0 = time.time() |
| up.enhance(im, tile=args.tile, tile_pad=args.tile_pad) |
| return time.time() - t0 |
|
|
|
|
| if __name__ == '__main__': |
| main() |
