import cv2

from face_analysis import FaceAnalysis
from models.inswapper import INSwapper

app = FaceAnalysis()

swapper = INSwapper('./inswapper_128.onnx')

source_img = "/home/leonel/Pictures/bun01.jpg"
target_img = "/home/leonel/Pictures/ac6a1e147711139.62c72d282c159.png"

# Load source and target images
src_img = cv2.imread(source_img)
dst_img = cv2.imread(target_img)

# Get faces
src_faces = app.get(src_img)
dst_faces = app.get(dst_img)

# Swap the first face found
if len(src_faces) > 0 and len(dst_faces) > 0:
    result = swapper.get(dst_img, dst_faces[0], src_faces[0], paste_back=True)
    cv2.imwrite("swapped2.jpg", result)