1010import cv2
1111import numpy as np
1212
13- from lib .logger import parse_class_init
13+ from lib .logger import format_array , parse_class_init
1414from lib .utils import FaceswapError , get_module_objects
1515
1616from .aligned_utils import get_adjusted_center , get_sub_crop_size
@@ -423,7 +423,7 @@ def from_dict(self, mask: MaskAlignmentsFile) -> T.Self:
423423 return self
424424
425425
426- class LandmarksMask (Mask ):
426+ class LandmarksMask ():
427427 """Create a single channel mask from aligned landmark points.
428428
429429 Landmarks masks are created on the fly, so the stored centering and size should be the same as
@@ -444,36 +444,59 @@ class LandmarksMask(Mask):
444444 The type of landmarks that this mask is being created from
445445 landmarks
446446 The landmarks to generate the mask from
447- affine_matrix
448- The transformation matrix required to transform the mask to the original frame.
449- storage_size
450- The size (in pixels) that the compressed mask should be stored at. Default: 128.
451- storage_centering
452- The centering to store the mask at. One of `"legacy"`, `"face"`, `"head"`.
453- Default: `"face"`
447+ size
448+ The size (in pixels) that the compressed mask should be
454449 dilation
455450 The amount of dilation to apply to the mask. as a percentage of the mask size. Default: 0.0
451+ blur_kernel
452+ The kernel size, in pixels to apply gaussian blurring to the mask. Set to 0 for no
453+ blurring. Should be odd, if an even number is passed in (outside of 0) then it is rounded
454+ up to the next odd number. Default: 0
455+ blur_type
456+ The blur type to use. ``gaussian`` or ``normalized`` box filter. Default: ``gaussian``
457+ blur_passes
458+ The number of passed to perform when blurring. Default: 1
456459 """
457460 def __init__ (self ,
458461 area : T .Literal ["eye" , "mouth" , "face" , "face_extended" ],
459462 landmark_type : LandmarkType ,
460463 landmarks : npt .NDArray [np .float32 ],
461- affine_matrix : npt .NDArray [np .float32 ],
462- storage_size : int = 128 ,
463- storage_centering : CenteringType = "face" ,
464- dilation : float = 0.0 ) -> None :
465- super ().__init__ (storage_size = storage_size , storage_centering = storage_centering )
464+ size : int ,
465+ dilation : float = 0.0 ,
466+ blur_kernel : int = 0 ,
467+ blur_type : T .Literal ["gaussian" , "normalized" ] | None = "gaussian" ,
468+ blur_passes : int = 1 ) -> None :
469+ logger .debug (parse_class_init (locals ()))
466470 self ._area = area
467471 self ._landmark_type = landmark_type
468- self ._lm_matrix = affine_matrix
469- self ._points = self ._get_points (landmarks )
470- self .set_dilation (dilation )
472+ self ._landmarks = landmarks
473+ self ._size = size
474+ self ._original_mask : npt .NDArray [np .uint8 ] | None = None
475+
476+ self .dilation = dilation
477+ """The amount of dilation to apply to the mask. as a percentage of the mask size.
478+ Default: 0.0"""
479+ self .blur_kernel = blur_kernel
480+ """The kernel size, in pixels to apply gaussian blurring to the mask. Set to 0 for no
481+ blurring. Should be odd, if an even number is passed in (outside of 0) then it is rounded
482+ up to the next odd number. Default: 0"""
483+ self .blur_type : T .Literal ["gaussian" , "normalized" ] | None = blur_type
484+ """The blur type to use. ``gaussian``, ``normalized`` box filter or ``None`` for no blur.
485+ Default: ``gaussian``"""
486+ self .blur_passes = blur_passes
487+ """The number of passed to perform when blurring. Default: 1"""
488+ self .mask = self .generate_mask ()
489+ """The mask at the size of :attr:`size` with any requested blurring, threshold amount and
490+ centering applied."""
471491
472- @property
473- def mask (self ) -> npt .NDArray [np .uint8 ]:
474- """Overrides the default mask property, creating the processed mask at first call and
475- compressing it. The decompressed mask is returned from this property."""
476- return self .stored_mask
492+ def __repr__ (self ) -> str :
493+ """Pretty print for logging"""
494+ params = {f"{ k [1 :]} : format_array (v ) if isinstance (v , np .ndarray ) else v
495+ for k , v in self .__dict__ .items ()
496+ if k in ("_area" , "_landmark_type" , "_landmarks" , "_size" ,
497+ "_dilation" , "_blur_kernel" , "_blur_type" , "blur_passes" )}
498+ s_params = ", " .join (f"{ k } { repr (v )} for k , v in params .items ())
499+ return f"{ self .__class__ .__name__ } { s_params }
477500
478501 def _get_slices (self ) -> list [slice ] | list [list [slice ]]:
479502 """Obtain the slices that will extract the points for the given area and landmark type
@@ -543,19 +566,15 @@ def _extend_face_landmarks(self,
543566 retval [22 :27 ] = top_r + ((top_r - bot_r ) // 2 )
544567 return retval
545568
546- def _get_points (self , landmarks : npt . NDArray [ np . float32 ] ) -> list [npt .NDArray [np .int32 ]]:
569+ def _get_points (self ) -> list [npt .NDArray [np .int32 ]]:
547570 """Obtain the points required to create the mask
548571
549- Parameters
550- ----------
551- landmarks
552- The landmarks to obtain the points from
553-
554572 Returns
555573 -------
556574 The list of points for creating each section of the mask
557575 """
558576 slices = self ._get_slices ()
577+ landmarks = self ._landmarks
559578 if self ._area == "face_extended" :
560579 landmarks = self ._extend_face_landmarks (landmarks )
561580
@@ -567,26 +586,49 @@ def _get_points(self, landmarks: npt.NDArray[np.float32]) -> list[npt.NDArray[np
567586 for zone in T .cast (list [list [slice ]], slices )]
568587 return retval
569588
570- def generate_mask (self ) -> None :
589+ def _dilate (self , mask : npt .NDArray [np .uint8 ]):
590+ """Perform dilation on the mask
591+
592+ Parameters
593+ ----------
594+ mask
595+ The mask to dilate
596+ """
597+ if self .dilation == 0.0 :
598+ return
599+ kernel_size = int (round (self ._size * abs (self .dilation / 100. ), 0 ))
600+ element = cv2 .getStructuringElement (cv2 .MORPH_ELLIPSE , (kernel_size , kernel_size ))
601+ func = cv2 .erode if self .dilation < 0 else cv2 .dilate
602+ func (mask , element , dst = mask , iterations = 1 )
603+
604+ def generate_mask (self ) -> npt .NDArray [np .uint8 ]:
571605 """Generate the mask.
572606
573- Creates the mask applying any requested dilation and blurring and assigns compressed mask
574- to :attr:`_mask`
607+ Creates the mask applying any requested dilation and blurring
608+
609+ Returns
610+ -------
611+ The landmarks based mask
575612 """
576- mask = np .zeros ((self .stored_size , self .stored_size , 1 ), dtype = np .uint8 )
577- for pts in self ._points :
578- lms = np .rint (pts ).astype ("int" )
579- cv2 .fillConvexPoly (mask , cv2 .convexHull (lms ), [255 ], lineType = cv2 .LINE_AA )
580- if self ._dilation [- 1 ] is not None :
581- self ._dilate_mask (mask )
582- if self ._blur_kernel != 0 and self ._blur_type is not None :
583- mask = BlurMask (self ._blur_type ,
613+ if self ._original_mask is None :
614+ points = self ._get_points ()
615+ mask = np .zeros ((self ._size , self ._size , 1 ), dtype = np .uint8 )
616+ for pts in points :
617+ lms = np .rint (pts ).astype ("int" )
618+ cv2 .fillConvexPoly (mask , cv2 .convexHull (lms ), [255 ], lineType = cv2 .LINE_AA )
619+ self ._original_mask = mask
620+
621+ mask = self ._original_mask .copy ()
622+ self ._dilate (mask )
623+
624+ if self .blur_kernel != 0 and self .blur_type is not None :
625+ mask = BlurMask (self .blur_type ,
584626 mask ,
585- self ._blur_kernel ,
586- passes = self ._blur_passes ).blurred
627+ self .blur_kernel ,
628+ passes = self .blur_passes ).blurred
587629 logger .trace ("[LM_MASK] mask: (shape: %s, dtype: %s)" , # type:ignore[attr-defined]
588630 mask .shape , mask .dtype )
589- self . add ( mask , self . _lm_matrix )
631+ return mask
590632
591633
592634class BlurMask ():
0 commit comments