| |
| """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 |
|
|
| 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' |
|
|
| |
| 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] |
|
|
| |
| 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) |
| 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 |
|
|
| |
| 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() |
