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