Blame for tools/make_video_sample.py

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