Spaces:
Runtime error
Runtime error
File size: 5,587 Bytes
83d8d3c |
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 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 |
import argparse
import os
import cv2
import kornia
import numpy as np
import torch
from loguru import logger
from benchmark.face_pipeline import alignFace
from benchmark.face_pipeline import FaceDetector
from benchmark.face_pipeline import inverse_transform_batch
from benchmark.face_pipeline import SoftErosion
from configs.train_config import TrainConfig
from models.model import HifiFace
class ImageSwap:
def __init__(self, cfg):
self.source_face = cfg.source_face
self.target_face = cfg.target_face
self.device = cfg.device
self.facedetector = FaceDetector(cfg.face_detector_weights, device=self.device)
self.alignface = alignFace()
self.work_dir = cfg.work_dir
opt = TrainConfig()
opt.use_ddp = False
checkpoint = (cfg.model_path, cfg.model_idx)
self.model = HifiFace(
opt.identity_extractor_config, is_training=False, device=self.device, load_checkpoint=checkpoint
)
self.model.eval()
os.makedirs(self.work_dir, exist_ok=True)
# model-idx_swapped_src-image-name_target-face-name.jpg
swapped_image_name = (
str(cfg.model_idx)
+ "_"
+ "swapped"
+ "_"
+ os.path.basename(self.source_face).split(".")[0]
+ "_"
+ os.path.basename(self.target_face).split(".")[0]
+ ".jpg"
)
self.swapped_image = os.path.join(self.work_dir, swapped_image_name)
self.smooth_mask = SoftErosion(kernel_size=7, threshold=0.9, iterations=7).to(self.device)
def _geometry_transfrom_warp_affine(self, swapped_image, inv_att_transforms, frame_size, square_mask):
swapped_image = kornia.geometry.transform.warp_affine(
swapped_image,
inv_att_transforms,
frame_size,
mode="bilinear",
padding_mode="border",
align_corners=True,
fill_value=torch.zeros(3),
)
square_mask = kornia.geometry.transform.warp_affine(
square_mask,
inv_att_transforms,
frame_size,
mode="bilinear",
padding_mode="zeros",
align_corners=True,
fill_value=torch.zeros(3),
)
return swapped_image, square_mask
def detect_and_align(self, image):
detection = self.facedetector(image)
if detection.score is None:
self.kps_window = []
return None, None
max_score_ind = np.argmax(detection.score, axis=0)
kps = detection.key_points[max_score_ind]
align_img, warp_mat = self.alignface.align_face(image, kps, 256)
align_img = cv2.resize(align_img, (256, 256))
align_img = align_img.transpose(2, 0, 1)
align_img = torch.from_numpy(align_img).unsqueeze(0).to(self.device).float()
align_img = align_img / 255.0
return align_img, warp_mat
def inference(self):
src = cv2.cvtColor(cv2.imread(self.source_face), cv2.COLOR_BGR2RGB)
src, _ = self.detect_and_align(src)
if src is None:
print("no face in src_img")
return
target = cv2.cvtColor(cv2.imread(self.target_face), cv2.COLOR_BGR2RGB)
align_target, warp_mat = self.detect_and_align(target)
if align_target is None:
print("no face in target_img")
return
logger.info("start swapping")
frame_size = (target.shape[0], target.shape[1])
with torch.no_grad():
swapped_face, m_r = self.model.forward(src, align_target)
swapped_face = torch.clamp(swapped_face, 0, 1)
smooth_face_mask, _ = self.smooth_mask(m_r)
warp_mat = torch.from_numpy(warp_mat).float().unsqueeze(0)
inverse_warp_mat = inverse_transform_batch(warp_mat, device=self.device)
swapped_face, smooth_face_mask = self._geometry_transfrom_warp_affine(
swapped_face, inverse_warp_mat, frame_size, smooth_face_mask
)
target = torch.from_numpy(target.transpose(2, 0, 1)).unsqueeze(0).to(self.device).float() / 255.0
result_face = (1 - smooth_face_mask) * target + smooth_face_mask * swapped_face
result_face = torch.clamp(result_face * 255.0, 0.0, 255.0, out=None).type(dtype=torch.uint8)
result_face = result_face.detach().cpu().numpy()
img = result_face.transpose(0, 2, 3, 1)[0]
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
cv2.imwrite(self.swapped_image, img)
class ConfigPath:
source_face = ""
target_face = ""
work_dir = ""
face_detector_weights = "/mnt/c/yangguo/useful_ckpt/face_detector/face_detector_scrfd_10g_bnkps.onnx"
model_path = ""
model_idx = 80000
device = "cuda"
def main():
cfg = ConfigPath()
parser = argparse.ArgumentParser(
prog="benchmark", description="What the program does", epilog="Text at the bottom of help"
)
parser.add_argument("-m", "--model_path")
parser.add_argument("-i", "--model_idx")
parser.add_argument("-s", "--source_face")
parser.add_argument("-t", "--target_face")
parser.add_argument("-w", "--work_dir")
parser.add_argument("-d", "--device", default="cuda")
args = parser.parse_args()
cfg.source_face = args.source_face
cfg.target_face = args.target_face
cfg.model_path = args.model_path
cfg.model_idx = int(args.model_idx)
cfg.work_dir = args.work_dir
cfg.device = args.device
infer = ImageSwap(cfg)
infer.inference()
if __name__ == "__main__":
main()
|