1
0
Fork 0
mirror of https://github.com/deepfakes/faceswap synced 2025-06-07 10:43:27 -04:00
faceswap/plugins/extract/mask/extended.py
torzdf 6a3b674bef
Rebase code (#1326)
* Remove tensorflow_probability requirement

* setup.py - fix progress bars

* requirements.txt: Remove pre python 3.9 packages

* update apple requirements.txt

* update INSTALL.md

* Remove python<3.9 code

* setup.py - fix Windows Installer

* typing: python3.9 compliant

* Update pytest and readthedocs python versions

* typing fixes

* Python Version updates
  - Reduce max version to 3.10
  - Default to 3.10 in installers
  - Remove incompatible 3.11 tests

* Update dependencies

* Downgrade imageio dep for Windows

* typing: merge optional unions and fixes

* Updates
  - min python version 3.10
  - typing to python 3.10 spec
  - remove pre-tf2.10 code
  - Add conda tests

* train: re-enable optimizer saving

* Update dockerfiles

* Update setup.py
  - Apple Conda deps to setup.py
  - Better Cuda + dependency handling

* bugfix: Patch logging to prevent Autograph errors

* Update dockerfiles

* Setup.py - Setup.py - stdout to utf-8

* Add more OSes to github Actions

* suppress mac-os end to end test
2023-06-27 11:27:47 +01:00

101 lines
3.5 KiB
Python

#!/usr/bin/env python3
""" Extended Mask for faceswap.py """
from __future__ import annotations
import logging
import typing as T
import cv2
import numpy as np
from ._base import BatchType, Masker
logger = logging.getLogger(__name__)
if T.TYPE_CHECKING:
from lib.align.aligned_face import AlignedFace
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 = "Extended"
self.vram = 0 # Doesn't use GPU
self.vram_per_batch = 0
self.batchsize = 1
def init_model(self) -> None:
logger.debug("No mask model to initialize")
def process_input(self, batch: BatchType) -> None:
""" Compile the detected faces for prediction """
batch.feed = np.zeros((self.batchsize, self.input_size, self.input_size, 1),
dtype="float32")
def predict(self, feed: np.ndarray) -> np.ndarray:
""" Run model to get predictions """
faces: list[AlignedFace] = feed[1]
feed = feed[0]
for mask, face in zip(feed, faces):
parts = self.parse_parts(np.array(face.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)
return feed
def process_output(self, batch: BatchType) -> None:
""" Compile found faces for output """
return
@classmethod
def _adjust_mask_top(cls, landmarks: np.ndarray) -> None:
""" Adjust the top of the mask to extend above eyebrows
Parameters
----------
landmarks: :class:`numpy.ndarray`
The 68 point landmarks to be adjusted
"""
# mid points between the side of face and eye point
ml_pnt = (landmarks[36] + landmarks[0]) // 2
mr_pnt = (landmarks[16] + landmarks[45]) // 2
# mid points between the mid points and eye
ql_pnt = (landmarks[36] + ml_pnt) // 2
qr_pnt = (landmarks[45] + mr_pnt) // 2
# Top of the eye arrays
bot_l = np.array((ql_pnt, landmarks[36], landmarks[37], landmarks[38], landmarks[39]))
bot_r = np.array((landmarks[42], landmarks[43], landmarks[44], landmarks[45], qr_pnt))
# Eyebrow arrays
top_l = landmarks[17:22]
top_r = landmarks[22:27]
# Adjust eyebrow arrays
landmarks[17:22] = top_l + ((top_l - bot_l) // 2)
landmarks[22:27] = top_r + ((top_r - bot_r) // 2)
def parse_parts(self, landmarks: np.ndarray) -> list[tuple[np.ndarray, ...]]:
""" Extended face hull mask """
self._adjust_mask_top(landmarks)
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