mirror of
https://github.com/iperov/DeepFaceLive.git
synced 2024-12-26 07:51:13 -08:00
73 lines
2.5 KiB
Python
73 lines
2.5 KiB
Python
import numpy as np
|
|
|
|
from ..AShape import AShape
|
|
from ..initializer import InitCoords2DArange
|
|
from ..SCacheton import SCacheton
|
|
from ..Tensor import Tensor
|
|
from .matmul import matmul
|
|
from .remap import remap
|
|
|
|
def warp_affine (input_t : Tensor, affine_t : Tensor, output_size=None, dtype=None) -> Tensor:
|
|
"""
|
|
arguments
|
|
|
|
input_t Tensor(...,H,W)
|
|
|
|
affine_t Tensor(...,2,3)
|
|
affine matrix
|
|
|
|
example of identity affine matrix
|
|
[1,0,0],
|
|
[0,1,0]
|
|
|
|
...-head part of shapes will be broadcasted to each other
|
|
|
|
output_size(None)
|
|
|
|
tuple of 2 ints (HW)
|
|
of output size
|
|
if None , size will not be changed
|
|
|
|
"""
|
|
op = SCacheton.get(_WarpAffineOp, input_t.shape, input_t.dtype, affine_t.shape, affine_t.dtype, output_size)
|
|
|
|
affine_t = affine_t.transpose( op.affine_transpose_axes, dtype=np.float32 ).reshape( (-1,3,2) )
|
|
|
|
coords_t = Tensor(op.coords_shape, np.float32, device=input_t.get_device(), initializer=op.coords_init )
|
|
coords_t = coords_t.reshape(op.coords_reshape)
|
|
coords_t = matmul(coords_t, affine_t).reshape(op.coords_affined_shape)
|
|
|
|
output_t = remap(input_t, coords_t, dtype=dtype)
|
|
return output_t
|
|
|
|
|
|
class _WarpAffineOp():
|
|
def __init__(self, i_shape : AShape, i_dtype, a_shape : AShape, a_dtype, o_size):
|
|
if np.dtype(i_dtype).type == np.bool_:
|
|
raise ValueError('np.bool_ dtype of i_dtype is not supported.')
|
|
if np.dtype(a_dtype).type == np.bool_:
|
|
raise ValueError('np.bool_ dtype of a_dtype is not supported.')
|
|
if i_shape.ndim < 2:
|
|
raise ValueError('i_shape.ndim must be >= 2 (...,H,W)')
|
|
if a_shape.ndim < 2:
|
|
raise ValueError(f'a_shape.ndim must be >= 2 (...,2,3)')
|
|
if a_shape[-2] != 2 or a_shape[-1] != 3:
|
|
raise ValueError('Last a_shape dims must be == (...,2,3)')
|
|
|
|
IH,IW = i_shape[-2:]
|
|
if o_size is not None:
|
|
OH,OW = o_size
|
|
else:
|
|
OH,OW = IH,IW
|
|
|
|
self.coords_shape = AShape( (OH,OW,3) )
|
|
self.coords_affined_shape = AShape( (OH,OW,2) )
|
|
|
|
if a_shape.ndim > 2:
|
|
self.coords_shape = a_shape[:-2] + self.coords_shape
|
|
self.coords_affined_shape = a_shape[:-2] + self.coords_affined_shape
|
|
|
|
self.coords_init = InitCoords2DArange(0,OH-1,0,OW-1)
|
|
self.coords_reshape = (-1,OH*OW,3)
|
|
self.affine_transpose_axes = a_shape.axes_arange().swapped_axes(-2,-1)
|