make_sample.py

#!/usr/bin/env python3
"""Build the `samples/` gallery of the repo.

Takes the public-domain NASA astronaut portraits, swaps one face into the other with
every OpenVINO precision and records quality (vs. the ONNX Runtime FP32 reference)
and latency.

    python make_sample.py --onnx-dir onnx --models-dir models --out-dir samples
"""
import argparse
import json
import os
import pathlib
import sys
import time

import cv2
import numpy as np
import onnx
import onnxruntime as ort

HERE = pathlib.Path(__file__).resolve().parent
sys.path.insert(0, os.path.dirname(HERE))
import inswapper_ov as iov  # noqa: E402

PRECISIONS = ['fp16', 'int8', 'int4', 'int4-mixed']


class ORTSession:
    """Mimics ov.CompiledModel so that the very same code path runs on ONNX Runtime."""

    def __init__(self, path):
        self.sess = ort.InferenceSession(path, providers=['CPUExecutionProvider'])
        self.names = [i.name for i in self.sess.get_inputs()]

    def __call__(self, inputs):
        return self.sess.run(None, {n: inputs[i] for i, n in enumerate(self.names)})


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
    ma, mb = cv2.GaussianBlur(a, (11, 11), 1.5), cv2.GaussianBlur(b, (11, 11), 1.5)
    sa = cv2.GaussianBlur(a * a, (11, 11), 1.5) - ma ** 2
    sb = cv2.GaussianBlur(b * b, (11, 11), 1.5) - mb ** 2
    sab = cv2.GaussianBlur(a * b, (11, 11), 1.5) - ma * mb
    return (((2 * ma * mb + c1) * (2 * sab + c2)) / ((ma ** 2 + mb ** 2 + c1) * (sa + sb + c2))).mean()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--onnx-dir', default='onnx')
    ap.add_argument('--models-dir', default='models')
    ap.add_argument('--out-dir', default='samples')
    ap.add_argument('--source', default='input/source_face_nasa_astronaut.jpg')
    ap.add_argument('--target', default='input/target_photo_nasa_astronaut.jpg')
    args = ap.parse_args()

    out = pathlib.Path(args.out_dir)
    out.mkdir(parents=True, exist_ok=True)
    src, tgt = cv2.imread(args.source), cv2.imread(args.target)
    assert src is not None and tgt is not None, 'source/target images not found'

    # ---------------- ONNX Runtime FP32 reference (same code path, ORT backend)
    emap = onnx.numpy_helper.to_array(onnx.load(os.path.join(args.onnx_dir, 'inswapper_128.onnx'))
                                      .graph.initializer[-1]).astype(np.float32)
    ref = iov.FaceSwapperOpenVINO(precision='fp16', model_dir=args.models_dir)
    ref.detector = iov.RetinaFace(ORTSession(os.path.join(args.onnx_dir, 'retinaface_10g.onnx')))
    ref.recognizer = iov.ArcFace(ORTSession(os.path.join(args.onnx_dir, 'arcface_w600k_r50.onnx')))
    ref.swapper = iov.INSwapper(ORTSession(os.path.join(args.onnx_dir, 'inswapper_128.onnx')), emap)
    ref.compiled_swap = ref.swapper.compiled

    t0 = time.time()
    ref_image = ref.swap(src, tgt, verbose=False)
    ref_total = time.time() - t0
    ref_emb, ref_kps = ref.embedding(src)
    _, ref_kpss = ref.detect(tgt)
    ref_crop, _ = iov.norm_crop(tgt, ref_kpss[0], 128)
    ref_blob = cv2.dnn.blobFromImage(ref_crop, 1 / 255.0, (128, 128), (0, 0, 0), swapRB=True)
    ref_face = ref.compiled_swap([ref_blob, ref.swapper.project(ref_emb)])[0][0]

    # crops used in the montage
    src_crop, _ = iov.norm_crop(src, ref.embedding(src)[1], 128)
    tgt_crop, _ = iov.norm_crop(tgt, ref_kpss[0], 128)
    cv2.imwrite(str(out / '00_source_face.jpg'), cv2.resize(src_crop, (256, 256), interpolation=cv2.INTER_CUBIC),
                [cv2.IMWRITE_JPEG_QUALITY, 95])
    cv2.imwrite(str(out / '01_target_original_face.jpg'), cv2.resize(tgt_crop, (256, 256), interpolation=cv2.INTER_CUBIC),
                [cv2.IMWRITE_JPEG_QUALITY, 95])
    cv2.imwrite(str(out / '02_target_original.jpg'), tgt, [cv2.IMWRITE_JPEG_QUALITY, 95])
    cv2.imwrite(str(out / '03_onnxruntime_fp32_reference.jpg'), ref_image, [cv2.IMWRITE_JPEG_QUALITY, 95])

    rows = [{'model': 'onnxruntime fp32 (reference)', 'swapper_mb': round(os.path.getsize(
        os.path.join(args.onnx_dir, 'inswapper_128.onnx')) / 1e6, 1),
        'face_psnr': 99.0, 'face_ssim': 1.0, 'image_psnr': 99.0, 'image_ssim': 1.0,
        'embedding_cos': 1.0, 'swap_seconds': round(ref_total, 3)}]

    ref_face_u8_montage = np.clip(ref_face.transpose(1, 2, 0), 0, 1)
    ref_face_u8_montage = (ref_face_u8_montage * 255).round().astype(np.uint8)[:, :, ::-1]
    tiles = {'source face': src_crop, 'target (before)': tgt_crop,
             'ORT FP32 (reference)': ref_face_u8_montage}
    for i, prec in enumerate(PRECISIONS, start=4):
        sw = iov.FaceSwapperOpenVINO(precision=prec, model_dir=args.models_dir)
        sw.swap(src, tgt, verbose=False)                       # warm-up
        t0 = time.time()
        image = sw.swap(src, tgt, verbose=False)
        swap_time = time.time() - t0
        emb, _ = sw.embedding(src)
        kpss = sw.detect(tgt)[1]
        crop, _ = iov.norm_crop(tgt, kpss[0], 128)
        blob = cv2.dnn.blobFromImage(crop, 1 / 255.0, (128, 128), (0, 0, 0), swapRB=True)
        face = sw.compiled_swap([blob, sw.swapper.project(emb)])[0][0]
        face_u8 = np.clip(face.transpose(1, 2, 0), 0, 1)
        face_u8 = (face_u8 * 255).round().astype(np.uint8)[:, :, ::-1]
        ref_face_u8 = np.clip(ref_face.transpose(1, 2, 0), 0, 1)
        ref_face_u8 = (ref_face_u8 * 255).round().astype(np.uint8)[:, :, ::-1]
        mb = sum(os.path.getsize(os.path.join(args.models_dir, f'{n}{e}'))
                 for n in (f'inswapper_{iov.PRECISIONS[prec][0]}',) for e in ('.xml', '.bin')) / 1e6
        row = {'model': f'openvino {prec}', 'swapper_mb': round(mb, 1),
               'face_psnr': round(psnr(ref_face_u8, face_u8), 2), 'face_ssim': round(ssim(ref_face_u8, face_u8), 4),
               'image_psnr': round(psnr(ref_image, image), 2), 'image_ssim': round(ssim(ref_image, image), 4),
               'embedding_cos': round(float(np.dot(emb, ref_emb) / (np.linalg.norm(emb) * np.linalg.norm(ref_emb))), 5),
               'swap_seconds': round(swap_time, 3)}
        rows.append(row)
        print(row, flush=True)
        cv2.imwrite(str(out / f'{i}0_openvino_{prec.replace("-", "_")}.jpg'), image, [cv2.IMWRITE_JPEG_QUALITY, 95])
        tiles[f'OV {prec}'] = cv2.resize(face_u8, (256, 256), interpolation=cv2.INTER_NEAREST)
        del sw

    # ---------------- montage of the 128x128 swapped faces ----------------
    names = list(tiles)
    panel = 256
    cols, rows_n = 4, 2
    montage = np.full((rows_n * (panel + 34), cols * panel, 3), 255, np.uint8)
    for idx, (name, img) in enumerate(tiles.items()):
        if idx >= rows_n * cols:
            break
        r, c = divmod(idx, cols)
        bar = np.full((34, panel, 3), 255, np.uint8)
        cv2.putText(bar, name, (8, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 0, 0), 2, cv2.LINE_AA)
        tile = cv2.resize(img, (panel, panel), interpolation=cv2.INTER_NEAREST)
        montage[r * (panel + 34):r * (panel + 34) + panel + 34, c * panel:(c + 1) * panel] = np.vstack([bar, tile])
    cv2.imwrite(str(out / '20_face_montage.jpg'), montage, [cv2.IMWRITE_JPEG_QUALITY, 95])

    (out / 'metrics.json').write_text(json.dumps({
        'source': args.source, 'target': args.target,
        'reference': 'onnxruntime 1.20 (FP32) running the same preprocessing/paste-back code',
        'results': rows}, indent=2))
    print(json.dumps(rows, indent=2))


if __name__ == '__main__':
    main()