177 lines
5.1 KiB
Python
177 lines
5.1 KiB
Python
import os
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
import math
|
|
import PIL
|
|
from PIL import Image
|
|
import folder_paths
|
|
|
|
|
|
_BBREGRESSOR_PARAM_FILE = os.path.join(folder_paths.models_dir, "alkaid/BBRegressorParam.npy")
|
|
|
|
BBREGRESSOR_PARAM = None
|
|
|
|
|
|
|
|
def tensorToNP(image):
|
|
out = torch.clamp(255. * image.detach().cpu(), 0, 255).to(torch.uint8)
|
|
if out.shape[1] == 3:
|
|
out = out.permute(0, 2, 3, 1)
|
|
out = out.numpy()
|
|
return out
|
|
|
|
def tensor2pil(image):
|
|
return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
|
|
|
|
def pil2tensor(image):
|
|
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|
|
|
|
def BBRegression(points):
|
|
global BBREGRESSOR_PARAM
|
|
if BBREGRESSOR_PARAM is None:
|
|
BBREGRESSOR_PARAM = np.load(_BBREGRESSOR_PARAM_FILE, allow_pickle=True).item()
|
|
|
|
w1 = BBREGRESSOR_PARAM['W1']
|
|
b1 = BBREGRESSOR_PARAM['B1']
|
|
w2 = BBREGRESSOR_PARAM['W2']
|
|
b2 = BBREGRESSOR_PARAM['B2']
|
|
data = points.copy()
|
|
data = data.reshape([5, 2])
|
|
data_mean = np.mean(data, axis=0)
|
|
x_mean = data_mean[0]
|
|
y_mean = data_mean[1]
|
|
data[:, 0] = data[:, 0] - x_mean
|
|
data[:, 1] = data[:, 1] - y_mean
|
|
|
|
rms = np.sqrt(np.sum(data ** 2) / 5)
|
|
data = data / rms
|
|
data = data.reshape([1, 10])
|
|
data = np.transpose(data)
|
|
inputs = np.matmul(w1, data) + b1
|
|
inputs = 2 / (1 + np.exp(-2 * inputs)) - 1
|
|
inputs = np.matmul(w2, inputs) + b2
|
|
inputs = np.transpose(inputs)
|
|
x = inputs[:, 0] * rms + x_mean
|
|
y = inputs[:, 1] * rms + y_mean
|
|
w = 224 / inputs[:, 2] * rms
|
|
rects = [x, y, w, w]
|
|
return np.array(rects).reshape([4])
|
|
|
|
|
|
def extract_5p(lm):
|
|
lm_idx = np.array([31, 37, 40, 43, 46, 49, 55]) - 1
|
|
value_list = [
|
|
lm[lm_idx[0], :],
|
|
np.mean(lm[lm_idx[[1, 2]], :], 0),
|
|
np.mean(lm[lm_idx[[3, 4]], :], 0), lm[lm_idx[5], :], lm[lm_idx[6], :]
|
|
]
|
|
lm5p = np.stack(value_list, axis=0)
|
|
lm5p = lm5p[[1, 2, 0, 3, 4], :]
|
|
return lm5p
|
|
|
|
|
|
def get_estimate_bbox(image, insightface, target_size=224., rescale_factor=102):
|
|
face = insightface.get(image)
|
|
h, w = image.shape[0], image.shape[1]
|
|
if face.__len__() == 0:
|
|
print('Detected no face, use all image')
|
|
return image, np.array([0, 0, w, h])
|
|
face = sorted(face, key=lambda x: (x['bbox'][2] - x['bbox'][0]) * (x['bbox'][3] - x['bbox'][1]))[-1]
|
|
kps = face['kps']
|
|
global BBREGRESSOR_PARAM
|
|
|
|
if BBREGRESSOR_PARAM is None:
|
|
BBREGRESSOR_PARAM = np.load(_BBREGRESSOR_PARAM_FILE, allow_pickle=True).item()
|
|
|
|
lm3d_std = BBREGRESSOR_PARAM['lm3d5_std']
|
|
t, s = POS(kps.transpose(), lm3d_std.transpose())
|
|
t = t.squeeze()
|
|
s = rescale_factor / s
|
|
|
|
l = int(t[0] - target_size / (2 * s))
|
|
r = l + int(target_size / s)
|
|
u = int(t[1] - target_size / (2 * s))
|
|
b = u + int(target_size / s)
|
|
bbox = np.array([u, l, b, r])
|
|
image_crop = image[bbox[0]:bbox[2], bbox[1]:bbox[3]]
|
|
image_crop = cv2.resize(image_crop, (224, 224))
|
|
return image_crop, bbox
|
|
|
|
|
|
def POS(xp, x):
|
|
npts = xp.shape[1]
|
|
|
|
A = np.zeros([2 * npts, 8])
|
|
|
|
A[0:2 * npts - 1:2, 0:3] = x.transpose()
|
|
A[0:2 * npts - 1:2, 3] = 1
|
|
|
|
A[1:2 * npts:2, 4:7] = x.transpose()
|
|
A[1:2 * npts:2, 7] = 1
|
|
|
|
b = np.reshape(xp.transpose(), [2 * npts, 1])
|
|
|
|
k, _, _, _ = np.linalg.lstsq(A, b)
|
|
|
|
R1 = k[0:3]
|
|
R2 = k[4:7]
|
|
sTx = k[3]
|
|
sTy = k[7]
|
|
s = (np.linalg.norm(R1) + np.linalg.norm(R2)) / 2
|
|
t = np.stack([sTx, sTy], axis=0)
|
|
|
|
return t, s
|
|
|
|
|
|
def draw_kps(h, w, kps, color_list=[(255,0,0), (0,255,0), (0,0,255), (255,255,0), (255,0,255)]):
|
|
stickwidth = 4
|
|
limbSeq = np.array([[0, 2], [1, 2], [3, 2], [4, 2]])
|
|
kps = np.array(kps)
|
|
|
|
out_img = np.zeros([h, w, 3])
|
|
|
|
for i in range(len(limbSeq)):
|
|
index = limbSeq[i]
|
|
color = color_list[index[0]]
|
|
|
|
x = kps[index][:, 0]
|
|
y = kps[index][:, 1]
|
|
length = ((x[0] - x[1]) ** 2 + (y[0] - y[1]) ** 2) ** 0.5
|
|
angle = math.degrees(math.atan2(y[0] - y[1], x[0] - x[1]))
|
|
polygon = cv2.ellipse2Poly((int(np.mean(x)), int(np.mean(y))), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
|
|
out_img = cv2.fillConvexPoly(out_img.copy(), polygon, color)
|
|
out_img = (out_img * 0.6).astype(np.uint8)
|
|
|
|
for idx_kp, kp in enumerate(kps):
|
|
color = color_list[idx_kp]
|
|
x, y = kp
|
|
out_img = cv2.circle(out_img.copy(), (int(x), int(y)), 10, color, -1)
|
|
|
|
out_img_pil = Image.fromarray(out_img.astype(np.uint8))
|
|
return out_img_pil
|
|
|
|
|
|
def tensor_to_size(source, dest_size):
|
|
if isinstance(dest_size, torch.Tensor):
|
|
dest_size = dest_size.shape[0]
|
|
source_size = source.shape[0]
|
|
|
|
if source_size < dest_size:
|
|
shape = [dest_size - source_size] + [1]*(source.dim()-1)
|
|
source = torch.cat((source, source[-1:].repeat(shape)), dim=0)
|
|
elif source_size > dest_size:
|
|
source = source[:dest_size]
|
|
|
|
return source
|
|
|
|
def tensor_to_image(tensor):
|
|
image = tensor.mul(255).clamp(0, 255).byte().cpu()
|
|
image = image[..., [2, 1, 0]].numpy()
|
|
return image
|
|
|
|
def image_to_tensor(image):
|
|
tensor = torch.clamp(torch.from_numpy(image).float() / 255., 0, 1)
|
|
tensor = tensor[..., [2, 1, 0]]
|
|
return tensor
|