OSSForks/faceswap
1
0
Fork 0
mirror of https://github.com/zebrajr/faceswap.git synced 2025-12-06 00:20:09 +01:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add mask dilation/erosion option for training

Browse Source
This commit is contained in:
torzdf 2024-03-26 18:06:48 +00:00
parent 2b08b8d2a6
commit 4502824481
6 changed files with 156 additions and 78 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

169
lib/align/detected_face.py
View File

@ -88,8 +88,8 @@ class DetectedFace():
landmarks_xy: np.ndarray | 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: filename: str | None = None) -> None:
logger.trace("Initializing %s: (image: %s, left: %s, width: %s, top: %s, " # type: ignore logger.trace("Initializing %s: (image: %s, left: %s, " # type:ignore[attr-defined]
"height: %s, landmarks_xy: %s, mask: %s, filename: %s)", "width: %s, top: %s, height: %s, landmarks_xy: %s, mask: %s, filename: %s)",
self.__class__.__name__, self.__class__.__name__,
image.shape if image is not None and image.any() else image, left, width, top, image.shape if image is not None and image.any() else image, left, width, top,
height, landmarks_xy, mask, filename) height, landmarks_xy, mask, filename)
@ -105,7 +105,7 @@ class DetectedFace():
self._training_masks: tuple[bytes, tuple[int, int, int]] | None = None self._training_masks: tuple[bytes, tuple[int, int, int]] | None = None
self._aligned: AlignedFace | None = None self._aligned: AlignedFace | None = None
logger.trace("Initialized %s", self.__class__.__name__) # type: ignore logger.trace("Initialized %s", self.__class__.__name__) # type:ignore[attr-defined]
@property @property
def aligned(self) -> AlignedFace: def aligned(self) -> AlignedFace:
@ -167,7 +167,7 @@ class DetectedFace():
The centering to store the mask at. One of `"legacy"`, `"face"`, `"head"`. The centering to store the mask at. One of `"legacy"`, `"face"`, `"head"`.
Default: `"face"` Default: `"face"`
""" """
logger.trace("name: '%s', mask shape: %s, affine_matrix: %s, " # type: ignore logger.trace("name: '%s', mask shape: %s, affine_matrix: %s, " # type:ignore[attr-defined]
"interpolator: %s, storage_size: %s, storage_centering: %s)", name, "interpolator: %s, storage_size: %s, storage_centering: %s)", name,
mask.shape, affine_matrix, interpolator, storage_size, storage_centering) mask.shape, affine_matrix, interpolator, storage_size, storage_centering)
fsmask = Mask(storage_size=storage_size, storage_centering=storage_centering) fsmask = Mask(storage_size=storage_size, storage_centering=storage_centering)
@ -183,7 +183,7 @@ class DetectedFace():
landmarks: :class:`numpy.ndarray` landmarks: :class:`numpy.ndarray`
The 68 point face landmarks to add for the face The 68 point face landmarks to add for the face
""" """
logger.trace("landmarks shape: '%s'", landmarks.shape) # type: ignore logger.trace("landmarks shape: '%s'", landmarks.shape) # type:ignore[attr-defined]
self._landmarks_xy = landmarks self._landmarks_xy = landmarks
def add_identity(self, name: str, embedding: np.ndarray, ) -> None: def add_identity(self, name: str, embedding: np.ndarray, ) -> None:
@ -197,7 +197,7 @@ class DetectedFace():
embedding: numpy.ndarray embedding: numpy.ndarray
The identity embedding The identity embedding
""" """
logger.trace("name: '%s', embedding shape: %s", # type: ignore logger.trace("name: '%s', embedding shape: %s", # type:ignore[attr-defined]
name, embedding.shape) name, embedding.shape)
assert name == "vggface2" assert name == "vggface2"
assert embedding.shape[0] == 512 assert embedding.shape[0] == 512
@ -210,7 +210,7 @@ class DetectedFace():
def get_landmark_mask(self, def get_landmark_mask(self,
area: T.Literal["eye", "face", "mouth"], area: T.Literal["eye", "face", "mouth"],
blur_kernel: int, blur_kernel: int,
dilation: int) -> np.ndarray: dilation: float) -> np.ndarray:
""" Add a :class:`LandmarksMask` to this detected face """ Add a :class:`LandmarksMask` to this detected face
Landmark based masks are generated from face Aligned Face landmark points. An aligned Landmark based masks are generated from face Aligned Face landmark points. An aligned
@ -224,8 +224,8 @@ class DetectedFace():
specific areas specific areas
blur_kernel: int blur_kernel: int
The size of the kernel for blurring the mask edges The size of the kernel for blurring the mask edges
dilation: int dilation: float
The amount of dilation to apply to the mask. `0` for none. Default: `0` The amount of dilation to apply to the mask. as a percentage of the mask size
Returns Returns
------- -------
@ -233,7 +233,7 @@ class DetectedFace():
The generated landmarks mask for the selected area The generated landmarks mask for the selected area
""" """
# TODO Face mask generation from landmarks # TODO Face mask generation from landmarks
logger.trace("area: %s, dilation: %s", area, dilation) # type: ignore logger.trace("area: %s, dilation: %s", area, dilation) # type:ignore[attr-defined]
areas = {"mouth": [slice(48, 60)], "eye": [slice(36, 42), slice(42, 48)]} areas = {"mouth": [slice(48, 60)], "eye": [slice(36, 42), slice(42, 48)]}
points = [self.aligned.landmarks[zone] points = [self.aligned.landmarks[zone]
for zone in areas[area]] for zone in areas[area]]
@ -307,7 +307,7 @@ class DetectedFace():
for name, mask in self.mask.items()}, for name, mask in self.mask.items()},
identity={k: v.tolist() for k, v in self._identity.items()}, identity={k: v.tolist() for k, v in self._identity.items()},
thumb=self.thumbnail) thumb=self.thumbnail)
logger.trace("Returning: %s", alignment) # type: ignore logger.trace("Returning: %s", alignment) # type:ignore[attr-defined]
return alignment return alignment
def from_alignment(self, alignment: AlignmentFileDict, def from_alignment(self, alignment: AlignmentFileDict,
@ -332,8 +332,8 @@ class DetectedFace():
Default: ``False`` Default: ``False``
""" """
logger.trace("Creating from alignment: (alignment: %s, has_image: %s)", # type: ignore logger.trace("Creating from alignment: (alignment: %s," # type:ignore[attr-defined]
alignment, bool(image is not None)) " has_image: %s)", alignment, bool(image is not None))
self.left = alignment["x"] self.left = alignment["x"]
self.width = alignment["w"] self.width = alignment["w"]
self.top = alignment["y"] self.top = alignment["y"]
@ -358,9 +358,9 @@ class DetectedFace():
self.mask[name].from_dict(mask_dict) self.mask[name].from_dict(mask_dict)
if image is not None and image.any(): if image is not None and image.any():
self._image_to_face(image) self._image_to_face(image)
logger.trace("Created from alignment: (left: %s, width: %s, top: %s, " # type: ignore logger.trace("Created from alignment: (left: %s, width: %s, " # type:ignore[attr-defined]
"height: %s, landmarks: %s, mask: %s)", self.left, self.width, self.top, "top: %s, height: %s, landmarks: %s, mask: %s)",
self.height, self.landmarks_xy, self.mask) self.left, self.width, self.top, self.height, self.landmarks_xy, self.mask)
def to_png_meta(self) -> PNGHeaderAlignmentsDict: def to_png_meta(self) -> PNGHeaderAlignmentsDict:
""" Return the detected face formatted for insertion into a png itxt header. """ Return the detected face formatted for insertion into a png itxt header.
@ -403,14 +403,14 @@ class DetectedFace():
for key, val in alignment.get("identity", {}).items(): for key, val in alignment.get("identity", {}).items():
assert key in ["vggface2"] assert key in ["vggface2"]
self._identity[T.cast(T.Literal["vggface2"], key)] = np.array(val, dtype="float32") self._identity[T.cast(T.Literal["vggface2"], key)] = np.array(val, dtype="float32")
logger.trace("Created from png exif header: (left: %s, width: %s, top: %s " # type: ignore logger.trace("Created from png exif header: (left: %s, " # type:ignore[attr-defined]
" height: %s, landmarks: %s, mask: %s, identity: %s)", self.left, self.width, "width: %s, top: %s height: %s, landmarks: %s, mask: %s, identity: %s)",
self.top, self.height, self.landmarks_xy, self.mask, self.left, self.width, self.top, self.height, self.landmarks_xy, self.mask,
{k: v.shape for k, v in self._identity.items()}) {k: v.shape for k, v in self._identity.items()})
def _image_to_face(self, image: np.ndarray) -> None: def _image_to_face(self, image: np.ndarray) -> None:
""" set self.image to be the cropped face from detected bounding box """ """ set self.image to be the cropped face from detected bounding box """
logger.trace("Cropping face from image") # type: ignore logger.trace("Cropping face from image") # type:ignore[attr-defined]
self.image = image[self.top: self.bottom, self.image = image[self.top: self.bottom,
self.left: self.right] self.left: self.right]
@ -467,10 +467,11 @@ class DetectedFace():
""" """
if self._aligned and not force: if self._aligned and not force:
# Don't reload an already aligned face # Don't reload an already aligned face
logger.trace("Skipping alignment calculation for already aligned face") # type: ignore logger.trace("Skipping alignment calculation for already " # type:ignore[attr-defined]
"aligned face")
else: else:
logger.trace("Loading aligned face: (size: %s, dtype: %s)", # type: ignore logger.trace("Loading aligned face: (size: %s, " # type:ignore[attr-defined]
size, dtype) "dtype: %s)", size, dtype)
self._aligned = AlignedFace(self.landmarks_xy, self._aligned = AlignedFace(self.landmarks_xy,
image=image, image=image,
centering=centering, centering=centering,
@ -507,7 +508,8 @@ class Mask():
def __init__(self, def __init__(self,
storage_size: int = 128, storage_size: int = 128,
storage_centering: CenteringType = "face") -> None: storage_centering: CenteringType = "face") -> None:
logger.trace("Initializing: %s (storage_size: %s, storage_centering: %s)", # type: ignore logger.trace("Initializing: %s (storage_size: %s, " # type:ignore[attr-defined]
"storage_centering: %s)",
self.__class__.__name__, storage_size, storage_centering) self.__class__.__name__, storage_size, storage_centering)
self.stored_size = storage_size self.stored_size = storage_size
self.stored_centering = storage_centering self.stored_centering = storage_centering
@ -520,19 +522,21 @@ class Mask():
self._blur_passes: int = 0 self._blur_passes: int = 0
self._blur_kernel: float | int = 0 self._blur_kernel: float | int = 0
self._threshold = 0.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_size = 0
self._sub_crop_slices: dict[T.Literal["in", "out"], list[slice]] = {} self._sub_crop_slices: dict[T.Literal["in", "out"], list[slice]] = {}
self.set_blur_and_threshold() self.set_blur_and_threshold()
logger.trace("Initialized: %s", self.__class__.__name__) # type: ignore logger.trace("Initialized: %s", self.__class__.__name__) # type:ignore[attr-defined]
@property @property
def mask(self) -> np.ndarray: def mask(self) -> np.ndarray:
""" :class:`numpy.ndarray`: The mask at the size of :attr:`stored_size` with any requested """ :class:`numpy.ndarray`: The mask at the size of :attr:`stored_size` with any requested
blurring, threshold amount and centering applied.""" blurring, threshold amount and centering applied."""
mask = self.stored_mask mask = self.stored_mask
if self._threshold != 0.0 or self._blur_kernel != 0: if self._dilation[-1] is not None or self._threshold != 0.0 or self._blur_kernel != 0:
mask = mask.copy() mask = mask.copy()
self._dilate_mask(mask)
if self._threshold != 0.0: if self._threshold != 0.0:
mask[mask < self._threshold] = 0.0 mask[mask < self._threshold] = 0.0
mask[mask > 255.0 - self._threshold] = 255.0 mask[mask > 255.0 - self._threshold] = 255.0
@ -546,7 +550,7 @@ class Mask():
slice_in, slice_out = self._sub_crop_slices["in"], self._sub_crop_slices["out"] 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], :] out[slice_out[0], slice_out[1], :] = mask[slice_in[0], slice_in[1], :]
mask = out mask = out
logger.trace("mask shape: %s", mask.shape) # type: ignore logger.trace("mask shape: %s", mask.shape) # type:ignore[attr-defined]
return mask return mask
@property @property
@ -556,7 +560,7 @@ class Mask():
assert self._mask is not None assert self._mask is not None
dims = (self.stored_size, self.stored_size, 1) dims = (self.stored_size, self.stored_size, 1)
mask = np.frombuffer(decompress(self._mask), dtype="uint8").reshape(dims) mask = np.frombuffer(decompress(self._mask), dtype="uint8").reshape(dims)
logger.trace("stored mask shape: %s", mask.shape) # type: ignore logger.trace("stored mask shape: %s", mask.shape) # type:ignore[attr-defined]
return mask return mask
@property @property
@ -567,9 +571,9 @@ class Mask():
[0, self.stored_size - 1], [0, self.stored_size - 1],
[self.stored_size - 1, self.stored_size - 1], [self.stored_size - 1, self.stored_size - 1],
[self.stored_size - 1, 0]], np.int32).reshape((-1, 1, 2)) [self.stored_size - 1, 0]], np.int32).reshape((-1, 1, 2))
matrix = cv2.invertAffineTransform(self._affine_matrix) matrix = cv2.invertAffineTransform(self.affine_matrix)
roi = cv2.transform(points, matrix).reshape((4, 2)) roi = cv2.transform(points, matrix).reshape((4, 2))
logger.trace("Returning: %s", roi) # type: ignore logger.trace("Returning: %s", roi) # type:ignore[attr-defined]
return roi return roi
@property @property
@ -584,6 +588,22 @@ class Mask():
assert self._interpolator is not None assert self._interpolator is not None
return self._interpolator 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: 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 """ Return the stored mask in a full size frame of the given dimensions
@ -600,13 +620,13 @@ class Mask():
""" """
frame = np.zeros((width, height, 1), dtype="uint8") frame = np.zeros((width, height, 1), dtype="uint8")
mask = cv2.warpAffine(self.mask, mask = cv2.warpAffine(self.mask,
self._affine_matrix, self.affine_matrix,
(width, height), (width, height),
frame, frame,
flags=cv2.WARP_INVERSE_MAP | self._interpolator, flags=cv2.WARP_INVERSE_MAP | self.interpolator,
borderMode=cv2.BORDER_CONSTANT) borderMode=cv2.BORDER_CONSTANT)
logger.trace("mask shape: %s, mask dtype: %s, mask min: %s, mask max: %s", # type: ignore logger.trace("mask shape: %s, mask dtype: %s, mask min: %s, " # type:ignore[attr-defined]
mask.shape, mask.dtype, mask.min(), mask.max()) "mask max: %s", mask.shape, mask.dtype, mask.min(), mask.max())
return mask return mask
def add(self, mask: np.ndarray, affine_matrix: np.ndarray, interpolator: int) -> None: def add(self, mask: np.ndarray, affine_matrix: np.ndarray, interpolator: int) -> None:
@ -624,9 +644,9 @@ class Mask():
interpolator, int: interpolator, int:
The CV2 interpolator required to transform this mask to it's original frame The CV2 interpolator required to transform this mask to it's original frame
""" """
logger.trace("mask shape: %s, mask dtype: %s, mask min: %s, mask max: %s, " # type: ignore logger.trace("mask shape: %s, mask dtype: %s, mask min: %s, " # type:ignore[attr-defined]
"affine_matrix: %s, interpolator: %s)", mask.shape, mask.dtype, mask.min(), "mask max: %s, affine_matrix: %s, interpolator: %s)",
affine_matrix, mask.max(), interpolator) mask.shape, mask.dtype, mask.min(), affine_matrix, mask.max(), interpolator)
self._affine_matrix = self._adjust_affine_matrix(mask.shape[0], affine_matrix) self._affine_matrix = self._adjust_affine_matrix(mask.shape[0], affine_matrix)
self._interpolator = interpolator self._interpolator = interpolator
self.replace_mask(mask) self.replace_mask(mask)
@ -645,6 +665,26 @@ class Mask():
interpolation=cv2.INTER_AREA) * 255.0).astype("uint8") interpolation=cv2.INTER_AREA) * 255.0).astype("uint8")
self._mask = compress(mask.tobytes()) 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, def set_blur_and_threshold(self,
blur_kernel: int = 0, blur_kernel: int = 0,
blur_type: T.Literal["gaussian", "normalized"] | None = "gaussian", blur_type: T.Literal["gaussian", "normalized"] | None = "gaussian",
@ -666,8 +706,9 @@ class Mask():
The threshold amount to minimize/maximize mask values to 0 and 100. Percentage value. The threshold amount to minimize/maximize mask values to 0 and 100. Percentage value.
Default: 0 Default: 0
""" """
logger.trace("blur_kernel: %s, blur_type: %s, blur_passes: %s, " # type: ignore logger.trace("blur_kernel: %s, blur_type: %s, " # type:ignore[attr-defined]
"threshold: %s", blur_kernel, blur_type, blur_passes, threshold) "blur_passes: %s, threshold: %s",
blur_kernel, blur_type, blur_passes, threshold)
if blur_type is not None: if blur_type is not None:
blur_kernel += 0 if blur_kernel == 0 or blur_kernel % 2 == 1 else 1 blur_kernel += 0 if blur_kernel == 0 or blur_kernel % 2 == 1 else 1
self._blur_kernel = blur_kernel self._blur_kernel = blur_kernel
@ -719,9 +760,9 @@ class Mask():
slice(max(roi[0] * -1, 0), slice(max(roi[0] * -1, 0),
crop_size - min(crop_size, max(0, roi[2] - self.stored_size)))] crop_size - min(crop_size, max(0, roi[2] - self.stored_size)))]
logger.trace("src_size: %s, coverage_ratio: %s, sub_crop_size: %s, " # type: ignore logger.trace("src_size: %s, coverage_ratio: %s, " # type:ignore[attr-defined]
"sub_crop_slices: %s", roi, coverage_ratio, self._sub_crop_size, "sub_crop_size: %s, sub_crop_slices: %s",
self._sub_crop_slices) 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: def _adjust_affine_matrix(self, mask_size: int, affine_matrix: np.ndarray) -> np.ndarray:
""" Adjust the affine matrix for the mask's storage size """ Adjust the affine matrix for the mask's storage size
@ -741,7 +782,7 @@ class Mask():
zoom = self.stored_size / mask_size zoom = self.stored_size / mask_size
zoom_mat = np.array([[zoom, 0, 0.], [0, zoom, 0.]]) 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.]])))) 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 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, "original matrix: %s, adjusted_matrix: %s", self.stored_size, mask_size, zoom,
affine_matrix.shape, adjust_mat.shape) affine_matrix.shape, adjust_mat.shape)
return adjust_mat return adjust_mat
@ -768,7 +809,8 @@ class Mask():
interpolator=self.interpolator, interpolator=self.interpolator,
stored_size=self.stored_size, stored_size=self.stored_size,
stored_centering=self.stored_centering) stored_centering=self.stored_centering)
logger.trace({k: v if k != "mask" else type(v) for k, v in retval.items()}) # type: ignore logger.trace({k: v if k != "mask" else type(v) # type:ignore[attr-defined]
for k, v in retval.items()})
return retval return retval
def to_png_meta(self) -> MaskAlignmentsFileDict: def to_png_meta(self) -> MaskAlignmentsFileDict:
@ -799,7 +841,7 @@ class Mask():
self.stored_size = mask_dict["stored_size"] self.stored_size = mask_dict["stored_size"]
centering = mask_dict.get("stored_centering") centering = mask_dict.get("stored_centering")
self.stored_centering = "face" if centering is None else centering self.stored_centering = "face" if centering is None else centering
logger.trace({k: v if k != "mask" else type(v) # type: ignore logger.trace({k: v if k != "mask" else type(v) # type:ignore[attr-defined]
for k, v in mask_dict.items()}) for k, v in mask_dict.items()})
@ -826,17 +868,17 @@ class LandmarksMask(Mask):
storage_centering, str (optional): storage_centering, str (optional):
The centering to store the mask at. One of `"legacy"`, `"face"`, `"head"`. The centering to store the mask at. One of `"legacy"`, `"face"`, `"head"`.
Default: `"face"` Default: `"face"`
dilation: int, optional dilation: float, optional
The amount of dilation to apply to the mask. `0` for none. Default: `0` The amount of dilation to apply to the mask. as a percentage of the mask size. Default: 0.0
""" """
def __init__(self, def __init__(self,
points: list[np.ndarray], points: list[np.ndarray],
storage_size: int = 128, storage_size: int = 128,
storage_centering: CenteringType = "face", storage_centering: CenteringType = "face",
dilation: int = 0) -> None: dilation: float = 0.0) -> None:
super().__init__(storage_size=storage_size, storage_centering=storage_centering) super().__init__(storage_size=storage_size, storage_centering=storage_centering)
self._points = points self._points = points
self._dilation = dilation self.set_dilation(dilation)
@property @property
def mask(self) -> np.ndarray: def mask(self) -> np.ndarray:
@ -862,17 +904,15 @@ class LandmarksMask(Mask):
for landmarks in self._points: for landmarks in self._points:
lms = np.rint(landmarks).astype("int") lms = np.rint(landmarks).astype("int")
cv2.fillConvexPoly(mask, cv2.convexHull(lms), 1.0, lineType=cv2.LINE_AA) cv2.fillConvexPoly(mask, cv2.convexHull(lms), 1.0, lineType=cv2.LINE_AA)
if self._dilation != 0: if self._dilation[-1] is not None:
mask = cv2.dilate(mask, self._dilate_mask(mask)
cv2.getStructuringElement(cv2.MORPH_ELLIPSE,
(self._dilation, self._dilation)),
iterations=1)
if self._blur_kernel != 0 and self._blur_type is not None: if self._blur_kernel != 0 and self._blur_type is not None:
mask = BlurMask(self._blur_type, mask = BlurMask(self._blur_type,
mask, mask,
self._blur_kernel, self._blur_kernel,
passes=self._blur_passes).blurred passes=self._blur_passes).blurred
logger.trace("mask: (shape: %s, dtype: %s)", mask.shape, mask.dtype) # type: ignore logger.trace("mask: (shape: %s, dtype: %s)", # type:ignore[attr-defined]
mask.shape, mask.dtype)
self.add(mask, affine_matrix, interpolator) self.add(mask, affine_matrix, interpolator)
@ -911,15 +951,16 @@ class BlurMask(): # pylint:disable=too-few-public-methods
kernel: int | float, kernel: int | float,
is_ratio: bool = False, is_ratio: bool = False,
passes: int = 1) -> None: passes: int = 1) -> None:
logger.trace("Initializing %s: (blur_type: '%s', mask_shape: %s, " # type: ignore logger.trace("Initializing %s: (blur_type: '%s', " # type:ignore[attr-defined]
"kernel: %s, is_ratio: %s, passes: %s)", self.__class__.__name__, blur_type, "mask_shape: %s, kernel: %s, is_ratio: %s, passes: %s)",
self.__class__.__name__, blur_type,
mask.shape, kernel, is_ratio, passes) mask.shape, kernel, is_ratio, passes)
self._blur_type = blur_type self._blur_type = blur_type
self._mask = mask self._mask = mask
self._passes = passes self._passes = passes
kernel_size = self._get_kernel_size(kernel, is_ratio) kernel_size = self._get_kernel_size(kernel, is_ratio)
self._kernel_size = self._get_kernel_tuple(kernel_size) self._kernel_size = self._get_kernel_tuple(kernel_size)
logger.trace("Initialized %s", self.__class__.__name__) # type: ignore logger.trace("Initialized %s", self.__class__.__name__) # type:ignore[attr-defined]
@property @property
def blurred(self) -> np.ndarray: def blurred(self) -> np.ndarray:
@ -930,12 +971,14 @@ class BlurMask(): # pylint:disable=too-few-public-methods
for i in range(self._passes): for i in range(self._passes):
assert isinstance(kwargs["ksize"], tuple) assert isinstance(kwargs["ksize"], tuple)
ksize = int(kwargs["ksize"][0]) ksize = int(kwargs["ksize"][0])
logger.trace("Pass: %s, kernel_size: %s", i + 1, (ksize, ksize)) # type: ignore logger.trace("Pass: %s, kernel_size: %s", # type:ignore[attr-defined]
i + 1, (ksize, ksize))
blurred = func(blurred, **kwargs) blurred = func(blurred, **kwargs)
ksize = int(round(ksize * self._multipass_factor)) ksize = int(round(ksize * self._multipass_factor))
kwargs["ksize"] = self._get_kernel_tuple(ksize) kwargs["ksize"] = self._get_kernel_tuple(ksize)
blurred = blurred[..., None] blurred = blurred[..., None]
logger.trace("Returning blurred mask. Shape: %s", blurred.shape) # type: ignore logger.trace("Returning blurred mask. Shape: %s", # type:ignore[attr-defined]
blurred.shape)
return blurred return blurred
@property @property
@ -991,7 +1034,7 @@ class BlurMask(): # pylint:disable=too-few-public-methods
mask_diameter = np.sqrt(np.sum(self._mask)) mask_diameter = np.sqrt(np.sum(self._mask))
radius = round(max(1., mask_diameter * kernel / 100.)) radius = round(max(1., mask_diameter * kernel / 100.))
kernel_size = int(radius * 2 + 1) kernel_size = int(radius * 2 + 1)
logger.trace("kernel_size: %s", kernel_size) # type: ignore logger.trace("kernel_size: %s", kernel_size) # type:ignore[attr-defined]
return kernel_size return kernel_size
@staticmethod @staticmethod
@ -1010,14 +1053,14 @@ class BlurMask(): # pylint:disable=too-few-public-methods
""" """
kernel_size += 1 if kernel_size % 2 == 0 else 0 kernel_size += 1 if kernel_size % 2 == 0 else 0
retval = (kernel_size, kernel_size) retval = (kernel_size, kernel_size)
logger.trace(retval) # type: ignore logger.trace(retval) # type:ignore[attr-defined]
return retval return retval
def _get_kwargs(self) -> dict[str, int | tuple[int, int]]: def _get_kwargs(self) -> dict[str, int | tuple[int, int]]:
""" dict: the valid keyword arguments for the requested :attr:`_blur_type` """ """ dict: the valid keyword arguments for the requested :attr:`_blur_type` """
retval = {kword: self._kwarg_mapping[kword] retval = {kword: self._kwarg_mapping[kword]
for kword in self._kwarg_requirements[self._blur_type]} for kword in self._kwarg_requirements[self._blur_type]}
logger.trace("BlurMask kwargs: %s", retval) # type: ignore logger.trace("BlurMask kwargs: %s", retval) # type:ignore[attr-defined]
return retval return retval

