Files
insightface/alignment/img_helper.py
2019-01-22 10:53:37 +08:00

81 lines
2.6 KiB
Python

import numpy as np
import math
import cv2
from skimage import transform as stf
def transform(data, center, output_size, scale, rotation):
scale_ratio = float(output_size)/scale
rot = float(rotation)*np.pi/180.0
#translation = (output_size/2-center[0]*scale_ratio, output_size/2-center[1]*scale_ratio)
t1 = stf.SimilarityTransform(scale=scale_ratio)
cx = center[0]*scale_ratio
cy = center[1]*scale_ratio
t2 = stf.SimilarityTransform(translation=(-1*cx, -1*cy))
t3 = stf.SimilarityTransform(rotation=rot)
t4 = stf.SimilarityTransform(translation=(output_size/2, output_size/2))
t = t1+t2+t3+t4
trans = t.params[0:2]
#print('M', scale, rotation, trans)
cropped = cv2.warpAffine(data,trans,(output_size, output_size), borderValue = 0.0)
return cropped, trans
def transform_pt(pt, trans):
new_pt = np.array([pt[0], pt[1], 1.]).T
new_pt = np.dot(trans, new_pt)
#print('new_pt', new_pt.shape, new_pt)
return new_pt[:2]
def gaussian(img, pt, sigma):
# Draw a 2D gaussian
assert(sigma>=0)
if sigma==0:
img[pt[1], pt[0]] = 1.0
return True
#assert pt[0]<=img.shape[1]
#assert pt[1]<=img.shape[0]
# Check that any part of the gaussian is in-bounds
ul = [int(pt[0] - 3 * sigma), int(pt[1] - 3 * sigma)]
br = [int(pt[0] + 3 * sigma + 1), int(pt[1] + 3 * sigma + 1)]
if (ul[0] > img.shape[1] or ul[1] >= img.shape[0] or
br[0] < 0 or br[1] < 0):
# If not, just return the image as is
#print('gaussian error')
return False
#return img
# Generate gaussian
size = 6 * sigma + 1
x = np.arange(0, size, 1, float)
y = x[:, np.newaxis]
x0 = y0 = size // 2
# The gaussian is not normalized, we want the center value to equal 1
g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * sigma ** 2))
# Usable gaussian range
g_x = max(0, -ul[0]), min(br[0], img.shape[1]) - ul[0]
g_y = max(0, -ul[1]), min(br[1], img.shape[0]) - ul[1]
# Image range
img_x = max(0, ul[0]), min(br[0], img.shape[1])
img_y = max(0, ul[1]), min(br[1], img.shape[0])
img[img_y[0]:img_y[1], img_x[0]:img_x[1]] = g[g_y[0]:g_y[1], g_x[0]:g_x[1]]
return True
#return img
def estimate_trans_bbox(face, input_size, s = 2.0):
w = face[2] - face[0]
h = face[3] - face[1]
wc = int( (face[2]+face[0])/2 )
hc = int( (face[3]+face[1])/2 )
im_size = max(w, h)
#size = int(im_size*1.2)
scale = input_size/(max(w,h)*s)
M = [
[scale, 0, input_size/2-wc*scale],
[0, scale, input_size/2-hc*scale],
]
M = np.array(M)
return M