File size: 3,472 Bytes
d435aa8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 |
import torch
import numpy as np
import cv2
from skimage import transform as trans
import cv2
def split_network_output(align_out):
anchor_bbox_pred, anchor_cls_pred, anchor_ldmk_pred, merged, _ = align_out
bbox, cls, ldmk = torch.split(merged, [4, 2, 10], dim=1)
return ldmk, bbox, cls
def get_cv2_affine_from_landmark(ldmks, reference_ldmk, image_width, image_height, ):
assert ldmks.ndim == 2 # batchdim
assert ldmks.shape[1] == 10
assert isinstance(ldmks, torch.Tensor)
assert reference_ldmk.ndim == 2
assert reference_ldmk.shape[0] == 5
assert reference_ldmk.shape[1] == 2
assert isinstance(reference_ldmk, np.ndarray)
to_img_size = np.array([[[image_width, image_height]]])
ldmks = ldmks.view(ldmks.shape[0], 5, 2).detach().cpu().numpy()
ldmks = ldmks * to_img_size
transforms = []
for ldmk in ldmks:
tform = trans.SimilarityTransform()
tform.estimate(ldmk, reference_ldmk)
M = tform.params[0:2, :]
transforms.append(M)
transforms = np.stack(transforms, axis=0)
return transforms
def cv2_param_to_torch_theta(cv2_tfms, image_width, image_height, output_width, output_height):
# https://github.com/wuneng/WarpAffine2GridSample
"""4.Affine Transformation Matrix to theta"""
assert cv2_tfms.ndim == 3 # N, 2, 3
assert cv2_tfms.shape[1] == 2
assert cv2_tfms.shape[2] == 3
srcs = np.array([[0, 0], [0, 1], [1, 1]], dtype=np.float32)
srcs = np.expand_dims(srcs, axis=0).repeat(cv2_tfms.shape[0], axis=0)
dsts = np.matmul(srcs, cv2_tfms[:, :, :2].transpose(0, 2, 1)) + cv2_tfms[:, :, 2:3].transpose(0, 2, 1)
# normalize to [-1, 1]
srcs = srcs / np.array([[[image_width, image_height]]]) * 2 - 1
dsts = dsts / np.array([[[output_width, output_height]]]) * 2 - 1
thetas = []
for src, dst in zip(srcs, dsts):
theta = trans.estimate_transform("affine", src=dst, dst=src).params[:2]
thetas.append(theta)
thetas = np.stack(thetas, axis=0)
thetas = torch.from_numpy(thetas).float()
return thetas
def adjust_ldmks(ldmks, thetas):
inv_thetas = inv_matrix(thetas).to(ldmks.device).float()
_ldmks = torch.cat([ldmks, torch.ones((ldmks.shape[0], 5, 1)).to(ldmks.device)], dim=2)
ldmk_aligned = (((_ldmks) * 2 - 1) @ inv_thetas.permute(0,2,1)) / 2 + 0.5
return ldmk_aligned
def inv_matrix(theta):
# torch batched version
assert theta.ndim == 3
a, b, t1 = theta[:, 0,0], theta[:, 0,1], theta[:, 0,2]
c, d, t2 = theta[:, 1,0], theta[:, 1,1], theta[:, 1,2]
det = a * d - b * c
inv_det = 1.0 / det
inv_mat = torch.stack([
torch.stack([d * inv_det, -b * inv_det, (b * t2 - d * t1) * inv_det], dim=1),
torch.stack([-c * inv_det, a * inv_det, (c * t1 - a * t2) * inv_det], dim=1)
], dim=1)
return inv_mat
def reference_landmark():
return np.array([[38.29459953, 51.69630051],
[73.53179932, 51.50139999],
[56.02519989, 71.73660278],
[41.54930115, 92.3655014],
[70.72990036, 92.20410156]])
def draw_ldmk(img, ldmk):
if ldmk is None:
return img
colors = [(0, 255, 0), (255, 0, 0), (0, 0, 255), (255, 255, 0), (0, 255, 255), (255, 0, 255)]
img = img.copy()
for i in range(5):
color = colors[i]
cv2.circle(img, (int(ldmk[i*2] * img.shape[1]), int(ldmk[i*2+1] * img.shape[0])), 1, color, 4)
return img |