File size: 7,332 Bytes
3f5e1ac | 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 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | #!/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() |