HelloSun

Add public-domain eagle sample processed by PyTorch + all OpenVINO variants; fix tile padding offset for non-multiple-of-tile sizes

3f5e1ac verified about 1 month ago

Raw History Blame Contribute Delete

7.33 kB

#!/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()

	#!/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()