Blame for inswapper_ov.py

#!/usr/bin/env python3
"""InsightFace `inswapper_128` face swap with OpenVINO - single file example.

Everything runs on OpenVINO (no PyTorch / onnxruntime needed at inference time):

    1. retinaface_10g   -> face detection + 5 point landmarks
    2. arcface_w600k_r50-> 512-d face embedding of the *source* face
    3. inswapper_128    -> the swap itself (128x128 target crop + source embedding)

The OpenVINO IRs are published in the Hugging Face repo `HelloSun/inswapper`
in 16-bit (FP16), 8-bit (W8A8) and 4-bit (u4 weight-only) precisions. They are
downloaded automatically on first use and cached in `~/.cache/inswapper_openvino`.

IMPORTANT: the InsightFace models are published for **non-commercial research
use only** (see the README). Do not use them to impersonate people.

Command line
------------
    python inswapper_ov.py -s source.jpg -t target.jpg -o out.jpg
    python inswapper_ov.py -s source.jpg -t target.jpg -o out.jpg -m int8
    python inswapper_ov.py -s source.jpg -t target.jpg -o out.jpg --all-faces
    python inswapper_ov.py -s source.jpg -t target.jpg -o out.jpg -d GPU

Python API
----------
    from inswapper_ov import FaceSwapperOpenVINO
    swapper = FaceSwapperOpenVINO(precision='int8')
    out = swapper.swap('source.jpg', 'target.jpg')
"""

import argparse
import os
import pathlib
import shutil
import sys
import time

import cv2
import numpy as np
import openvino as ov

HF_REPO = 'HelloSun/inswapper'
CACHE_DIR = pathlib.Path(os.environ.get('INSWAPPER_OV_CACHE', pathlib.Path.home() / '.cache/inswapper_openvino'))

# precision -> (file stem, description)
PRECISIONS = {
    'fp16': ('fp16', '16-bit FP16 weights/activations'),
    'int8': ('int8', '8-bit W8A8 (NNCF PTQ + SmoothQuant)'),
    'int4': ('int4', '4-bit u4 weight-only (NNCF LWC), FP16 activations'),
    'int4-mixed': ('int4_mixed', '4-bit/8-bit mixed precision (75% of the weights INT4)'),
}
AUX_FILES = {'det': 'retinaface', 'rec': 'arcface'}

# ---------------------------------------------------------------------------
# ported from insightface (Apache-2.0): face_align.norm_crop / estimate_norm
# ---------------------------------------------------------------------------
ARCFACE_DST = np.array([[38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366],
                        [41.5493, 92.3655], [70.7299, 92.2041]], dtype=np.float32)


def estimate_norm(lmk, image_size=112):
    """Similarity transform that maps the 5 landmarks onto the arcface template."""
    assert lmk.shape == (5, 2)
    assert image_size % 112 == 0 or image_size % 128 == 0
    if image_size % 112 == 0:
        ratio, diff_x = float(image_size) / 112.0, 0.0
    else:
        ratio, diff_x = float(image_size) / 128.0, 8.0 * (float(image_size) / 128.0)
    dst = ARCFACE_DST * ratio
    dst[:, 0] += diff_x
    src = np.asarray(lmk, dtype=np.float32)
    # least-squares similarity transform (closed form, same as skimage's SimilarityTransform)
    src_mean, dst_mean = src.mean(axis=0), dst.mean(axis=0)
    src_c, dst_c = src - src_mean, dst - dst_mean
    A = dst_c.T @ src_c
    U, S, Vt = np.linalg.svd(A)
    R = U @ Vt
    if np.linalg.det(R) < 0:          # reflection -> flip the smallest singular direction
        U[:, -1] *= -1
        R = U @ Vt
    scale = S.sum() / (src_c ** 2).sum()
    t = dst_mean - scale * R @ src_mean
    return np.concatenate([scale * R, t[:, None]], axis=1).astype(np.float32)


def norm_crop(img, landmark, image_size=128):
    """Crop + warp the face described by `landmark` to `image_size` x `image_size` (BGR)."""
    M = estimate_norm(landmark, image_size)
    return cv2.warpAffine(img, M, (image_size, image_size), borderValue=0.0), M


# ---------------------------------------------------------------------------
# detectors / recognisers
# ---------------------------------------------------------------------------
def _blob(img, mean, std, size, swap_rb=True):
    return cv2.dnn.blobFromImage(img, 1.0 / std, size, (mean, mean, mean), swapRB=swap_rb)


def distance2bbox(points, distance):
    x1 = points[:, 0] - distance[:, 0]
    y1 = points[:, 1] - distance[:, 1]
    x2 = points[:, 0] + distance[:, 2]
    y2 = points[:, 1] + distance[:, 3]
    return np.stack([x1, y1, x2, y2], axis=-1)


