Spaces:
Running
Running
File size: 2,496 Bytes
6cc375a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
#!/usr/bin/env python3
"""
swap.py — Simple 128px face-swap using INSwapperONNX + InsightFace
Usage:
python swap.py --src <your.jpg> --dst <target.jpg> --out <out.png> --device cpu|gpu
"""
import argparse
import os
import sys
import cv2
def swap_faces(
src_path: str,
dst_path: str,
out_path: str,
device: str,
):
import onnxruntime as ort
from insightface.app import FaceAnalysis
from insightface.model_zoo.inswapper import INSwapper
providers = ["CUDAExecutionProvider"] if device == "gpu" else [
"CPUExecutionProvider"
]
app = FaceAnalysis(name="buffalo_l", providers=providers)
app.prepare(ctx_id=0 if device == "gpu" else -1, det_size=(640, 640))
# path to the ONNX model
model_path = os.path.expanduser(
"~/.insightface/models/inswapper_128.onnx"
)
if not os.path.isfile(model_path):
sys.exit(f"❌ Model file not found: {model_path}")
session = ort.InferenceSession(model_path, providers=providers)
swapper = INSwapper(model_file=model_path, session=session)
src_img = cv2.imread(src_path)
dst_img = cv2.imread(dst_path)
if src_img is None or dst_img is None:
sys.exit("❌ Failed to load one of the images.")
faces_src = app.get(src_img)
faces_dst = app.get(dst_img)
if not faces_src or not faces_dst:
sys.exit("❌ Could not detect a face in one of the images.")
result = swapper.get(
dst_img,
faces_dst[0],
faces_src[0],
paste_back=True,
)
ext = os.path.splitext(out_path)[1].lower()
params = [cv2.IMWRITE_JPEG_QUALITY, 100] if ext in (
".jpg",
".jpeg",
) else []
cv2.imwrite(out_path, result, params)
print(f"✅ Swap complete → {out_path}")
def main():
parser = argparse.ArgumentParser(description="Face-swap 128px CLI")
parser.add_argument(
"--src",
required=True,
help="Path to your source image",
)
parser.add_argument(
"--dst",
required=True,
help="Path to target image",
)
parser.add_argument(
"--out",
default="result.png",
help="Output file path",
)
parser.add_argument(
"--device",
choices=["cpu", "gpu"],
default="cpu",
help="Run on CPU or GPU (needs onnxruntime-gpu)",
)
args = parser.parse_args()
swap_faces(args.src, args.dst, args.out, args.device)
if __name__ == "__main__":
main()
|