|
- import cv2
- import numpy as np
- from numpy.linalg import norm
-
- from facex import FaceX
- from facex.model_zoo import get_model, compute_sim
- from facex.face_align import norm_crop, norm_crop2
-
- target_image = "/root/WangLiDan.jpg"
- source_image = "/root/yangmi.jpg"
-
- app = FaceX(
- tasks=[
- "detection",
- # "recognition",
- ],
- providers=[
- 'CUDAExecutionProvider',
- 'CPUExecutionProvider',
- ],
- )
-
- app.prepare(ctx_id=0, det_size=(640, 640))
- face_restore = get_model("restoration")
- face_swapper = get_model("inswapper")
- face_recognition = get_model("recognition")
- face_parser = get_model("parsing")
-
- target_image = cv2.imread(target_image)
- assert target_image is not None
- target_face = app.get(target_image)[0]
- target_face_image, M = norm_crop2(target_image, landmark=target_face.kps, image_size=512)
- cv2.imwrite("target_face.jpg", target_face_image)
-
- source_image = cv2.imread(source_image)
- source_face = app.get(source_image)[0]
- source_face_image = norm_crop(source_image, landmark=source_face.kps, image_size=512)
- cv2.imwrite("source_face.jpg", source_face_image)
-
- embedding = face_recognition(source_face_image)
- norm_embedding = embedding / norm(embedding)
-
- swapped_face_image = face_swapper(target_face_image, norm_embedding)
- swapped_face_image = cv2.resize(swapped_face_image, (512, 512), interpolation=cv2.INTER_LANCZOS4)
- swapped_face_image = face_restore(swapped_face_image)
- cv2.imwrite("result1.jpg", swapped_face_image)
-
- parse_mask1 = face_parser(swapped_face_image, color_map=False)
- parse_mask2 = face_parser(target_face_image, color_map=False)
- face_reigon1 = np.zeros_like(parse_mask1).astype(bool)
- face_reigon2 = np.zeros_like(parse_mask2).astype(bool)
-
- face_reigon = {
- 'skin': 1,
- 'left-eyebrow': 2,
- 'right-eyebrow': 3,
- 'left-eye': 4,
- 'right-eye': 5,
- 'eye-glasses': 6,
- 'nose': 10,
- 'mouth': 11,
- 'upper-lip': 12,
- 'lower-lip': 13
- }
-
- for idx in range(19):
- if idx in face_reigon.values():
- face_reigon1[parse_mask1 == idx] = True
- face_reigon2[parse_mask2 == idx] = True
-
- parse_mask = np.zeros(face_reigon1.shape + (3, ))
- parse_mask[face_reigon1 & face_reigon2] = [204, 0, 0]
- parse_mask[~face_reigon1 & face_reigon2] = [0, 204, 0]
- parse_mask[face_reigon1 & ~face_reigon2] = [0, 0, 204]
- # parse_mask[~face_reigon1 & ~face_reigon2] = [255, 51, 153]
- parse_mask = parse_mask.astype(np.uint8)
- print(parse_mask.shape, parse_mask.dtype)
- assert cv2.imwrite("parse_mask.jpg", parse_mask)
-
- face_reigon = face_reigon1 | face_reigon2
- face_reigon = face_reigon[..., None]
- swapped_face_image = np.where(face_reigon, swapped_face_image, target_face_image)
- swapped_result = face_swapper.paste_back(swapped_face_image, source_face_image, target_image, M)
-
- cv2.imwrite("result2.jpg", swapped_result)
-
- embedding = np.stack([face_recognition(i) for i in [target_face_image, source_face_image, swapped_face_image]], axis=0)
- norm_embedding: np.ndarray = embedding / norm(embedding, axis=-1, keepdims=True)
- sim = np.dot(norm_embedding, norm_embedding.T)
- print(sim)
|