def distance2kps(points, distance):
    preds = []
    for i in range(0, distance.shape[1], 2):
        preds.append(points[:, i % 2] + distance[:, i])
        preds.append(points[:, i % 2 + 1] + distance[:, i + 1])
    return np.stack(preds, axis=-1)


class RetinaFace:
    """retinaface_10g detection (single input scale, NMS, 5 landmarks)."""

    FEAT_STRIDES = [8, 16, 32]
    FMC = 3
    NUM_ANCHORS = 2

    def __init__(self, compiled, det_size=(640, 640), det_thresh=0.5, nms_thresh=0.4):
        self.compiled = compiled
        self.det_size = tuple(det_size)
        self.det_thresh = det_thresh
        self.nms_thresh = nms_thresh
        self.center_cache = {}

    def _anchor_centers(self, height, width, stride):
        key = (height, width, stride)
        if key not in self.center_cache:
            centers = np.stack(np.mgrid[:height, :width][::-1], axis=-1).astype(np.float32)
            centers = (centers * stride).reshape(-1, 2)
            if self.NUM_ANCHORS > 1:
                centers = np.stack([centers] * self.NUM_ANCHORS, axis=1).reshape(-1, 2)
            self.center_cache[key] = centers
        return self.center_cache[key]

    def detect(self, img):
        """Return (bboxes[n,5] = x1,y1,x2,y2,score, kpss[n,5,2])."""
        in_w, in_h = self.det_size
        im_ratio = img.shape[0] / img.shape[1]
        model_ratio = in_h / in_w
        if im_ratio > model_ratio:
            new_h, new_w = in_h, int(in_h / im_ratio)
        else:
            new_w, new_h = in_w, int(in_w * im_ratio)
        det_scale = new_h / img.shape[0]
        resized = cv2.resize(img, (new_w, new_h))
        det_img = np.zeros((in_h, in_w, 3), dtype=np.uint8)
        det_img[:new_h, :new_w, :] = resized
        blob = _blob(det_img, 127.5, 128.0, (in_w, in_h))
        outs = self.compiled([blob])

        scores_all, bboxes_all, kpss_all = [], [], []
        for idx, stride in enumerate(self.FEAT_STRIDES):
            scores = outs[idx]
            bbox_preds = outs[idx + self.FMC] * stride
            kps_preds = outs[idx + self.FMC * 2] * stride
            height, width = blob.shape[2] // stride, blob.shape[3] // stride
            centers = self._anchor_centers(height, width, stride)
            pos = np.where(scores >= self.det_thresh)[0]
            if pos.size == 0:
                continue
            scores_all.append(scores[pos])
            bboxes_all.append(distance2bbox(centers, bbox_preds)[pos])
            kps = distance2kps(centers, kps_preds).reshape(-1, 5, 2)
            kpss_all.append(kps[pos])
        if not scores_all:
            return np.empty((0, 5), np.float32), np.empty((0, 5, 2), np.float32)
        scores = np.vstack(scores_all)
        order = scores.ravel().argsort()[::-1]
        bboxes = np.vstack(bboxes_all) / det_scale
        kpss = np.vstack(kpss_all) / det_scale
        pre_det = np.hstack((bboxes, scores)).astype(np.float32)[order, :]
        kpss = kpss[order, :, :]
        keep = self._nms(pre_det)
        return pre_det[keep, :], kpss[keep, :, :]

    def _nms(self, dets):
        x1, y1, x2, y2, scores = dets[:, 0], dets[:, 1], dets[:, 2], dets[:, 3], dets[:, 4]
        areas = (x2 - x1 + 1) * (y2 - y1 + 1)
        order = scores.argsort()[::-1]
        keep = []
        while order.size > 0:
            i = order[0]
            keep.append(i)
            xx1 = np.maximum(x1[i], x1[order[1:]])
            yy1 = np.maximum(y1[i], y1[order[1:]])
            xx2 = np.minimum(x2[i], x2[order[1:]])
            yy2 = np.minimum(y2[i], y2[order[1:]])
            w = np.maximum(0.0, xx2 - xx1 + 1)
            h = np.maximum(0.0, yy2 - yy1 + 1)
            inter = w * h
            ovr = inter / (areas[i] + areas[order[1:]] - inter)
            order = order[np.where(ovr <= self.nms_thresh)[0] + 1]
        return keep


class ArcFace:
    """arcface_w600k_r50 embedding (512-d, L2-normalised)."""

    def __init__(self, compiled):
        self.compiled = compiled

    def embedding(self, img_bgr, landmark):
        aimg, _ = norm_crop(img_bgr, landmark, 112)
        blob = _blob(aimg, 127.5, 127.5, (112, 112))
        feat = self.compiled([blob])[0][0]
        return feat / np.linalg.norm(feat)


