#!/usr/bin/env python3 """Build the `samples/video_airglow/` gallery: one public-domain NASA clip processed by the original PyTorch model and by every OpenVINO variant (INT4 / INT8 / FP16). Source: NASA/Goddard "The Secrets behind Earth's Multi-colored Glow" (ICON mission), airglow filmed from the ISS - NASA media is public domain. The clip (`input/nasa_ICON_airglow_640x360_63s-66s.mp4`, 640x360) is the LOW RESOLUTION input; every model upscales it to 2560x1440. Per model the script writes a 4x mp4, records the wall-clock time of every frame and computes PSNR/SSIM against the PyTorch FP32 reference. Frames are streamed (never all kept in RAM): the reference is stored JPEG-compressed and the side-by-side grid is assembled in a second pass by decoding the encoded videos. """ import json import os import pathlib import subprocess 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 # The clip is cut from the NASA video asset (see samples/video_airglow/README.md): # ffmpeg -ss 63 -t 3 -i GSFC_20181022_ICON_m12902_Airglow~small.mp4 -an -c:v libx264 ... SRC = ROOT / 'input' / 'nasa_ICON_airglow_640x360_63s-66s.mp4' OUT = ROOT / 'samples' / 'video_airglow' SS, DUR = '0', '1.5' # clip start (s) / duration (s) FPS = 30 TILE, TILE_PAD = 128, 16 GRID = '20_side_by_side_1920x780.mp4' POSTER = '21_poster_frame.jpg' PANELS = [('00_source_lr.mp4', 'input (LR) bicubic x4'), ('10_pytorch_fp32_4x.mp4', 'PyTorch FP32'), ('11_openvino_fp16_4x.mp4', 'OV FP16'), ('12_openvino_int8_w8a8_4x.mp4', 'OV INT8 W8A8'), ('13_openvino_int4_mixed_4x.mp4', 'OV INT4 mixed'), ('14_openvino_int4_4x.mp4', 'OV INT4 (W4A16)')] 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() class Writer: """Raw BGR frames -> H.264 mp4 through ffmpeg.""" def __init__(self, path, size, fps=FPS, crf=26): self.w, self.h = size cmd = ['ffmpeg', '-hide_banner', '-loglevel', 'error', '-y', '-f', 'rawvideo', '-pix_fmt', 'bgr24', '-s', f'{self.w}x{self.h}', '-r', str(fps), '-i', '-', '-c:v', 'libx264', '-crf', str(crf), '-preset', 'veryfast', '-pix_fmt', 'yuv420p', '-movflags', '+faststart', str(path)] self.proc = subprocess.Popen(cmd, stdin=subprocess.PIPE) def write(self, frame): self.proc.stdin.write(np.ascontiguousarray(frame, dtype=np.uint8).tobytes()) def close(self): self.proc.stdin.close() self.proc.wait() def cut_clip(): OUT.mkdir(parents=True, exist_ok=True) dst = OUT / '00_source_lr.mp4' if not SRC.exists(): raise SystemExit(f'source clip not found: {SRC}') subprocess.run(['ffmpeg', '-hide_banner', '-loglevel', 'error', '-y', '-ss', SS, '-t', DUR, '-i', str(SRC), '-an', '-c:v', 'libx264', '-crf', '20', '-preset', 'veryfast', '-pix_fmt', 'yuv420p', '-movflags', '+faststart', str(dst)], check=True) return dst def build_torch_runner(): """Original PyTorch model behind the same API/tiling as realesrgan_ov.""" 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' ' curl -L -o weights/RealESRGAN_x4plus.pth \\\n' ' 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() class Shim: # mimics ov.CompiledModel: compiled([NCHW float32]) -> [(1, C, H, W)] def __call__(self, inputs): with torch.inference_mode(): return [net(torch.from_numpy(inputs[0])).numpy()] runner = RealESRGAN.__new__(RealESRGAN) # reuse _forward_tiled without compiling a model runner.compiled = Shim() runner.scale = 4 return runner LABEL_H = 30 def build_zoom_poster(frame_idx=22, crop=(150, 1010, 640, 360)): """2x3 poster with a 1:1 zoom on the city-lights strip of every variant.""" x, y, cw, ch = crop colors = [(235, 235, 235), (255, 255, 255), (255, 255, 255), (235, 235, 235), (235, 235, 235), (235, 235, 235)] tiles = [] for panel in range(len(PANELS)): cap = cv2.VideoCapture(str(OUT / PANELS[panel][0])) cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx) ok, img = cap.read() cap.release() if not ok: raise SystemExit(f'cannot read frame {frame_idx} of {PANELS[panel][0]}') if panel == 0: # the LR source: show its bicubic x4 rendering, same geometry as the others img = cv2.resize(img, (img.shape[1] * 4, img.shape[0] * 4), interpolation=cv2.INTER_CUBIC) tiles.append(label(img[y:y + ch, x:x + cw], PANELS[panel][1] + ' (1:1 zoom)', colors[panel])) grid = np.vstack([np.hstack(tiles[:3]), np.hstack(tiles[3:])]) cv2.imwrite(str(OUT / '22_poster_detail_zoom.jpg'), grid, [cv2.IMWRITE_JPEG_QUALITY, 93]) print('zoom poster written', flush=True) def label(img, text, bar_color=(255, 255, 255)): bar = np.full((LABEL_H, img.shape[1], 3), bar_color, np.uint8) cv2.putText(bar, text, (8, 21), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (0, 0, 0), 2, cv2.LINE_AA) return np.vstack([bar, img]) def main(): clip = cut_clip() cap = cv2.VideoCapture(str(clip)) ok, f = cap.read() cap.release() w, h = f.shape[1], f.shape[0] n = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) cap = cv2.VideoCapture(str(clip)) frames = [] while True: ok, f = cap.read() if not ok: break frames.append(f) cap.release() h, w = frames[0].shape[:2] n = len(frames) print(f'clip: {n} frames @ {FPS} fps, {w}x{h} -> 4x = {w * 4}x{h * 4}, tile={TILE}', flush=True) def make_grid(): caps = [cv2.VideoCapture(str(OUT / name)) for name, _ in PANELS] writer = Writer(OUT / GRID, (w * 3, (h + LABEL_H) * 2), crf=24) poster = None i = 0 while True: imgs = [] for cap in caps: ok, img = cap.read() if not ok: imgs = [] break imgs.append(cv2.resize(img, (w, h), interpolation=cv2.INTER_AREA)) if not imgs: break row1 = np.hstack([label(imgs[k], PANELS[k][1]) for k in (0, 1, 2)]) row2 = np.hstack([label(imgs[k], PANELS[k][1]) for k in (3, 4, 5)]) grid = np.vstack([row1, row2]) writer.write(grid) if i == n // 2: poster = grid i += 1 for cap in caps: cap.release() writer.close() cv2.imwrite(str(OUT / POSTER), poster, [cv2.IMWRITE_JPEG_QUALITY, 93]) print(f'grid video: {i} frames', flush=True) if '--grid-only' in sys.argv: # rebuild the comparison grid from the existing videos make_grid() return if '--poster-only' in sys.argv: build_zoom_poster() return results = {} def process(name, xml, precision, out_name, reference_jpeg=None): runner = build_torch_runner() if xml is None else RealESRGAN( model=str(ROOT / xml), precision=precision or 'f32') writer = Writer(OUT / out_name, (w * 4, h * 4)) times, ps, ss, jpegs = [], [], [], [] for i, frame in enumerate(frames): t0 = time.time() out = runner.enhance(frame, tile=TILE, tile_pad=TILE_PAD)[0] times.append(time.time() - t0) writer.write(out) if reference_jpeg is not None: ref = cv2.imdecode(reference_jpeg[i], cv2.IMREAD_COLOR) ps.append(psnr(ref, out)) ss.append(ssim(ref, out)) else: jpegs.append(cv2.imencode('.jpg', out, [cv2.IMWRITE_JPEG_QUALITY, 95])[1]) if (i + 1) % 20 == 0: print(f' {name}: {i + 1}/{n} frames, {np.mean(times):.2f} s/frame', flush=True) writer.close() total = sum(times) row = {'file': out_name, 'seconds_total': round(total, 2), 'ms_per_frame': round(total / n * 1000, 1), 'fps': round(n / total, 2), 'input_mpix_per_s': round(w * h / 1e6 * n / total, 2)} if ps: row['psnr_vs_pytorch'] = round(float(np.mean(ps)), 2) row['ssim_vs_pytorch'] = round(float(np.mean(ss)), 4) results[name] = row print(name, row, flush=True) del runner return jpegs ref_jpeg = process('pytorch_fp32', None, None, PANELS[1][0]) variants = [('openvino_fp16', 'models/realesrgan_x4plus_fp16.xml', 'f32', PANELS[2][0]), ('openvino_int8_w8a8', 'models/realesrgan_x4plus_int8_w8a8.xml', None, PANELS[3][0]), ('openvino_int4_mixed_w4a16', 'models/realesrgan_x4plus_int4_mixed_w4a16.xml', 'f32', PANELS[4][0]), ('openvino_int4_w4a16', 'models/realesrgan_x4plus_int4_w4a16.xml', 'f32', PANELS[5][0])] for name, xml, prec, out_name in variants: process(name, xml, prec, out_name, reference_jpeg=ref_jpeg) make_grid() build_zoom_poster() (OUT / 'metrics.json').write_text(json.dumps( {'clip': {'source': SRC.name, 'frames': n, 'fps': FPS, 'input': f'{w}x{h}', 'output': f'{w * 4}x{h * 4}', 'tile': TILE, 'tile_pad': TILE_PAD, 'seconds': [SS, DUR]}, 'results': results}, indent=2)) print(json.dumps(results, indent=2)) if __name__ == '__main__': main()