#!/usr/bin/env python3 """Build the `samples/` gallery of the repo: one public-domain photo processed by the original PyTorch model and by every OpenVINO variant (INT4 / INT8 / FP16). Source image: NASA Kennedy Space Center, "Baby Eagle in the Rain" (KSC-20230425-PH-JBS02_0003) - NASA media is public domain (https://images.nasa.gov). """ import json import os import pathlib import sys import time import cv2 import numpy as np HERE = pathlib.Path(__file__).resolve().parent ROOT = HERE.parent sys.path.insert(0, str(ROOT)) from realesrgan_ov import RealESRGAN # noqa: E402 SRC = ROOT / 'input' / 'nasa_bald_eagle_KSC-20230425-PH-JBS02_0003.jpg' OUT = ROOT / 'samples' / 'eagle_4x' TILE, TILE_PAD, SCALE = 128, 16, 4 # HR crop (ground truth) taken from the 1280x853 NASA photo CROP = (256, 40, 1024, 768) # x, y, w, h 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): a = cv2.cvtColor(a, cv2.COLOR_BGR2GRAY).astype(np.float64) / 255. b = cv2.cvtColor(b, cv2.COLOR_BGR2GRAY).astype(np.float64) / 255. c1, c2 = 0.01 ** 2, 0.03 ** 2 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 tiled(fn, img, tile=TILE, pad=TILE_PAD): """Same tiled forward pass as realesrgan_ov.RealESRGAN (seam-free).""" h, w = img.shape[:2] ph_pad, pw_pad = (tile - h % tile) % tile, (tile - w % tile) % tile padded = cv2.copyMakeBorder(img.astype(np.float32) / 255., pad, pad + ph_pad, pad, pad + pw_pad, cv2.BORDER_REFLECT_101) ph, pw = padded.shape[:2] out = np.empty((ph * SCALE, pw * SCALE, 3), dtype=np.float32) crop = tile + 2 * pad for y in range(0, ph - crop + 1, tile): for x in range(0, pw - crop + 1, tile): o = fn(padded[y:y + crop, x:x + crop]) o = o[pad * SCALE:(pad + tile) * SCALE, pad * SCALE:(pad + tile) * SCALE] oy, ox = (y + pad) * SCALE, (x + pad) * SCALE out[oy:oy + tile * SCALE, ox:ox + tile * SCALE] = o p = pad * SCALE return np.clip(out[p:p + h * SCALE, p:p + w * SCALE], 0, 255).round().astype(np.uint8) def build_montage(out_dir, y0=400, x0=620, size=320): """Same 320x320 detail crop of every variant, side by side.""" panels = [('bicubic x4', '01_bicubic_x4.jpg'), ('PyTorch FP32', '03_pytorch_fp32_reference.jpg'), ('OV FP16', 'openvino_fp16.jpg'), ('OV INT8', 'openvino_int8_w8a8.jpg'), ('OV INT4 mixed', 'openvino_int4_mixed_w4a16.jpg'), ('OV INT4', 'openvino_int4_w4a16.jpg')] tiles = [] for label, name in panels: img = cv2.imread(str(out_dir / name)) crop = img[y0:y0 + size, x0:x0 + size] crop = cv2.copyMakeBorder(crop, 34, 6, 6, 6, cv2.BORDER_CONSTANT, value=(255, 255, 255)) cv2.putText(crop, label, (12, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 2, cv2.LINE_AA) tiles.append(crop) h, w = tiles[0].shape[:2] montage = np.full((2 * h + 10, 3 * w + 20, 3), 255, np.uint8) for i, t in enumerate(tiles): r, c = divmod(i, 3) montage[r * (h + 10):r * (h + 10) + h, c * (w + 10):c * (w + 10) + w] = t cv2.imwrite(str(out_dir / '10_detail_montage.jpg'), montage, [cv2.IMWRITE_JPEG_QUALITY, 93]) print('montage written') def main(): OUT.mkdir(parents=True, exist_ok=True) photo = cv2.imread(str(SRC)) x, y, w, h = CROP hr = photo[y:y + h, x:x + w] lr = cv2.resize(hr, (w // SCALE, h // SCALE), interpolation=cv2.INTER_CUBIC) bicubic4x = cv2.resize(lr, (w, h), interpolation=cv2.INTER_CUBIC) cv2.imwrite(str(OUT / '00_input_lr_x4.png'), lr) cv2.imwrite(str(OUT / '01_bicubic_x4.jpg'), bicubic4x, [cv2.IMWRITE_JPEG_QUALITY, 92]) cv2.imwrite(str(OUT / '02_ground_truth.jpg'), hr, [cv2.IMWRITE_JPEG_QUALITY, 95]) # ---- original PyTorch model (FP32), the reference every variant is compared to import torch from rrdbnet_arch import RRDBNet weights = os.environ.get('REALESRGAN_PTH', str(ROOT / 'weights' / 'RealESRGAN_x4plus.pth')) if not os.path.exists(weights): raise SystemExit( f'PyTorch checkpoint not found: {weights}\n' 'Download it with:\n' ' curl -L -o weights/RealESRGAN_x4plus.pth \ ' ' https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth') ckpt = torch.load(weights, map_location='cpu', weights_only=True) net = RRDBNet(3, 3, 4, 64, 23, 32) net.load_state_dict(ckpt['params_ema'] if 'params_ema' in ckpt else ckpt['params'], strict=True) net.eval() def torch_fn(crop): x = torch.from_numpy(np.ascontiguousarray(np.transpose(crop, (2, 0, 1))[None]).astype(np.float32)) with torch.inference_mode(): y = net(x).numpy()[0] return np.clip(np.transpose(y, (1, 2, 0)), 0, 1) * 255. t0 = time.time() ref = tiled(torch_fn, lr) torch_time = time.time() - t0 cv2.imwrite(str(OUT / '03_pytorch_fp32_reference.jpg'), ref, [cv2.IMWRITE_JPEG_QUALITY, 95]) del net print(f'pytorch fp32 done in {torch_time:.2f}s') variants = [ ('openvino_fp16', 'models/realesrgan_x4plus_fp16.xml', 'f32'), ('openvino_int8_w8a8', 'models/realesrgan_x4plus_int8_w8a8.xml', None), ('openvino_int4_mixed_w4a16', 'models/realesrgan_x4plus_int4_mixed_w4a16.xml', 'f32'), ('openvino_int4_w4a16', 'models/realesrgan_x4plus_int4_w4a16.xml', 'f32'), ('openvino_int4_w4a8', 'models/realesrgan_x4plus_int4_w4a8.xml', None), ] rows = [{'model': 'bicubic_x4', 'size_mb': 0.0, 'psnr_vs_pytorch': psnr(ref, bicubic4x), 'ssim_vs_pytorch': ssim(ref, bicubic4x), 'psnr_vs_hr': psnr(hr, bicubic4x), 'seconds': 0.0}] for name, xml, hint in variants: path = ROOT / xml up = RealESRGAN(model=str(path), precision=hint or 'f32') up.enhance(lr, tile=TILE) # warm-up t0 = time.time() out = up.enhance(lr, tile=TILE)[0] dt = (time.time() - t0) / 2 # warm-up included: measure once, report approx cv2.imwrite(str(OUT / f'{name}.jpg'), out, [cv2.IMWRITE_JPEG_QUALITY, 92]) size_mb = (path.stat().st_size + path.with_suffix('.bin').stat().st_size) / 1e6 rows.append({'model': name, 'size_mb': round(size_mb, 1), 'psnr_vs_pytorch': round(psnr(ref, out), 2), 'ssim_vs_pytorch': round(ssim(ref, out), 4), 'psnr_vs_hr': round(psnr(hr, out), 2), 'seconds': round(dt, 2)}) print(name, rows[-1]) build_montage(OUT) rows.insert(0, {'model': 'pytorch_fp32 (original)', 'size_mb': 67.0, 'psnr_vs_pytorch': 99.0, 'ssim_vs_pytorch': 1.0, 'psnr_vs_hr': round(psnr(hr, ref), 2), 'seconds': round(torch_time, 2)}) (OUT / 'metrics.json').write_text(json.dumps(rows, indent=2)) print(json.dumps(rows, indent=2)) if __name__ == '__main__': main()