diff --git a/docs/full/lib/align.rst b/docs/full/lib/align.rst index cb01196739..7baf843e23 100644 --- a/docs/full/lib/align.rst +++ b/docs/full/lib/align.rst @@ -30,6 +30,27 @@ Handles aligned faces and corresponding pose estimates :show-inheritance: +aligned\_mask module +==================== + +Handles aligned storage and retrieval of Faceswap generated masks + +.. rubric:: Module Summary + +.. autosummary:: + :nosignatures: + + ~lib.align.aligned_mask.BlurMask + ~lib.align.aligned_mask.Mask + +.. rubric:: Module + +.. automodule:: lib.align.aligned_mask + :members: + :undoc-members: + :show-inheritance: + + alignments module ================= @@ -71,9 +92,7 @@ Handles detected face objects and their associated masks. .. autosummary:: :nosignatures: - ~lib.align.detected_face.BlurMask ~lib.align.detected_face.DetectedFace - ~lib.align.detected_face.Mask ~lib.align.detected_face.update_legacy_png_header .. rubric:: Module diff --git a/lib/align/__init__.py b/lib/align/__init__.py index ec00ec7798..3f5887bcd6 100644 --- a/lib/align/__init__.py +++ b/lib/align/__init__.py @@ -3,6 +3,7 @@ associated objects. """ from .aligned_face import (AlignedFace, get_adjusted_center, get_matrix_scaling, get_centered_size, transform_image) +from .aligned_mask import BlurMask, LandmarksMask, Mask from .alignments import Alignments from .constants import CenteringType, EXTRACT_RATIOS, LANDMARK_PARTS, LandmarkType -from .detected_face import BlurMask, DetectedFace, Mask, update_legacy_png_header +from .detected_face import DetectedFace, update_legacy_png_header diff --git a/lib/align/aligned_mask.py b/lib/align/aligned_mask.py new file mode 100644 index 0000000000..6a34060653 --- /dev/null +++ b/lib/align/aligned_mask.py @@ -0,0 +1,599 @@ +#!/usr/bin python3 +""" Handles retrieval and storage of Faceswap aligned masks """ + +from __future__ import annotations +import logging +import typing as T + +from zlib import compress, decompress + +import cv2 +import numpy as np + +from lib.logger import parse_class_init + +from .alignments import MaskAlignmentsFileDict +from . import get_adjusted_center, get_centered_size + +if T.TYPE_CHECKING: + from collections.abc import Callable + from .aligned_face import CenteringType + +logger = logging.getLogger(__name__) + + +class Mask(): + """ Face Mask information and convenience methods + + Holds a Faceswap mask as generated from :mod:`plugins.extract.mask` and the information + required to transform it to its original frame. + + Holds convenience methods to handle the warping, storing and retrieval of the mask. + + Parameters + ---------- + storage_size: int, optional + The size (in pixels) that the mask should be stored at. Default: 128. + storage_centering, str (optional): + The centering to store the mask at. One of `"legacy"`, `"face"`, `"head"`. + Default: `"face"` + + Attributes + ---------- + stored_size: int + The size, in pixels, of the stored mask across its height and width. + stored_centering: str + The centering that the mask is stored at. One of `"legacy"`, `"face"`, `"head"` + """ + def __init__(self, + storage_size: int = 128, + storage_centering: CenteringType = "face") -> None: + logger.trace(parse_class_init(locals())) # type:ignore[attr-defined] + self.stored_size = storage_size + self.stored_centering = storage_centering + + self._mask: bytes | None = None + self._affine_matrix: np.ndarray | None = None + self._interpolator: int | None = None + + self._blur_type: T.Literal["gaussian", "normalized"] | None = None + self._blur_passes: int = 0 + self._blur_kernel: float | int = 0 + self._threshold = 0.0 + self._dilation: tuple[T.Literal["erode", "dilate"], np.ndarray | None] = ("erode", None) + self._sub_crop_size = 0 + self._sub_crop_slices: dict[T.Literal["in", "out"], list[slice]] = {} + + self.set_blur_and_threshold() + logger.trace("Initialized: %s", self.__class__.__name__) # type:ignore[attr-defined] + + @property + def mask(self) -> np.ndarray: + """ :class:`numpy.ndarray`: The mask at the size of :attr:`stored_size` with any requested + blurring, threshold amount and centering applied.""" + mask = self.stored_mask + if self._dilation[-1] is not None or self._threshold != 0.0 or self._blur_kernel != 0: + mask = mask.copy() + self._dilate_mask(mask) + if self._threshold != 0.0: + mask[mask < self._threshold] = 0.0 + mask[mask > 255.0 - self._threshold] = 255.0 + if self._blur_kernel != 0 and self._blur_type is not None: + mask = BlurMask(self._blur_type, + mask, + self._blur_kernel, + passes=self._blur_passes).blurred + if self._sub_crop_size: # Crop the mask to the given centering + out = np.zeros((self._sub_crop_size, self._sub_crop_size, 1), dtype=mask.dtype) + slice_in, slice_out = self._sub_crop_slices["in"], self._sub_crop_slices["out"] + out[slice_out[0], slice_out[1], :] = mask[slice_in[0], slice_in[1], :] + mask = out + logger.trace("mask shape: %s", mask.shape) # type:ignore[attr-defined] + return mask + + @property + def stored_mask(self) -> np.ndarray: + """ :class:`numpy.ndarray`: The mask at the size of :attr:`stored_size` as it is stored + (i.e. with no blurring/centering applied). """ + assert self._mask is not None + dims = (self.stored_size, self.stored_size, 1) + mask = np.frombuffer(decompress(self._mask), dtype="uint8").reshape(dims) + logger.trace("stored mask shape: %s", mask.shape) # type:ignore[attr-defined] + return mask + + @property + def original_roi(self) -> np.ndarray: + """ :class: `numpy.ndarray`: The original region of interest of the mask in the + source frame. """ + points = np.array([[0, 0], + [0, self.stored_size - 1], + [self.stored_size - 1, self.stored_size - 1], + [self.stored_size - 1, 0]], np.int32).reshape((-1, 1, 2)) + matrix = cv2.invertAffineTransform(self.affine_matrix) + roi = cv2.transform(points, matrix).reshape((4, 2)) + logger.trace("Returning: %s", roi) # type:ignore[attr-defined] + return roi + + @property + def affine_matrix(self) -> np.ndarray: + """ :class: `numpy.ndarray`: The affine matrix to transpose the mask to a full frame. """ + assert self._affine_matrix is not None + return self._affine_matrix + + @property + def interpolator(self) -> int: + """ int: The cv2 interpolator required to transpose the mask to a full frame. """ + assert self._interpolator is not None + return self._interpolator + + def _dilate_mask(self, mask: np.ndarray) -> None: + """ Erode/Dilate the mask. The action is performed in-place on the given mask. + + No action is performed if a dilation amount has not been set + + Parameters + ---------- + mask: :class:`numpy.ndarray` + The mask to be eroded/dilated + """ + if self._dilation[-1] is None: + return + + func = cv2.erode if self._dilation[0] == "erode" else cv2.dilate + func(mask, self._dilation[-1], dst=mask, iterations=1) + + def get_full_frame_mask(self, width: int, height: int) -> np.ndarray: + """ Return the stored mask in a full size frame of the given dimensions + + Parameters + ---------- + width: int + The width of the original frame that the mask was extracted from + height: int + The height of the original frame that the mask was extracted from + + Returns + ------- + :class:`numpy.ndarray`: The mask affined to the original full frame of the given dimensions + """ + frame = np.zeros((width, height, 1), dtype="uint8") + mask = cv2.warpAffine(self.mask, + self.affine_matrix, + (width, height), + frame, + flags=cv2.WARP_INVERSE_MAP | self.interpolator, + borderMode=cv2.BORDER_CONSTANT) + logger.trace("mask shape: %s, mask dtype: %s, mask min: %s, " # type:ignore[attr-defined] + "mask max: %s", mask.shape, mask.dtype, mask.min(), mask.max()) + return mask + + def add(self, mask: np.ndarray, affine_matrix: np.ndarray, interpolator: int) -> None: + """ Add a Faceswap mask to this :class:`Mask`. + + The mask should be the original output from :mod:`plugins.extract.mask` + + Parameters + ---------- + mask: :class:`numpy.ndarray` + The mask that is to be added as output from :mod:`plugins.extract.mask` + It should be in the range 0.0 - 1.0 ideally with a ``dtype`` of ``float32`` + affine_matrix: :class:`numpy.ndarray` + The transformation matrix required to transform the mask to the original frame. + interpolator, int: + The CV2 interpolator required to transform this mask to it's original frame + """ + logger.trace("mask shape: %s, mask dtype: %s, mask min: %s, " # type:ignore[attr-defined] + "mask max: %s, affine_matrix: %s, interpolator: %s)", + mask.shape, mask.dtype, mask.min(), affine_matrix, mask.max(), interpolator) + self._affine_matrix = self._adjust_affine_matrix(mask.shape[0], affine_matrix) + self._interpolator = interpolator + self.replace_mask(mask) + + def replace_mask(self, mask: np.ndarray) -> None: + """ Replace the existing :attr:`_mask` with the given mask. + + Parameters + ---------- + mask: :class:`numpy.ndarray` + The mask that is to be added as output from :mod:`plugins.extract.mask`. + It should be in the range 0.0 - 1.0 ideally with a ``dtype`` of ``float32`` + """ + mask = (cv2.resize(mask * 255.0, + (self.stored_size, self.stored_size), + interpolation=cv2.INTER_AREA)).astype("uint8") + self._mask = compress(mask.tobytes()) + + def set_dilation(self, amount: float) -> None: + """ Set the internal dilation object for returned masks + + Parameters + ---------- + amount: float + The amount of erosion/dilation to apply as a percentage of the total mask size. + Negative values erode the mask. Positive values dilate the mask + """ + if amount == 0: + self._dilation = ("erode", None) + return + + action: T.Literal["erode", "dilate"] = "erode" if amount < 0 else "dilate" + kernel = int(round(self.stored_size * abs(amount / 100.), 0)) + self._dilation = (action, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel, kernel))) + + logger.trace("action: '%s', amount: %s, kernel: %s, ", # type:ignore[attr-defined] + action, amount, kernel) + + def set_blur_and_threshold(self, + blur_kernel: int = 0, + blur_type: T.Literal["gaussian", "normalized"] | None = "gaussian", + blur_passes: int = 1, + threshold: int = 0) -> None: + """ Set the internal blur kernel and threshold amount for returned masks + + Parameters + ---------- + blur_kernel: int, optional + The kernel size, in pixels to apply gaussian blurring to the mask. Set to 0 for no + blurring. Should be odd, if an even number is passed in (outside of 0) then it is + rounded up to the next odd number. Default: 0 + blur_type: ["gaussian", "normalized"], optional + The blur type to use. ``gaussian`` or ``normalized`` box filter. Default: ``gaussian`` + blur_passes: int, optional + The number of passed to perform when blurring. Default: 1 + threshold: int, optional + The threshold amount to minimize/maximize mask values to 0 and 100. Percentage value. + Default: 0 + """ + logger.trace("blur_kernel: %s, blur_type: %s, " # type:ignore[attr-defined] + "blur_passes: %s, threshold: %s", + blur_kernel, blur_type, blur_passes, threshold) + if blur_type is not None: + blur_kernel += 0 if blur_kernel == 0 or blur_kernel % 2 == 1 else 1 + self._blur_kernel = blur_kernel + self._blur_type = blur_type + self._blur_passes = blur_passes + self._threshold = (threshold / 100.0) * 255.0 + + def set_sub_crop(self, + source_offset: np.ndarray, + target_offset: np.ndarray, + centering: CenteringType, + coverage_ratio: float = 1.0) -> None: + """ Set the internal crop area of the mask to be returned. + + This impacts the returned mask from :attr:`mask` if the requested mask is required for + different face centering than what has been stored. + + Parameters + ---------- + source_offset: :class:`numpy.ndarray` + The (x, y) offset for the mask at its stored centering + target_offset: :class:`numpy.ndarray` + The (x, y) offset for the mask at the requested target centering + centering: str + The centering to set the sub crop area for. One of `"legacy"`, `"face"`. `"head"` + coverage_ratio: float, optional + The coverage ratio to be applied to the target image. ``None`` for default (1.0). + Default: ``None`` + """ + if centering == self.stored_centering and coverage_ratio == 1.0: + return + + center = get_adjusted_center(self.stored_size, + source_offset, + target_offset, + self.stored_centering) + crop_size = get_centered_size(self.stored_centering, + centering, + self.stored_size, + coverage_ratio=coverage_ratio) + roi = np.array([center - crop_size // 2, center + crop_size // 2]).ravel() + + self._sub_crop_size = crop_size + self._sub_crop_slices["in"] = [slice(max(roi[1], 0), max(roi[3], 0)), + slice(max(roi[0], 0), max(roi[2], 0))] + self._sub_crop_slices["out"] = [ + slice(max(roi[1] * -1, 0), + crop_size - min(crop_size, max(0, roi[3] - self.stored_size))), + slice(max(roi[0] * -1, 0), + crop_size - min(crop_size, max(0, roi[2] - self.stored_size)))] + + logger.trace("src_size: %s, coverage_ratio: %s, " # type:ignore[attr-defined] + "sub_crop_size: %s, sub_crop_slices: %s", + roi, coverage_ratio, self._sub_crop_size, self._sub_crop_slices) + + def _adjust_affine_matrix(self, mask_size: int, affine_matrix: np.ndarray) -> np.ndarray: + """ Adjust the affine matrix for the mask's storage size + + Parameters + ---------- + mask_size: int + The original size of the mask. + affine_matrix: :class:`numpy.ndarray` + The affine matrix to transform the mask at original size to the parent frame. + + Returns + ------- + affine_matrix: :class:`numpy,ndarray` + The affine matrix adjusted for the mask at its stored dimensions. + """ + zoom = self.stored_size / mask_size + zoom_mat = np.array([[zoom, 0, 0.], [0, zoom, 0.]]) + adjust_mat = np.dot(zoom_mat, np.concatenate((affine_matrix, np.array([[0., 0., 1.]])))) + logger.trace("storage_size: %s, mask_size: %s, zoom: %s, " # type:ignore[attr-defined] + "original matrix: %s, adjusted_matrix: %s", self.stored_size, mask_size, zoom, + affine_matrix.shape, adjust_mat.shape) + return adjust_mat + + def to_dict(self, is_png=False) -> MaskAlignmentsFileDict: + """ Convert the mask to a dictionary for saving to an alignments file + + Parameters + ---------- + is_png: bool + ``True`` if the dictionary is being created for storage in a png header otherwise + ``False``. Default: ``False`` + + Returns + ------- + dict: + The :class:`Mask` for saving to an alignments file. Contains the keys ``mask``, + ``affine_matrix``, ``interpolator``, ``stored_size``, ``stored_centering`` + """ + assert self._mask is not None + affine_matrix = self.affine_matrix.tolist() if is_png else self.affine_matrix + retval = MaskAlignmentsFileDict(mask=self._mask, + affine_matrix=affine_matrix, + interpolator=self.interpolator, + stored_size=self.stored_size, + stored_centering=self.stored_centering) + logger.trace({k: v if k != "mask" else type(v) # type:ignore[attr-defined] + for k, v in retval.items()}) + return retval + + def to_png_meta(self) -> MaskAlignmentsFileDict: + """ Convert the mask to a dictionary supported by png itxt headers. + + Returns + ------- + dict: + The :class:`Mask` for saving to an alignments file. Contains the keys ``mask``, + ``affine_matrix``, ``interpolator``, ``stored_size``, ``stored_centering`` + """ + return self.to_dict(is_png=True) + + def from_dict(self, mask_dict: MaskAlignmentsFileDict) -> None: + """ Populates the :class:`Mask` from a dictionary loaded from an alignments file. + + Parameters + ---------- + mask_dict: dict + A dictionary stored in an alignments file containing the keys ``mask``, + ``affine_matrix``, ``interpolator``, ``stored_size``, ``stored_centering`` + """ + self._mask = mask_dict["mask"] + affine_matrix = mask_dict["affine_matrix"] + self._affine_matrix = (affine_matrix if isinstance(affine_matrix, np.ndarray) + else np.array(affine_matrix, dtype="float64")) + self._interpolator = mask_dict["interpolator"] + self.stored_size = mask_dict["stored_size"] + centering = mask_dict.get("stored_centering") + self.stored_centering = "face" if centering is None else centering + logger.trace({k: v if k != "mask" else type(v) # type:ignore[attr-defined] + for k, v in mask_dict.items()}) + + +class LandmarksMask(Mask): + """ Create a single channel mask from aligned landmark points. + + Landmarks masks are created on the fly, so the stored centering and size should be the same as + the aligned face that the mask will be applied to. As the masks are created on the fly, blur + + dilation is applied to the mask at creation (prior to compression) rather than after + decompression when requested. + + Note + ---- + Threshold is not used for Landmarks mask as the mask is binary + + Parameters + ---------- + points: list + A list of landmark points that correspond to the given storage_size to create + the mask. Each item in the list should be a :class:`numpy.ndarray` that a filled + convex polygon will be created from + storage_size: int, optional + The size (in pixels) that the compressed mask should be stored at. Default: 128. + storage_centering, str (optional): + The centering to store the mask at. One of `"legacy"`, `"face"`, `"head"`. + Default: `"face"` + dilation: float, optional + The amount of dilation to apply to the mask. as a percentage of the mask size. Default: 0.0 + """ + def __init__(self, + points: list[np.ndarray], + storage_size: int = 128, + storage_centering: CenteringType = "face", + dilation: float = 0.0) -> None: + super().__init__(storage_size=storage_size, storage_centering=storage_centering) + self._points = points + self.set_dilation(dilation) + + @property + def mask(self) -> np.ndarray: + """ :class:`numpy.ndarray`: Overrides the default mask property, creating the processed + mask at first call and compressing it. The decompressed mask is returned from this + property. """ + return self.stored_mask + + def generate_mask(self, affine_matrix: np.ndarray, interpolator: int) -> None: + """ Generate the mask. + + Creates the mask applying any requested dilation and blurring and assigns compressed mask + to :attr:`_mask` + + Parameters + ---------- + affine_matrix: :class:`numpy.ndarray` + The transformation matrix required to transform the mask to the original frame. + interpolator, int: + The CV2 interpolator required to transform this mask to it's original frame + """ + mask = np.zeros((self.stored_size, self.stored_size, 1), dtype="float32") + for landmarks in self._points: + lms = np.rint(landmarks).astype("int") + cv2.fillConvexPoly(mask, cv2.convexHull(lms), [1.0], lineType=cv2.LINE_AA) + if self._dilation[-1] is not None: + self._dilate_mask(mask) + if self._blur_kernel != 0 and self._blur_type is not None: + mask = BlurMask(self._blur_type, + mask, + self._blur_kernel, + passes=self._blur_passes).blurred + logger.trace("mask: (shape: %s, dtype: %s)", # type:ignore[attr-defined] + mask.shape, mask.dtype) + self.add(mask, affine_matrix, interpolator) + + +class BlurMask(): + """ Factory class to return the correct blur object for requested blur type. + + Works for square images only. Currently supports Gaussian and Normalized Box Filters. + + Parameters + ---------- + blur_type: ["gaussian", "normalized"] + The type of blur to use + mask: :class:`numpy.ndarray` + The mask to apply the blur to + kernel: int or float + Either the kernel size (in pixels) or the size of the kernel as a ratio of mask size + is_ratio: bool, optional + Whether the given :attr:`kernel` parameter is a ratio or not. If ``True`` then the + actual kernel size will be calculated from the given ratio and the mask size. If + ``False`` then the kernel size will be set directly from the :attr:`kernel` parameter. + Default: ``False`` + passes: int, optional + The number of passes to perform when blurring. Default: ``1`` + + Example + ------- + >>> print(mask.shape) + (128, 128, 1) + >>> new_mask = BlurMask("gaussian", mask, 3, is_ratio=False, passes=1).blurred + >>> print(new_mask.shape) + (128, 128, 1) + """ + def __init__(self, + blur_type: T.Literal["gaussian", "normalized"], + mask: np.ndarray, + kernel: int | float, + is_ratio: bool = False, + passes: int = 1) -> None: + logger.trace(parse_class_init(locals())) # type:ignore[attr-defined] + self._blur_type = blur_type + self._mask = mask + self._passes = passes + kernel_size = self._get_kernel_size(kernel, is_ratio) + self._kernel_size = self._get_kernel_tuple(kernel_size) + logger.trace("Initialized %s", self.__class__.__name__) # type:ignore[attr-defined] + + @property + def blurred(self) -> np.ndarray: + """ :class:`numpy.ndarray`: The final mask with blurring applied. """ + func = self._func_mapping[self._blur_type] + kwargs = self._get_kwargs() + blurred = self._mask + for i in range(self._passes): + assert isinstance(kwargs["ksize"], tuple) + ksize = int(kwargs["ksize"][0]) + logger.trace("Pass: %s, kernel_size: %s", # type:ignore[attr-defined] + i + 1, (ksize, ksize)) + blurred = func(blurred, **kwargs) + ksize = int(round(ksize * self._multipass_factor)) + kwargs["ksize"] = self._get_kernel_tuple(ksize) + blurred = blurred[..., None] + logger.trace("Returning blurred mask. Shape: %s", # type:ignore[attr-defined] + blurred.shape) + return blurred + + @property + def _multipass_factor(self) -> float: + """ For multiple passes the kernel must be scaled down. This value is + different for box filter and gaussian """ + factor = {"gaussian": 0.8, "normalized": 0.5} + return factor[self._blur_type] + + @property + def _sigma(self) -> T.Literal[0]: + """ int: The Sigma for Gaussian Blur. Returns 0 to force calculation from kernel size. """ + return 0 + + @property + def _func_mapping(self) -> dict[T.Literal["gaussian", "normalized"], Callable]: + """ dict: :attr:`_blur_type` mapped to cv2 Function name. """ + return {"gaussian": cv2.GaussianBlur, "normalized": cv2.blur} + + @property + def _kwarg_requirements(self) -> dict[T.Literal["gaussian", "normalized"], list[str]]: + """ dict: :attr:`_blur_type` mapped to cv2 Function required keyword arguments. """ + return {"gaussian": ['ksize', 'sigmaX'], "normalized": ['ksize']} + + @property + def _kwarg_mapping(self) -> dict[str, int | tuple[int, int]]: + """ dict: cv2 function keyword arguments mapped to their parameters. """ + return {"ksize": self._kernel_size, "sigmaX": self._sigma} + + def _get_kernel_size(self, kernel: int | float, is_ratio: bool) -> int: + """ Set the kernel size to absolute value. + + If :attr:`is_ratio` is ``True`` then the kernel size is calculated from the given ratio and + the :attr:`_mask` size, otherwise the given kernel size is just returned. + + Parameters + ---------- + kernel: int or float + Either the kernel size (in pixels) or the size of the kernel as a ratio of mask size + is_ratio: bool, optional + Whether the given :attr:`kernel` parameter is a ratio or not. If ``True`` then the + actual kernel size will be calculated from the given ratio and the mask size. If + ``False`` then the kernel size will be set directly from the :attr:`kernel` parameter. + + Returns + ------- + int + The size (in pixels) of the blur kernel + """ + if not is_ratio: + return int(kernel) + + mask_diameter = np.sqrt(np.sum(self._mask)) + radius = round(max(1., mask_diameter * kernel / 100.)) + kernel_size = int(radius * 2 + 1) + logger.trace("kernel_size: %s", kernel_size) # type:ignore[attr-defined] + return kernel_size + + @staticmethod + def _get_kernel_tuple(kernel_size: int) -> tuple[int, int]: + """ Make sure kernel_size is odd and return it as a tuple. + + Parameters + ---------- + kernel_size: int + The size in pixels of the blur kernel + + Returns + ------- + tuple + The kernel size as a tuple of ('int', 'int') + """ + kernel_size += 1 if kernel_size % 2 == 0 else 0 + retval = (kernel_size, kernel_size) + logger.trace(retval) # type:ignore[attr-defined] + return retval + + def _get_kwargs(self) -> dict[str, int | tuple[int, int]]: + """ dict: the valid keyword arguments for the requested :attr:`_blur_type` """ + retval = {kword: self._kwarg_mapping[kword] + for kword in self._kwarg_requirements[self._blur_type]} + logger.trace("BlurMask kwargs: %s", retval) # type:ignore[attr-defined] + return retval diff --git a/lib/align/detected_face.py b/lib/align/detected_face.py index 8ebea9a599..fb8ef37f98 100644 --- a/lib/align/detected_face.py +++ b/lib/align/detected_face.py @@ -8,17 +8,18 @@ from hashlib import sha1 from zlib import compress, decompress -import cv2 import numpy as np from lib.image import encode_image, read_image +from lib.logger import parse_class_init from lib.utils import FaceswapError -from .alignments import (Alignments, AlignmentFileDict, MaskAlignmentsFileDict, - PNGHeaderAlignmentsDict, PNGHeaderDict, PNGHeaderSourceDict) -from . import AlignedFace, get_adjusted_center, get_centered_size, LANDMARK_PARTS +from .alignments import (Alignments, AlignmentFileDict, PNGHeaderAlignmentsDict, + PNGHeaderDict, PNGHeaderSourceDict) +from .aligned_face import AlignedFace +from .aligned_mask import LandmarksMask, Mask +from .constants import LANDMARK_PARTS if T.TYPE_CHECKING: - from collections.abc import Callable from .aligned_face import CenteringType logger = logging.getLogger(__name__) @@ -53,7 +54,7 @@ class DetectedFace(): of 68 `(x, y)` ``tuples`` with each of the landmark co-ordinates. mask: dict The generated mask(s) for the face as generated in :mod:`plugins.extract.mask`. Must be a - dict of {**name** (`str`): :class:`Mask`}. + dict of {**name** (`str`): :class:`~lib.align.aligned_mask.Mask`}. Attributes ---------- @@ -77,7 +78,7 @@ class DetectedFace(): The 68 point landmarks as discovered in :mod:`plugins.extract.align`. mask: dict The generated mask(s) for the face as generated in :mod:`plugins.extract.mask`. Is a - dict of {**name** (`str`): :class:`Mask`}. + dict of {**name** (`str`): :class:`~lib.align.aligned_mask.Mask`}. """ def __init__(self, image: np.ndarray | None = None, @@ -86,13 +87,9 @@ def __init__(self, top: int | None = None, height: int | None = None, landmarks_xy: np.ndarray | None = None, - mask: dict[str, "Mask"] | None = None, + mask: dict[str, Mask] | None = None, filename: str | None = None) -> None: - logger.trace("Initializing %s: (image: %s, left: %s, " # type:ignore[attr-defined] - "width: %s, top: %s, height: %s, landmarks_xy: %s, mask: %s, filename: %s)", - self.__class__.__name__, - image.shape if image is not None and image.any() else image, left, width, top, - height, landmarks_xy, mask, filename) + logger.trace(parse_class_init(locals())) # type:ignore[attr-defined] self.image = image self.left = left self.width = width @@ -143,7 +140,7 @@ def add_mask(self, interpolator: int, storage_size: int = 128, storage_centering: CenteringType = "face") -> None: - """ Add a :class:`Mask` to this detected face + """ Add a :class:`~lib.align.aligned_mask.Mask` to this detected face The mask should be the original output from :mod:`plugins.extract.mask` If a mask with this name already exists it will be overwritten by the given @@ -211,7 +208,7 @@ def get_landmark_mask(self, area: T.Literal["eye", "face", "mouth"], blur_kernel: int, dilation: float) -> np.ndarray: - """ Add a :class:`LandmarksMask` to this detected face + """ Add a :class:`L~lib.align.aligned_mask.LandmarksMask` to this detected face Landmark based masks are generated from face Aligned Face landmark points. An aligned face must be loaded. As the data is coming from the already aligned face, no further mask @@ -273,8 +270,8 @@ def store_training_masks(self, A list of training mask. Must be all be uint-8 3D arrays of the same size in 0-255 range delete_masks: bool, optional - ``True`` to delete any of the :class:`Mask` objects owned by this detected face. Use to - free up unrequired memory usage. Default: ``False`` + ``True`` to delete any of the :class:`~lib.align.aligned_mask.Mask` objects owned by + this detected face. Use to free up unrequired memory usage. Default: ``False`` """ if delete_masks: del self.mask @@ -496,588 +493,6 @@ def load_aligned(self, is_legacy=is_aligned and is_legacy) -class Mask(): - """ Face Mask information and convenience methods - - Holds a Faceswap mask as generated from :mod:`plugins.extract.mask` and the information - required to transform it to its original frame. - - Holds convenience methods to handle the warping, storing and retrieval of the mask. - - Parameters - ---------- - storage_size: int, optional - The size (in pixels) that the mask should be stored at. Default: 128. - storage_centering, str (optional): - The centering to store the mask at. One of `"legacy"`, `"face"`, `"head"`. - Default: `"face"` - - Attributes - ---------- - stored_size: int - The size, in pixels, of the stored mask across its height and width. - stored_centering: str - The centering that the mask is stored at. One of `"legacy"`, `"face"`, `"head"` - """ - def __init__(self, - storage_size: int = 128, - storage_centering: CenteringType = "face") -> None: - logger.trace("Initializing: %s (storage_size: %s, " # type:ignore[attr-defined] - "storage_centering: %s)", - self.__class__.__name__, storage_size, storage_centering) - self.stored_size = storage_size - self.stored_centering = storage_centering - - self._mask: bytes | None = None - self._affine_matrix: np.ndarray | None = None - self._interpolator: int | None = None - - self._blur_type: T.Literal["gaussian", "normalized"] | None = None - self._blur_passes: int = 0 - self._blur_kernel: float | int = 0 - self._threshold = 0.0 - self._dilation: tuple[T.Literal["erode", "dilate"], np.ndarray | None] = ("erode", None) - self._sub_crop_size = 0 - self._sub_crop_slices: dict[T.Literal["in", "out"], list[slice]] = {} - - self.set_blur_and_threshold() - logger.trace("Initialized: %s", self.__class__.__name__) # type:ignore[attr-defined] - - @property - def mask(self) -> np.ndarray: - """ :class:`numpy.ndarray`: The mask at the size of :attr:`stored_size` with any requested - blurring, threshold amount and centering applied.""" - mask = self.stored_mask - if self._dilation[-1] is not None or self._threshold != 0.0 or self._blur_kernel != 0: - mask = mask.copy() - self._dilate_mask(mask) - if self._threshold != 0.0: - mask[mask < self._threshold] = 0.0 - mask[mask > 255.0 - self._threshold] = 255.0 - if self._blur_kernel != 0 and self._blur_type is not None: - mask = BlurMask(self._blur_type, - mask, - self._blur_kernel, - passes=self._blur_passes).blurred - if self._sub_crop_size: # Crop the mask to the given centering - out = np.zeros((self._sub_crop_size, self._sub_crop_size, 1), dtype=mask.dtype) - slice_in, slice_out = self._sub_crop_slices["in"], self._sub_crop_slices["out"] - out[slice_out[0], slice_out[1], :] = mask[slice_in[0], slice_in[1], :] - mask = out - logger.trace("mask shape: %s", mask.shape) # type:ignore[attr-defined] - return mask - - @property - def stored_mask(self) -> np.ndarray: - """ :class:`numpy.ndarray`: The mask at the size of :attr:`stored_size` as it is stored - (i.e. with no blurring/centering applied). """ - assert self._mask is not None - dims = (self.stored_size, self.stored_size, 1) - mask = np.frombuffer(decompress(self._mask), dtype="uint8").reshape(dims) - logger.trace("stored mask shape: %s", mask.shape) # type:ignore[attr-defined] - return mask - - @property - def original_roi(self) -> np.ndarray: - """ :class: `numpy.ndarray`: The original region of interest of the mask in the - source frame. """ - points = np.array([[0, 0], - [0, self.stored_size - 1], - [self.stored_size - 1, self.stored_size - 1], - [self.stored_size - 1, 0]], np.int32).reshape((-1, 1, 2)) - matrix = cv2.invertAffineTransform(self.affine_matrix) - roi = cv2.transform(points, matrix).reshape((4, 2)) - logger.trace("Returning: %s", roi) # type:ignore[attr-defined] - return roi - - @property - def affine_matrix(self) -> np.ndarray: - """ :class: `numpy.ndarray`: The affine matrix to transpose the mask to a full frame. """ - assert self._affine_matrix is not None - return self._affine_matrix - - @property - def interpolator(self) -> int: - """ int: The cv2 interpolator required to transpose the mask to a full frame. """ - assert self._interpolator is not None - return self._interpolator - - def _dilate_mask(self, mask: np.ndarray) -> None: - """ Erode/Dilate the mask. The action is performed in-place on the given mask. - - No action is performed if a dilation amount has not been set - - Parameters - ---------- - mask: :class:`numpy.ndarray` - The mask to be eroded/dilated - """ - if self._dilation[-1] is None: - return - - func = cv2.erode if self._dilation[0] == "erode" else cv2.dilate - func(mask, self._dilation[-1], dst=mask, iterations=1) - - def get_full_frame_mask(self, width: int, height: int) -> np.ndarray: - """ Return the stored mask in a full size frame of the given dimensions - - Parameters - ---------- - width: int - The width of the original frame that the mask was extracted from - height: int - The height of the original frame that the mask was extracted from - - Returns - ------- - :class:`numpy.ndarray`: The mask affined to the original full frame of the given dimensions - """ - frame = np.zeros((width, height, 1), dtype="uint8") - mask = cv2.warpAffine(self.mask, - self.affine_matrix, - (width, height), - frame, - flags=cv2.WARP_INVERSE_MAP | self.interpolator, - borderMode=cv2.BORDER_CONSTANT) - logger.trace("mask shape: %s, mask dtype: %s, mask min: %s, " # type:ignore[attr-defined] - "mask max: %s", mask.shape, mask.dtype, mask.min(), mask.max()) - return mask - - def add(self, mask: np.ndarray, affine_matrix: np.ndarray, interpolator: int) -> None: - """ Add a Faceswap mask to this :class:`Mask`. - - The mask should be the original output from :mod:`plugins.extract.mask` - - Parameters - ---------- - mask: :class:`numpy.ndarray` - The mask that is to be added as output from :mod:`plugins.extract.mask` - It should be in the range 0.0 - 1.0 ideally with a ``dtype`` of ``float32`` - affine_matrix: :class:`numpy.ndarray` - The transformation matrix required to transform the mask to the original frame. - interpolator, int: - The CV2 interpolator required to transform this mask to it's original frame - """ - logger.trace("mask shape: %s, mask dtype: %s, mask min: %s, " # type:ignore[attr-defined] - "mask max: %s, affine_matrix: %s, interpolator: %s)", - mask.shape, mask.dtype, mask.min(), affine_matrix, mask.max(), interpolator) - self._affine_matrix = self._adjust_affine_matrix(mask.shape[0], affine_matrix) - self._interpolator = interpolator - self.replace_mask(mask) - - def replace_mask(self, mask: np.ndarray) -> None: - """ Replace the existing :attr:`_mask` with the given mask. - - Parameters - ---------- - mask: :class:`numpy.ndarray` - The mask that is to be added as output from :mod:`plugins.extract.mask`. - It should be in the range 0.0 - 1.0 ideally with a ``dtype`` of ``float32`` - """ - mask = (cv2.resize(mask * 255.0, - (self.stored_size, self.stored_size), - interpolation=cv2.INTER_AREA)).astype("uint8") - self._mask = compress(mask.tobytes()) - - def set_dilation(self, amount: float) -> None: - """ Set the internal dilation object for returned masks - - Parameters - ---------- - amount: float - The amount of erosion/dilation to apply as a percentage of the total mask size. - Negative values erode the mask. Positive values dilate the mask - """ - if amount == 0: - self._dilation = ("erode", None) - return - - action: T.Literal["erode", "dilate"] = "erode" if amount < 0 else "dilate" - kernel = int(round(self.stored_size * abs(amount / 100.), 0)) - self._dilation = (action, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel, kernel))) - - logger.trace("action: '%s', amount: %s, kernel: %s, ", # type:ignore[attr-defined] - action, amount, kernel) - - def set_blur_and_threshold(self, - blur_kernel: int = 0, - blur_type: T.Literal["gaussian", "normalized"] | None = "gaussian", - blur_passes: int = 1, - threshold: int = 0) -> None: - """ Set the internal blur kernel and threshold amount for returned masks - - Parameters - ---------- - blur_kernel: int, optional - The kernel size, in pixels to apply gaussian blurring to the mask. Set to 0 for no - blurring. Should be odd, if an even number is passed in (outside of 0) then it is - rounded up to the next odd number. Default: 0 - blur_type: ["gaussian", "normalized"], optional - The blur type to use. ``gaussian`` or ``normalized`` box filter. Default: ``gaussian`` - blur_passes: int, optional - The number of passed to perform when blurring. Default: 1 - threshold: int, optional - The threshold amount to minimize/maximize mask values to 0 and 100. Percentage value. - Default: 0 - """ - logger.trace("blur_kernel: %s, blur_type: %s, " # type:ignore[attr-defined] - "blur_passes: %s, threshold: %s", - blur_kernel, blur_type, blur_passes, threshold) - if blur_type is not None: - blur_kernel += 0 if blur_kernel == 0 or blur_kernel % 2 == 1 else 1 - self._blur_kernel = blur_kernel - self._blur_type = blur_type - self._blur_passes = blur_passes - self._threshold = (threshold / 100.0) * 255.0 - - def set_sub_crop(self, - source_offset: np.ndarray, - target_offset: np.ndarray, - centering: CenteringType, - coverage_ratio: float = 1.0) -> None: - """ Set the internal crop area of the mask to be returned. - - This impacts the returned mask from :attr:`mask` if the requested mask is required for - different face centering than what has been stored. - - Parameters - ---------- - source_offset: :class:`numpy.ndarray` - The (x, y) offset for the mask at its stored centering - target_offset: :class:`numpy.ndarray` - The (x, y) offset for the mask at the requested target centering - centering: str - The centering to set the sub crop area for. One of `"legacy"`, `"face"`. `"head"` - coverage_ratio: float, optional - The coverage ratio to be applied to the target image. ``None`` for default (1.0). - Default: ``None`` - """ - if centering == self.stored_centering and coverage_ratio == 1.0: - return - - center = get_adjusted_center(self.stored_size, - source_offset, - target_offset, - self.stored_centering) - crop_size = get_centered_size(self.stored_centering, - centering, - self.stored_size, - coverage_ratio=coverage_ratio) - roi = np.array([center - crop_size // 2, center + crop_size // 2]).ravel() - - self._sub_crop_size = crop_size - self._sub_crop_slices["in"] = [slice(max(roi[1], 0), max(roi[3], 0)), - slice(max(roi[0], 0), max(roi[2], 0))] - self._sub_crop_slices["out"] = [ - slice(max(roi[1] * -1, 0), - crop_size - min(crop_size, max(0, roi[3] - self.stored_size))), - slice(max(roi[0] * -1, 0), - crop_size - min(crop_size, max(0, roi[2] - self.stored_size)))] - - logger.trace("src_size: %s, coverage_ratio: %s, " # type:ignore[attr-defined] - "sub_crop_size: %s, sub_crop_slices: %s", - roi, coverage_ratio, self._sub_crop_size, self._sub_crop_slices) - - def _adjust_affine_matrix(self, mask_size: int, affine_matrix: np.ndarray) -> np.ndarray: - """ Adjust the affine matrix for the mask's storage size - - Parameters - ---------- - mask_size: int - The original size of the mask. - affine_matrix: :class:`numpy.ndarray` - The affine matrix to transform the mask at original size to the parent frame. - - Returns - ------- - affine_matrix: :class:`numpy,ndarray` - The affine matrix adjusted for the mask at its stored dimensions. - """ - zoom = self.stored_size / mask_size - zoom_mat = np.array([[zoom, 0, 0.], [0, zoom, 0.]]) - adjust_mat = np.dot(zoom_mat, np.concatenate((affine_matrix, np.array([[0., 0., 1.]])))) - logger.trace("storage_size: %s, mask_size: %s, zoom: %s, " # type:ignore[attr-defined] - "original matrix: %s, adjusted_matrix: %s", self.stored_size, mask_size, zoom, - affine_matrix.shape, adjust_mat.shape) - return adjust_mat - - def to_dict(self, is_png=False) -> MaskAlignmentsFileDict: - """ Convert the mask to a dictionary for saving to an alignments file - - Parameters - ---------- - is_png: bool - ``True`` if the dictionary is being created for storage in a png header otherwise - ``False``. Default: ``False`` - - Returns - ------- - dict: - The :class:`Mask` for saving to an alignments file. Contains the keys ``mask``, - ``affine_matrix``, ``interpolator``, ``stored_size``, ``stored_centering`` - """ - assert self._mask is not None - affine_matrix = self.affine_matrix.tolist() if is_png else self.affine_matrix - retval = MaskAlignmentsFileDict(mask=self._mask, - affine_matrix=affine_matrix, - interpolator=self.interpolator, - stored_size=self.stored_size, - stored_centering=self.stored_centering) - logger.trace({k: v if k != "mask" else type(v) # type:ignore[attr-defined] - for k, v in retval.items()}) - return retval - - def to_png_meta(self) -> MaskAlignmentsFileDict: - """ Convert the mask to a dictionary supported by png itxt headers. - - Returns - ------- - dict: - The :class:`Mask` for saving to an alignments file. Contains the keys ``mask``, - ``affine_matrix``, ``interpolator``, ``stored_size``, ``stored_centering`` - """ - return self.to_dict(is_png=True) - - def from_dict(self, mask_dict: MaskAlignmentsFileDict) -> None: - """ Populates the :class:`Mask` from a dictionary loaded from an alignments file. - - Parameters - ---------- - mask_dict: dict - A dictionary stored in an alignments file containing the keys ``mask``, - ``affine_matrix``, ``interpolator``, ``stored_size``, ``stored_centering`` - """ - self._mask = mask_dict["mask"] - affine_matrix = mask_dict["affine_matrix"] - self._affine_matrix = (affine_matrix if isinstance(affine_matrix, np.ndarray) - else np.array(affine_matrix, dtype="float64")) - self._interpolator = mask_dict["interpolator"] - self.stored_size = mask_dict["stored_size"] - centering = mask_dict.get("stored_centering") - self.stored_centering = "face" if centering is None else centering - logger.trace({k: v if k != "mask" else type(v) # type:ignore[attr-defined] - for k, v in mask_dict.items()}) - - -class LandmarksMask(Mask): - """ Create a single channel mask from aligned landmark points. - - Landmarks masks are created on the fly, so the stored centering and size should be the same as - the aligned face that the mask will be applied to. As the masks are created on the fly, blur + - dilation is applied to the mask at creation (prior to compression) rather than after - decompression when requested. - - Note - ---- - Threshold is not used for Landmarks mask as the mask is binary - - Parameters - ---------- - points: list - A list of landmark points that correspond to the given storage_size to create - the mask. Each item in the list should be a :class:`numpy.ndarray` that a filled - convex polygon will be created from - storage_size: int, optional - The size (in pixels) that the compressed mask should be stored at. Default: 128. - storage_centering, str (optional): - The centering to store the mask at. One of `"legacy"`, `"face"`, `"head"`. - Default: `"face"` - dilation: float, optional - The amount of dilation to apply to the mask. as a percentage of the mask size. Default: 0.0 - """ - def __init__(self, - points: list[np.ndarray], - storage_size: int = 128, - storage_centering: CenteringType = "face", - dilation: float = 0.0) -> None: - super().__init__(storage_size=storage_size, storage_centering=storage_centering) - self._points = points - self.set_dilation(dilation) - - @property - def mask(self) -> np.ndarray: - """ :class:`numpy.ndarray`: Overrides the default mask property, creating the processed - mask at first call and compressing it. The decompressed mask is returned from this - property. """ - return self.stored_mask - - def generate_mask(self, affine_matrix: np.ndarray, interpolator: int) -> None: - """ Generate the mask. - - Creates the mask applying any requested dilation and blurring and assigns compressed mask - to :attr:`_mask` - - Parameters - ---------- - affine_matrix: :class:`numpy.ndarray` - The transformation matrix required to transform the mask to the original frame. - interpolator, int: - The CV2 interpolator required to transform this mask to it's original frame - """ - mask = np.zeros((self.stored_size, self.stored_size, 1), dtype="float32") - for landmarks in self._points: - lms = np.rint(landmarks).astype("int") - cv2.fillConvexPoly(mask, cv2.convexHull(lms), [1.0], lineType=cv2.LINE_AA) - if self._dilation[-1] is not None: - self._dilate_mask(mask) - if self._blur_kernel != 0 and self._blur_type is not None: - mask = BlurMask(self._blur_type, - mask, - self._blur_kernel, - passes=self._blur_passes).blurred - logger.trace("mask: (shape: %s, dtype: %s)", # type:ignore[attr-defined] - mask.shape, mask.dtype) - self.add(mask, affine_matrix, interpolator) - - -class BlurMask(): - """ Factory class to return the correct blur object for requested blur type. - - Works for square images only. Currently supports Gaussian and Normalized Box Filters. - - Parameters - ---------- - blur_type: ["gaussian", "normalized"] - The type of blur to use - mask: :class:`numpy.ndarray` - The mask to apply the blur to - kernel: int or float - Either the kernel size (in pixels) or the size of the kernel as a ratio of mask size - is_ratio: bool, optional - Whether the given :attr:`kernel` parameter is a ratio or not. If ``True`` then the - actual kernel size will be calculated from the given ratio and the mask size. If - ``False`` then the kernel size will be set directly from the :attr:`kernel` parameter. - Default: ``False`` - passes: int, optional - The number of passes to perform when blurring. Default: ``1`` - - Example - ------- - >>> print(mask.shape) - (128, 128, 1) - >>> new_mask = BlurMask("gaussian", mask, 3, is_ratio=False, passes=1).blurred - >>> print(new_mask.shape) - (128, 128, 1) - """ - def __init__(self, - blur_type: T.Literal["gaussian", "normalized"], - mask: np.ndarray, - kernel: int | float, - is_ratio: bool = False, - passes: int = 1) -> None: - logger.trace("Initializing %s: (blur_type: '%s', " # type:ignore[attr-defined] - "mask_shape: %s, kernel: %s, is_ratio: %s, passes: %s)", - self.__class__.__name__, blur_type, - mask.shape, kernel, is_ratio, passes) - self._blur_type = blur_type - self._mask = mask - self._passes = passes - kernel_size = self._get_kernel_size(kernel, is_ratio) - self._kernel_size = self._get_kernel_tuple(kernel_size) - logger.trace("Initialized %s", self.__class__.__name__) # type:ignore[attr-defined] - - @property - def blurred(self) -> np.ndarray: - """ :class:`numpy.ndarray`: The final mask with blurring applied. """ - func = self._func_mapping[self._blur_type] - kwargs = self._get_kwargs() - blurred = self._mask - for i in range(self._passes): - assert isinstance(kwargs["ksize"], tuple) - ksize = int(kwargs["ksize"][0]) - logger.trace("Pass: %s, kernel_size: %s", # type:ignore[attr-defined] - i + 1, (ksize, ksize)) - blurred = func(blurred, **kwargs) - ksize = int(round(ksize * self._multipass_factor)) - kwargs["ksize"] = self._get_kernel_tuple(ksize) - blurred = blurred[..., None] - logger.trace("Returning blurred mask. Shape: %s", # type:ignore[attr-defined] - blurred.shape) - return blurred - - @property - def _multipass_factor(self) -> float: - """ For multiple passes the kernel must be scaled down. This value is - different for box filter and gaussian """ - factor = {"gaussian": 0.8, "normalized": 0.5} - return factor[self._blur_type] - - @property - def _sigma(self) -> T.Literal[0]: - """ int: The Sigma for Gaussian Blur. Returns 0 to force calculation from kernel size. """ - return 0 - - @property - def _func_mapping(self) -> dict[T.Literal["gaussian", "normalized"], Callable]: - """ dict: :attr:`_blur_type` mapped to cv2 Function name. """ - return {"gaussian": cv2.GaussianBlur, "normalized": cv2.blur} - - @property - def _kwarg_requirements(self) -> dict[T.Literal["gaussian", "normalized"], list[str]]: - """ dict: :attr:`_blur_type` mapped to cv2 Function required keyword arguments. """ - return {"gaussian": ['ksize', 'sigmaX'], "normalized": ['ksize']} - - @property - def _kwarg_mapping(self) -> dict[str, int | tuple[int, int]]: - """ dict: cv2 function keyword arguments mapped to their parameters. """ - return {"ksize": self._kernel_size, "sigmaX": self._sigma} - - def _get_kernel_size(self, kernel: int | float, is_ratio: bool) -> int: - """ Set the kernel size to absolute value. - - If :attr:`is_ratio` is ``True`` then the kernel size is calculated from the given ratio and - the :attr:`_mask` size, otherwise the given kernel size is just returned. - - Parameters - ---------- - kernel: int or float - Either the kernel size (in pixels) or the size of the kernel as a ratio of mask size - is_ratio: bool, optional - Whether the given :attr:`kernel` parameter is a ratio or not. If ``True`` then the - actual kernel size will be calculated from the given ratio and the mask size. If - ``False`` then the kernel size will be set directly from the :attr:`kernel` parameter. - - Returns - ------- - int - The size (in pixels) of the blur kernel - """ - if not is_ratio: - return int(kernel) - - mask_diameter = np.sqrt(np.sum(self._mask)) - radius = round(max(1., mask_diameter * kernel / 100.)) - kernel_size = int(radius * 2 + 1) - logger.trace("kernel_size: %s", kernel_size) # type:ignore[attr-defined] - return kernel_size - - @staticmethod - def _get_kernel_tuple(kernel_size: int) -> tuple[int, int]: - """ Make sure kernel_size is odd and return it as a tuple. - - Parameters - ---------- - kernel_size: int - The size in pixels of the blur kernel - - Returns - ------- - tuple - The kernel size as a tuple of ('int', 'int') - """ - kernel_size += 1 if kernel_size % 2 == 0 else 0 - retval = (kernel_size, kernel_size) - logger.trace(retval) # type:ignore[attr-defined] - return retval - - def _get_kwargs(self) -> dict[str, int | tuple[int, int]]: - """ dict: the valid keyword arguments for the requested :attr:`_blur_type` """ - retval = {kword: self._kwarg_mapping[kword] - for kword in self._kwarg_requirements[self._blur_type]} - logger.trace("BlurMask kwargs: %s", retval) # type:ignore[attr-defined] - return retval - - _HASHES_SEEN: dict[str, dict[str, int]] = {} diff --git a/tools/manual/frameviewer/editor/__init__.py b/tools/manual/frameviewer/editor/__init__.py index 8a7244abe4..7902b47227 100644 --- a/tools/manual/frameviewer/editor/__init__.py +++ b/tools/manual/frameviewer/editor/__init__.py @@ -1,8 +1,8 @@ #!/usr/bin/env python3 """ The Frame Viewer for Faceswap's Manual Tool. """ -from ._base import View # noqa -from .bounding_box import BoundingBox # noqa -from .extract_box import ExtractBox # noqa -from .landmarks import Landmarks, Mesh # noqa -from .mask import Mask # noqa +from ._base import View +from .bounding_box import BoundingBox +from .extract_box import ExtractBox +from .landmarks import Landmarks, Mesh +from .mask import Mask diff --git a/tools/manual/manual.py b/tools/manual/manual.py index c31683e3ee..46685766bc 100644 --- a/tools/manual/manual.py +++ b/tools/manual/manual.py @@ -27,8 +27,7 @@ from .thumbnails import ThumbsCreator if T.TYPE_CHECKING: - from lib.align import DetectedFace - from lib.align.detected_face import Mask + from lib.align import DetectedFace, Mask from lib.queue_manager import EventQueue logger = logging.getLogger(__name__) @@ -819,7 +818,7 @@ def get_masks(self, frame_index: int, face_index: int) -> dict[str, Mask]: Returns ------- - dict[str, :class:`~lib.align.detected_face.Mask`] + dict[str, :class:`~lib.align.aligned_mask.Mask`] The updated masks """ logger.trace("frame_index: %s, face_index: %s", # type:ignore[attr-defined]