4
lib/training/cache.py
View File

@ -428,8 +428,10 @@ class _Cache():
f"The masks that exist for this face are: {list(detected_face.mask)}") f"The masks that exist for this face are: {list(detected_face.mask)}")
mask = detected_face.mask[str(self._config["mask_type"])] mask = detected_face.mask[str(self._config["mask_type"])]
assert isinstance(self._config["mask_dilation"], float)
assert isinstance(self._config["mask_blur_kernel"], int) assert isinstance(self._config["mask_blur_kernel"], int)
assert isinstance(self._config["mask_threshold"], int) assert isinstance(self._config["mask_threshold"], int)
mask.set_dilation(self._config["mask_dilation"])
mask.set_blur_and_threshold(blur_kernel=self._config["mask_blur_kernel"], mask.set_blur_and_threshold(blur_kernel=self._config["mask_blur_kernel"],
threshold=self._config["mask_threshold"]) threshold=self._config["mask_threshold"])
@ -467,7 +469,7 @@ class _Cache():
assert isinstance(multiplier, int) assert isinstance(multiplier, int)
if not self._config["penalized_mask_loss"] or multiplier <= 1: if not self._config["penalized_mask_loss"] or multiplier <= 1:
return None return None
mask = detected_face.get_landmark_mask(area, self._size // 16, self._size // 32) mask = detected_face.get_landmark_mask(area, self._size // 16, 2.5)
logger.trace("Caching localized '%s' mask for: %s %s", # type: ignore logger.trace("Caching localized '%s' mask for: %s %s", # type: ignore
area, filename, mask.shape) area, filename, mask.shape)
return mask return mask

