refactoring

xlib/face/FaceWarper.py

@@ -0,0 +1,143 @@
+from typing import Iterable, Tuple, Union
+
+import cv2
+import numpy as np
+
+from ..math import Affine2DMat, Affine2DUniMat
+
+class FaceWarper:
+    def __init__(self,
+                 img_to_face_uni_mat : Affine2DUniMat,
+
+                 align_rot_deg : Union[None, float,  Tuple[float, float] ] = [-15,15],
+                 align_scale : Union[None, float,  Tuple[float, float] ]   = [-0.15, 0.15],
+                 align_tx : Union[None, float,  Tuple[float, float] ]      = [-0.05, 0.05],
+                 align_ty : Union[None, float,  Tuple[float, float] ]      = [-0.05, 0.05],
+
+                 rw_grid_cell_count : Union[None, int,  Tuple[int, int] ]    = [3,7],
+                 rw_grid_rot_deg : Union[None, float,  Tuple[float, float] ] = [-180,180],
+                 rw_grid_scale : Union[None, float,  Tuple[float, float] ]   = [-0.25, 0.25],
+                 rw_grid_tx : Union[None, float,  Tuple[float, float] ]      = [-0.25, 0.25],
+                 rw_grid_ty : Union[None, float,  Tuple[float, float] ]      = [-0.25, 0.25],
+
+                 rnd_state : np.random.RandomState = None,
+                ):
+        """
+        Max quality one-pass face augmentation via geometric transformations with provided values.
+        
+            img_to_face_uni_mat    Affine2DUniMat
+            
+        Affine2DUniMat given from FLandmarks2D.calc_cut
+        it is an uniform affineMat to transform original image to aligned face
+
+            align_* rw_grid_*
+            
+        exact augmentation parameters or range for random generation.
+        
+        
+        """
+        self._img_to_face_uni_mat = img_to_face_uni_mat
+        self._face_to_img_uni_mat = img_to_face_uni_mat.invert()
+
+        if rnd_state is None:
+            rnd_state = np.random
+        self._rnd_state_state = rnd_state.get_state()
+        
+        self._align_rot_deg = rnd_state.uniform(*align_rot_deg) if isinstance(align_rot_deg, Iterable) else align_rot_deg
+        self._align_scale = rnd_state.uniform(*align_scale) if isinstance(align_scale, Iterable) else align_scale
+        self._align_tx = rnd_state.uniform(*align_tx) if isinstance(align_tx, Iterable) else align_tx
+        self._align_ty = rnd_state.uniform(*align_ty) if isinstance(align_ty, Iterable) else align_ty
+        self._rw_grid_cell_count = rnd_state.randint(*rw_grid_cell_count) if isinstance(rw_grid_cell_count, Iterable) else rw_grid_cell_count
+        self._rw_grid_rot_deg = rnd_state.uniform(*rw_grid_rot_deg) if isinstance(rw_grid_rot_deg, Iterable) else rw_grid_rot_deg
+        self._rw_grid_scale = rnd_state.uniform(*rw_grid_scale) if isinstance(rw_grid_scale, Iterable) else rw_grid_scale
+        self._rw_grid_tx = rnd_state.uniform(*rw_grid_tx) if isinstance(rw_grid_tx, Iterable) else rw_grid_tx
+        self._rw_grid_ty = rnd_state.uniform(*rw_grid_ty) if isinstance(rw_grid_ty, Iterable) else rw_grid_ty
+
+        self._cached = {}
+        
+    def transform(self, img : np.ndarray, out_res : int, random_warp : bool = True) -> np.ndarray:
+        """
+        transform an image. Subsequent calls will output the same result for any img shape and out_res.
+        
+            img         np.ndarray  (HWC)
+            
+            out_res     int
+            
+            random_warp(True)   bool
+        """
+        H,W = img.shape[:2]
+
+        key = (H,W,random_warp)
+        data = self._cached.get(key, None)
+        if data is None:
+            rnd_state = np.random.RandomState()
+            rnd_state.set_state( self._rnd_state_state )
+            
+            image_grid, face_mask = self._cached[key] = self._gen(H,W, random_warp, out_res, rnd_state=rnd_state )
+            
+        new_img = cv2.remap(img, image_grid, None, interpolation=cv2.INTER_LANCZOS4)
+        new_img *= face_mask
+        return new_img
+        
+    def _gen(self, H, W, random_warp, out_res, rnd_state):
+
+        image_grid = np.stack(np.meshgrid(np.linspace(0., 1.0, H, dtype=np.float32),
+                                          np.linspace(0., 1.0, W, dtype=np.float32)), -1)
+        
+        if random_warp:
+            # make a random face_warp_grid in the space of the face
+            face_warp_grid = FaceWarper._gen_random_warp_uni_grid_diff(out_res, self._rw_grid_cell_count, 0.12, rnd_state)
+
+            # make a randomly transformable mat to transform face_warp_grid from face to image
+            face_warp_grid_mat = (self._face_to_img_uni_mat *
+                                  Affine2DUniMat.from_transformation(0.5, 0.5, self._rw_grid_rot_deg, 1.0+self._rw_grid_scale, self._rw_grid_tx, self._rw_grid_ty)
+                                  )
+
+            # warp face_warp_grid to the space of image and merge with image_grid
+            image_grid += cv2.warpAffine(face_warp_grid, face_warp_grid_mat.to_exact_mat(out_res,out_res, W, H), (W,H), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT)
+
+        # scale uniform grid from to image size
+        image_grid *= (H-1, W-1)
+
+        # apply random transormations for align mat
+        img_to_face_rnd_mat = (self._face_to_img_uni_mat * Affine2DMat.from_transformation(0.5, 0.5, self._align_rot_deg, 1.0+self._align_scale, self._align_tx, self._align_ty)
+                                ).invert().to_exact_mat(W,H,out_res,out_res)
+
+        # warp image_grid to face space
+        image_grid = cv2.warpAffine(image_grid, img_to_face_rnd_mat, (out_res,out_res), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE )
+
+        # One-pass remap from original image to aligned face with all transformations
+        #new_img = cv2.remap(img, image_grid, None, interpolation=cv2.INTER_LANCZOS4)
+
+        # make mask to refine image-boundary visible in face space
+        face_mask = cv2.warpAffine( np.ones( (H,W), dtype=np.uint8), img_to_face_rnd_mat, (out_res,out_res), flags=cv2.INTER_NEAREST)[...,None]
+
+        return image_grid, face_mask
+        
+    def _gen_random_warp_uni_grid_diff(size: int, cell_count, cell_mod, rnd_state) -> np.ndarray:
+        """
+        generates square uniform random warp
+        grid of shape (size, size, 2)  (x,y)
+
+        cell_count(3)        3+
+
+        cell_mod  (0.12)     [ 0 .. 0.24 ]
+        """
+        cell_count = max(3, cell_count)
+        cell_mod = np.clip(cell_mod, 0, 0.24)
+        cell_size = 1.0 / (cell_count-1)
+
+        grid = np.zeros( (cell_count,cell_count, 2), dtype=np.float32 )
+
+        grid[1:-1,1:-1, 0:2] += rnd_state.uniform (low=-cell_size*cell_mod, high=cell_size*cell_mod, size=(cell_count-2, cell_count-2, 2) )
+        grid = cv2.resize(grid, (size, size), interpolation=cv2.INTER_CUBIC ).astype(np.float32)
+
+        # Linear dump border cells to zero
+        border_size = size // cell_count
+        dumper = np.linspace(0, 1, border_size, dtype=np.float32)
+        grid[:border_size, :,:] *= dumper[:,None,None]
+        grid[-border_size:,:,:] *= dumper[::-1,None,None]
+        grid[:,:border_size ,:] *= dumper[None,:,None]
+        grid[:,-border_size:,:] *= dumper[None,::-1,None]
+
+        return grid