File size: 4,017 Bytes
316dec9 5a1f4db 316dec9 5a1f4db 316dec9 5a1f4db 316dec9 5a1f4db 316dec9 | 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 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | #!/usr/bin/env python3
"""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(): # not downloaded yet -> fetch it from the Hub
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() |