20
locales/plugins.train._config.pot
View File

@ -8,7 +8,7 @@ msgid ""
msgstr "" msgstr ""
"Project-Id-Version: PACKAGE VERSION\n" "Project-Id-Version: PACKAGE VERSION\n"
"Report-Msgid-Bugs-To: \n" "Report-Msgid-Bugs-To: \n"
"POT-Creation-Date: 2023-08-20 14:54+0100\n" "POT-Creation-Date: 2024-03-26 17:37+0000\n"
"PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n" "PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n"
"Last-Translator: FULL NAME <EMAIL@ADDRESS>\n" "Last-Translator: FULL NAME <EMAIL@ADDRESS>\n"
"Language-Team: LANGUAGE <LL@li.org>\n" "Language-Team: LANGUAGE <LL@li.org>\n"
@ -539,8 +539,9 @@ msgid ""
"attention on the core face area." "attention on the core face area."
msgstr "" msgstr ""
#: plugins/train/_config.py:600 plugins/train/_config.py:642 #: plugins/train/_config.py:600 plugins/train/_config.py:643
#: plugins/train/_config.py:656 plugins/train/_config.py:665 #: plugins/train/_config.py:656 plugins/train/_config.py:671
#: plugins/train/_config.py:680
msgid "mask" msgid "mask"
msgstr "" msgstr ""
@ -582,7 +583,14 @@ msgid ""
"performance." "performance."
msgstr "" msgstr ""
#: plugins/train/_config.py:644 #: plugins/train/_config.py:645
msgid ""
"Dilate or erode the mask. Negative values erode the mask (make it smaller). "
"Positive values dilate the mask (make it larger). The value given is a "
"percentage of the total mask size."
msgstr ""
#: plugins/train/_config.py:658
msgid "" msgid ""
"Apply gaussian blur to the mask input. This has the effect of smoothing the " "Apply gaussian blur to the mask input. This has the effect of smoothing the "
"edges of the mask, which can help with poorly calculated masks and give less " "edges of the mask, which can help with poorly calculated masks and give less "
@ -592,13 +600,13 @@ msgid ""
"number." "number."
msgstr "" msgstr ""
#: plugins/train/_config.py:658 #: plugins/train/_config.py:673
msgid "" msgid ""
"Sets pixels that are near white to white and near black to black. Set to 0 " "Sets pixels that are near white to white and near black to black. Set to 0 "
"for off." "for off."
msgstr "" msgstr ""
#: plugins/train/_config.py:667 #: plugins/train/_config.py:682
msgid "" msgid ""
"Dedicate a portion of the model to learning how to duplicate the input mask. " "Dedicate a portion of the model to learning how to duplicate the input mask. "
"Increases VRAM usage in exchange for learning a quick ability to try to " "Increases VRAM usage in exchange for learning a quick ability to try to "