class INSwapper:
    """inswapper_128: target crop (128x128 RGB) + source embedding -> swapped face."""

    def __init__(self, compiled, emap):
        self.compiled = compiled
        self.emap = emap

    def project(self, embedding):
        latent = np.dot(embedding.reshape(1, -1), self.emap)
        return (latent / np.linalg.norm(latent)).astype(np.float32)

    def swap_face(self, img_bgr, target_landmark, latent):
        """Swap one face and paste it back; returns a new BGR image."""
        aimg, M = norm_crop(img_bgr, target_landmark, 128)
        blob = _blob(aimg, 0.0, 255.0, (128, 128))          # RGB in [0, 1]
        pred = self.compiled([blob, latent])[0][0]
        img_fake = pred.transpose(1, 2, 0)
        bgr_fake = np.clip(255 * img_fake, 0, 255).astype(np.uint8)[:, :, ::-1]

        fake_diff = bgr_fake.astype(np.float32) - aimg.astype(np.float32)
        fake_diff = np.abs(fake_diff).mean(axis=2)
        fake_diff[:2, :] = 0
        fake_diff[-2:, :] = 0
        fake_diff[:, :2] = 0
        fake_diff[:, -2:] = 0
        IM = cv2.invertAffineTransform(M)
        size = (img_bgr.shape[1], img_bgr.shape[0])
        img_white = np.full(aimg.shape[:2], 255, dtype=np.float32)
        bgr_fake = cv2.warpAffine(bgr_fake, IM, size, borderValue=0.0)
        img_white = cv2.warpAffine(img_white, IM, size, borderValue=0.0)
        fake_diff = cv2.warpAffine(fake_diff, IM, size, borderValue=0.0)
        img_white[img_white > 20] = 255
        fake_diff[fake_diff < 10] = 0
        fake_diff[fake_diff >= 10] = 255
        mask_h_inds, mask_w_inds = np.where(img_white == 255)
        mask_h = np.max(mask_h_inds) - np.min(mask_h_inds)
        mask_w = np.max(mask_w_inds) - np.min(mask_w_inds)
        mask_size = int(np.sqrt(mask_h * mask_w))
        k = max(mask_size // 10, 10)
        img_mask = cv2.erode(img_white, np.ones((k, k), np.uint8), iterations=1)
        fake_diff = cv2.dilate(fake_diff, np.ones((2, 2), np.uint8), iterations=1)
        k = max(mask_size // 20, 5)
        img_mask = cv2.GaussianBlur(img_mask, (2 * k + 1, 2 * k + 1), 0)
        fake_diff = cv2.GaussianBlur(fake_diff, (11, 11), 0)
        img_mask = np.reshape(img_mask / 255., [img_mask.shape[0], img_mask.shape[1], 1])
        merged = img_mask * bgr_fake + (1 - img_mask) * img_bgr.astype(np.float32)
        return merged.astype(np.uint8)


# ---------------------------------------------------------------------------
def _download(rel_path):
    """Fetch a model file from the Hub into the local cache (xml + bin together)."""
    from huggingface_hub import hf_hub_download
    paths = [rel_path]
    if rel_path.endswith('.xml'):            # the IR is split into .xml + .bin
        paths.append(rel_path[:-4] + '.bin')
    CACHE_DIR.mkdir(parents=True, exist_ok=True)
    for rel in paths:
        dst = CACHE_DIR / pathlib.Path(rel).name
        if dst.exists() and dst.stat().st_size > 0:
            continue
        print(f'[inswapper_ov] downloading {rel} from {HF_REPO} ...')
        cached = hf_hub_download(repo_id=HF_REPO, filename=rel)
        shutil.copy(cached, dst)              # copy (not move): keep the hub cache intact
    return str(CACHE_DIR / pathlib.Path(rel_path).name)


def model_name(kind, precision):
    """kind in {'swap', 'det', 'rec'} -> file stem, e.g. 'inswapper_int4'.

    'int4-mixed' only exists for the swapper; the detector / recogniser fall back to
    their pure INT4 builds.
    """
    if precision not in PRECISIONS:
        raise ValueError(f'unknown precision "{precision}", choose from {list(PRECISIONS)}')
    suffix = PRECISIONS[precision][0]
    if kind == 'swap':
        return f'inswapper_{suffix}'
    return f'{AUX_FILES[kind]}_{"int4" if suffix == "int4_mixed" else suffix}'


def resolve_model(kind, precision, model_dir=None):
    """-> local .xml path (from model_dir, the cache, or the Hugging Face Hub)."""
    name = model_name(kind, precision)
    if model_dir:
        local = pathlib.Path(model_dir) / f'{name}.xml'
        if local.exists() and local.with_suffix('.bin').exists():
            return str(local)
    xml = CACHE_DIR / f'{name}.xml'
    if not xml.exists() or not xml.with_suffix('.bin').exists():
        xml = _download(f'models/{name}.xml')
    return str(xml)


def resolve_emap(model_dir=None):
    if model_dir:
        local = pathlib.Path(model_dir) / 'inswapper_emap.npy'
        if local.exists():
            return np.load(str(local))
    npy = CACHE_DIR / 'inswapper_emap.npy'
    if not npy.exists():
        npy = _download('models/inswapper_emap.npy')
    return np.load(str(npy))


class FaceSwapperOpenVINO:
    """High level API: `swap(source_img, target_img, face_index=0, all_faces=False)`."""

    def __init__(self, precision='fp16', device='CPU', det_size=(640, 640), model_dir=None,
                 num_threads=None, precision_hint='f32'):
        core = ov.Core()
        props = {}
        if num_threads:
            props['INFERENCE_NUM_THREADS'] = num_threads
        if precision_hint:
            props['INFERENCE_PRECISION_HINT'] = precision_hint

        def load(kind):
            path = resolve_model(kind, precision, model_dir)
            return core.compile_model(core.read_model(path), device, props)

        self.detector = RetinaFace(load('det'), det_size=det_size)
        self.recognizer = ArcFace(load('rec'))
        self.swapper = INSwapper(load('swap'), resolve_emap(model_dir))
        self.compiled_swap = self.swapper.compiled
        self.precision = precision
        print(f'[inswapper_ov] precision={PRECISIONS[precision][0]} ({PRECISIONS[precision][1]}) device={device}')

    # -- helpers -------------------------------------------------------------
    @staticmethod
    def _read(path):
        img = cv2.imread(str(path), cv2.IMREAD_COLOR)
        if img is None:
            raise FileNotFoundError(f'cannot read image: {path}')
        return img

    def detect(self, img):
        return self.detector.detect(img)

    def embedding(self, img, face_index=0):
        _, kpss = self.detector.detect(img)
        if len(kpss) == 0:
            raise RuntimeError('no face detected')
        kps = kpss[min(face_index, len(kpss) - 1)]
        return self.recognizer.embedding(img, kps), kps

    # -- main entry point ----------------------------------------------------
    def swap(self, source, target, out=None, face_index=0, all_faces=False, verbose=True):
        """Swap the face of `source` into `target`.

        source/target can be a path or a BGR ndarray.  Returns the BGR result.
        """
        src = self._read(source) if isinstance(source, (str, pathlib.Path)) else source
        tgt = self._read(target) if isinstance(target, (str, pathlib.Path)) else target

        _, src_kpss = self.detector.detect(src)
        if len(src_kpss) == 0:
            raise RuntimeError('no face detected in the source image')
        src_kps = src_kpss[0]
        embedding = self.recognizer.embedding(src, src_kps)
        latent = self.swapper.project(embedding)

        bboxes, kpss = self.detector.detect(tgt)
        if len(kpss) == 0:
            raise RuntimeError('no face detected in the target image')
        idx = range(len(kpss)) if all_faces else [min(face_index, len(kpss) - 1)]
        result = tgt
        for i in idx:
            result = self.swapper.swap_face(result, kpss[i], latent)
        if verbose:
            print(f'[inswapper_ov] detected {len(kpss)} face(s), swapped {len(list(idx))}')
        if out is not None:
            cv2.imwrite(str(out), result)
        return result


def main():
    parser = argparse.ArgumentParser(description='InsightFace inswapper_128 face swap with OpenVINO',
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('-s', '--source', required=True, help='image providing the face')
    parser.add_argument('-t', '--target', required=True, help='image to swap into')
    parser.add_argument('-o', '--output', default='result.jpg')
    parser.add_argument('-m', '--model', default='fp16', choices=list(PRECISIONS),
                        help='16-bit (fp16) / 8-bit (int8) / 4-bit (int4)')
    parser.add_argument('-d', '--device', default='CPU', help='CPU / GPU / NPU / AUTO')
    parser.add_argument('--all-faces', action='store_true', help='swap every detected face')
    parser.add_argument('--face-index', type=int, default=0)
    parser.add_argument('--det-size', type=int, nargs=2, default=[640, 640])
    parser.add_argument('--model-dir', default=None, help='folder with the local IR files')
    parser.add_argument('--num-threads', type=int, default=0)
    args = parser.parse_args()

    t0 = time.time()
    swapper = FaceSwapperOpenVINO(precision=args.model, device=args.device, det_size=tuple(args.det_size),
                                  model_dir=args.model_dir, num_threads=args.num_threads or None)
    swapper.swap(args.source, args.target, out=args.output, face_index=args.face_index,
                 all_faces=args.all_faces)
    print(f'[inswapper_ov] wrote {args.output} in {time.time() - t0:.2f}s')
    return 0


if __name__ == '__main__':
    sys.exit(main())