mirror of
https://github.com/deepfakes/faceswap
synced 2025-06-09 04:36:50 -04:00
63 lines
2.2 KiB
Python
63 lines
2.2 KiB
Python
#!/usr/bin/env python3
|
|
""" Components Mask for faceswap.py """
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from ._base import Masker, logger
|
|
|
|
|
|
class Mask(Masker):
|
|
""" Perform transformation to align and get landmarks """
|
|
def __init__(self, **kwargs):
|
|
git_model_id = None
|
|
model_filename = None
|
|
super().__init__(git_model_id=git_model_id, model_filename=model_filename, **kwargs)
|
|
self.input_size = 256
|
|
self.name = "Components"
|
|
self.vram = 0 # Doesn't use GPU
|
|
self.vram_per_batch = 0
|
|
self.batchsize = 1
|
|
|
|
def init_model(self):
|
|
logger.debug("No mask model to initialize")
|
|
|
|
def process_input(self, batch):
|
|
""" Compile the detected faces for prediction """
|
|
batch["feed"] = np.zeros((self.batchsize, self.input_size, self.input_size, 1),
|
|
dtype="float32")
|
|
return batch
|
|
|
|
def predict(self, batch):
|
|
""" Run model to get predictions """
|
|
for mask, face in zip(batch["feed"], batch["detected_faces"]):
|
|
parts = self.parse_parts(np.array(face.feed_landmarks))
|
|
for item in parts:
|
|
item = np.rint(np.concatenate(item)).astype("int32")
|
|
hull = cv2.convexHull(item)
|
|
cv2.fillConvexPoly(mask, hull, 1.0, lineType=cv2.LINE_AA)
|
|
batch["prediction"] = batch["feed"]
|
|
return batch
|
|
|
|
def process_output(self, batch):
|
|
""" Compile found faces for output """
|
|
return batch
|
|
|
|
@staticmethod
|
|
def parse_parts(landmarks):
|
|
""" Component face hull mask """
|
|
r_jaw = (landmarks[0:9], landmarks[17:18])
|
|
l_jaw = (landmarks[8:17], landmarks[26:27])
|
|
r_cheek = (landmarks[17:20], landmarks[8:9])
|
|
l_cheek = (landmarks[24:27], landmarks[8:9])
|
|
nose_ridge = (landmarks[19:25], landmarks[8:9],)
|
|
r_eye = (landmarks[17:22],
|
|
landmarks[27:28],
|
|
landmarks[31:36],
|
|
landmarks[8:9])
|
|
l_eye = (landmarks[22:27],
|
|
landmarks[27:28],
|
|
landmarks[31:36],
|
|
landmarks[8:9])
|
|
nose = (landmarks[27:31], landmarks[31:36])
|
|
parts = [r_jaw, l_jaw, r_cheek, l_cheek, nose_ridge, r_eye, l_eye, nose]
|
|
return parts
|