BIN
locales/ru/LC_MESSAGES/plugins.train._config.mo
View File

Binary file not shown.

26
locales/ru/LC_MESSAGES/plugins.train._config.po
View File

@ -7,15 +7,15 @@ msgid ""
msgstr "" msgstr ""
"Project-Id-Version: \n" "Project-Id-Version: \n"
"Report-Msgid-Bugs-To: \n" "Report-Msgid-Bugs-To: \n"
"POT-Creation-Date: 2023-08-20 14:54+0100\n" "POT-Creation-Date: 2024-03-26 17:37+0000\n"
"PO-Revision-Date: 2023-08-20 14:58+0100\n" "PO-Revision-Date: 2024-03-26 17:40+0000\n"
"Last-Translator: \n" "Last-Translator: \n"
"Language-Team: \n" "Language-Team: \n"
"Language: ru_RU\n" "Language: ru_RU\n"
"MIME-Version: 1.0\n" "MIME-Version: 1.0\n"
"Content-Type: text/plain; charset=UTF-8\n" "Content-Type: text/plain; charset=UTF-8\n"
"Content-Transfer-Encoding: 8bit\n" "Content-Transfer-Encoding: 8bit\n"
"X-Generator: Poedit 3.3.2\n" "X-Generator: Poedit 3.4.2\n"
#: plugins/train/_config.py:17 #: plugins/train/_config.py:17
msgid "" msgid ""
@ -851,8 +851,9 @@ msgstr ""
"время как область лица с маской является приоритетной. Может повысить общее " "время как область лица с маской является приоритетной. Может повысить общее "
"качество за счет концентрации внимания на основной области лица." "качество за счет концентрации внимания на основной области лица."
#: plugins/train/_config.py:600 plugins/train/_config.py:642 #: plugins/train/_config.py:600 plugins/train/_config.py:643
#: plugins/train/_config.py:656 plugins/train/_config.py:665 #: plugins/train/_config.py:656 plugins/train/_config.py:671
#: plugins/train/_config.py:680
msgid "mask" msgid "mask"
msgstr "маска" msgstr "маска"
@ -928,7 +929,16 @@ msgstr ""
"сообщества и для дальнейшего описания нуждается в тестировании. Профильные " "сообщества и для дальнейшего описания нуждается в тестировании. Профильные "
"лица могут иметь низкую производительность." "лица могут иметь низкую производительность."
#: plugins/train/_config.py:644 #: plugins/train/_config.py:645
msgid ""
"Dilate or erode the mask. Negative values erode the mask (make it smaller). "
"Positive values dilate the mask (make it larger). The value given is a "
"percentage of the total mask size."
msgstr ""
"Расширяет или сужает маску. Отрицательные значения сужают маску (делают её "
"меньше). Положительные значения расширяют маску (делают её больше). "
#: plugins/train/_config.py:658
msgid "" msgid ""
"Apply gaussian blur to the mask input. This has the effect of smoothing the " "Apply gaussian blur to the mask input. This has the effect of smoothing the "
"edges of the mask, which can help with poorly calculated masks and give less " "edges of the mask, which can help with poorly calculated masks and give less "
@ -944,7 +954,7 @@ msgstr ""
"должно быть нечетным, если передано четное число, то оно будет округлено до " "должно быть нечетным, если передано четное число, то оно будет округлено до "
"следующего нечетного числа." "следующего нечетного числа."
#: plugins/train/_config.py:658 #: plugins/train/_config.py:673
msgid "" msgid ""
"Sets pixels that are near white to white and near black to black. Set to 0 " "Sets pixels that are near white to white and near black to black. Set to 0 "
"for off." "for off."
@ -952,7 +962,7 @@ msgstr ""
"Устанавливает пиксели, которые почти белые - в белые и которые почти черные " "Устанавливает пиксели, которые почти белые - в белые и которые почти черные "
"- в черные. Установите 0, чтобы выключить." "- в черные. Установите 0, чтобы выключить."
#: plugins/train/_config.py:667 #: plugins/train/_config.py:682
msgid "" msgid ""
"Dedicate a portion of the model to learning how to duplicate the input mask. " "Dedicate a portion of the model to learning how to duplicate the input mask. "
"Increases VRAM usage in exchange for learning a quick ability to try to " "Increases VRAM usage in exchange for learning a quick ability to try to "

15
plugins/train/_config.py
View File

@ -632,6 +632,19 @@ class Config(FaceswapConfig):
"faces. The mask model has been trained by community members and will need " "faces. The mask model has been trained by community members and will need "
"testing for further description. Profile faces may result in sub-par " "testing for further description. Profile faces may result in sub-par "
"performance.")) "performance."))
self.add_item(
section=section,
title="mask_dilation",
datatype=float,
min_max=(-5.0, 5.0),
rounding=1,
default=0,
fixed=False,
group=_("mask"),
info=_(
"Dilate or erode the mask. Negative values erode the mask (make it smaller). "
"Positive values dilate the mask (make it larger). The value given is a "
"percentage of the total mask size."))
self.add_item( self.add_item(
section=section, section=section,
title="mask_blur_kernel", title="mask_blur_kernel",
@ -639,6 +652,7 @@ class Config(FaceswapConfig):
min_max=(0, 9), min_max=(0, 9),
rounding=1, rounding=1,
default=3, default=3,
fixed=False,
group=_("mask"), group=_("mask"),
info=_( info=_(
"Apply gaussian blur to the mask input. This has the effect of smoothing the " "Apply gaussian blur to the mask input. This has the effect of smoothing the "
@ -653,6 +667,7 @@ class Config(FaceswapConfig):
default=4, default=4,
min_max=(0, 50), min_max=(0, 50),
rounding=1, rounding=1,
fixed=False,
group=_("mask"), group=_("mask"),
info=_( info=_(
"Sets pixels that are near white to white and near black to black. Set to 0 for " "Sets pixels that are near white to white and near black to black. Set to 0 for "