mirror of
https://github.com/deepinsight/insightface.git
synced 2026-05-14 12:17:55 +00:00
237 lines
9.6 KiB
Python
237 lines
9.6 KiB
Python
import datetime
|
|
import os
|
|
import os.path as osp
|
|
import sys
|
|
|
|
import albumentations as A
|
|
import cv2
|
|
import numpy as np
|
|
from albumentations.core.transforms_interface import ImageOnlyTransform
|
|
|
|
from ..data.image import get_image as ins_get_image
|
|
from ..thirdparty import face3d
|
|
from ..utils import DEFAULT_MP_NAME, get_model_dir
|
|
from .face_analysis import FaceAnalysis
|
|
|
|
|
|
class MaskRenderer:
|
|
def __init__(self, name=DEFAULT_MP_NAME, root='~/.insightface', insfa=None):
|
|
#if insfa is None, enter render_only mode
|
|
self.mp_name = name
|
|
self.root = root
|
|
self.insfa = insfa
|
|
model_dir = get_model_dir(name, root)
|
|
bfm_file = osp.join(model_dir, 'BFM.mat')
|
|
assert osp.exists(bfm_file), 'should contains BFM.mat in your model directory'
|
|
self.bfm = face3d.morphable_model.MorphabelModel(bfm_file)
|
|
self.index_ind = self.bfm.kpt_ind
|
|
bfm_uv_file = osp.join(model_dir, 'BFM_UV.mat')
|
|
assert osp.exists(bfm_uv_file), 'should contains BFM_UV.mat in your model directory'
|
|
uv_coords = face3d.morphable_model.load.load_uv_coords(bfm_uv_file)
|
|
self.uv_size = (224,224)
|
|
self.mask_stxr = 0.1
|
|
self.mask_styr = 0.33
|
|
self.mask_etxr = 0.9
|
|
self.mask_etyr = 0.7
|
|
self.tex_h , self.tex_w, self.tex_c = self.uv_size[1] , self.uv_size[0],3
|
|
texcoord = np.zeros_like(uv_coords)
|
|
texcoord[:, 0] = uv_coords[:, 0] * (self.tex_h - 1)
|
|
texcoord[:, 1] = uv_coords[:, 1] * (self.tex_w - 1)
|
|
texcoord[:, 1] = self.tex_w - texcoord[:, 1] - 1
|
|
self.texcoord = np.hstack((texcoord, np.zeros((texcoord.shape[0], 1))))
|
|
self.X_ind = self.bfm.kpt_ind
|
|
self.mask_image_names = ['mask_white', 'mask_blue', 'mask_black', 'mask_green']
|
|
self.mask_aug_probs = [0.4, 0.4, 0.1, 0.1]
|
|
#self.mask_images = []
|
|
#self.mask_images_rgb = []
|
|
#for image_name in mask_image_names:
|
|
# mask_image = ins_get_image(image_name)
|
|
# self.mask_images.append(mask_image)
|
|
# mask_image_rgb = mask_image[:,:,::-1]
|
|
# self.mask_images_rgb.append(mask_image_rgb)
|
|
|
|
|
|
def prepare(self, ctx_id=0, det_thresh=0.5, det_size=(128, 128)):
|
|
self.pre_ctx_id = ctx_id
|
|
self.pre_det_thresh = det_thresh
|
|
self.pre_det_size = det_size
|
|
|
|
def transform(self, shape3D, R):
|
|
s = 1.0
|
|
shape3D[:2, :] = shape3D[:2, :]
|
|
shape3D = s * np.dot(R, shape3D)
|
|
return shape3D
|
|
|
|
def preprocess(self, vertices, w, h):
|
|
R1 = face3d.mesh.transform.angle2matrix([0, 180, 180])
|
|
t = np.array([-w // 2, -h // 2, 0])
|
|
vertices = vertices.T
|
|
vertices += t
|
|
vertices = self.transform(vertices.T, R1).T
|
|
return vertices
|
|
|
|
def project_to_2d(self,vertices,s,angles,t):
|
|
transformed_vertices = self.bfm.transform(vertices, s, angles, t)
|
|
projected_vertices = transformed_vertices.copy() # using stantard camera & orth projection
|
|
return projected_vertices[self.bfm.kpt_ind, :2]
|
|
|
|
def params_to_vertices(self,params , H , W):
|
|
fitted_sp, fitted_ep, fitted_s, fitted_angles, fitted_t = params
|
|
fitted_vertices = self.bfm.generate_vertices(fitted_sp, fitted_ep)
|
|
transformed_vertices = self.bfm.transform(fitted_vertices, fitted_s, fitted_angles,
|
|
fitted_t)
|
|
transformed_vertices = self.preprocess(transformed_vertices.T, W, H)
|
|
image_vertices = face3d.mesh.transform.to_image(transformed_vertices, H, W)
|
|
return image_vertices
|
|
|
|
def draw_lmk(self, face_image):
|
|
faces = self.insfa.get(face_image, max_num=1)
|
|
if len(faces)==0:
|
|
return face_image
|
|
return self.insfa.draw_on(face_image, faces)
|
|
|
|
def build_params(self, face_image):
|
|
#landmark = self.if3d68_handler.get(face_image)
|
|
#if landmark is None:
|
|
# return None #face not found
|
|
if self.insfa is None:
|
|
self.insfa = FaceAnalysis(name=self.mp_name, root=self.root, allowed_modules=['detection', 'landmark_3d_68'])
|
|
self.insfa.prepare(ctx_id=self.pre_ctx_id, det_thresh=self.pre_det_thresh, det_size=self.pre_det_size)
|
|
|
|
faces = self.insfa.get(face_image, max_num=1)
|
|
if len(faces)==0:
|
|
return None
|
|
landmark = faces[0].landmark_3d_68[:,:2]
|
|
fitted_sp, fitted_ep, fitted_s, fitted_angles, fitted_t = self.bfm.fit(landmark, self.X_ind, max_iter = 3)
|
|
return [fitted_sp, fitted_ep, fitted_s, fitted_angles, fitted_t]
|
|
|
|
def generate_mask_uv(self,mask, positions):
|
|
uv_size = (self.uv_size[1], self.uv_size[0], 3)
|
|
h, w, c = uv_size
|
|
uv = np.zeros(shape=(self.uv_size[1],self.uv_size[0], 3), dtype=np.uint8)
|
|
stxr, styr = positions[0], positions[1]
|
|
etxr, etyr = positions[2], positions[3]
|
|
stx, sty = int(w * stxr), int(h * styr)
|
|
etx, ety = int(w * etxr), int(h * etyr)
|
|
height = ety - sty
|
|
width = etx - stx
|
|
mask = cv2.resize(mask, (width, height))
|
|
uv[sty:ety, stx:etx] = mask
|
|
return uv
|
|
|
|
def render_mask(self,face_image, mask_image, params, input_is_rgb=False, auto_blend = True, positions=[0.1, 0.33, 0.9, 0.7]):
|
|
if isinstance(mask_image, str):
|
|
to_rgb = True if input_is_rgb else False
|
|
mask_image = ins_get_image(mask_image, to_rgb=to_rgb)
|
|
uv_mask_image = self.generate_mask_uv(mask_image, positions)
|
|
h,w,c = face_image.shape
|
|
image_vertices = self.params_to_vertices(params ,h,w)
|
|
output = (1-face3d.mesh.render.render_texture(image_vertices, self.bfm.full_triangles , uv_mask_image, self.texcoord, self.bfm.full_triangles, h , w ))*255
|
|
output = output.astype(np.uint8)
|
|
if auto_blend:
|
|
mask_bd = (output==255).astype(np.uint8)
|
|
final = face_image*mask_bd + (1-mask_bd)*output
|
|
return final
|
|
return output
|
|
|
|
#def mask_augmentation(self, face_image, label, input_is_rgb=False, p=0.1):
|
|
# if np.random.random()<p:
|
|
# assert isinstance(label, (list, np.ndarray)), 'make sure the rec dataset includes mask params'
|
|
# assert len(label)==237 or len(lable)==235, 'make sure the rec dataset includes mask params'
|
|
# if len(label)==237:
|
|
# if label[1]<0.0: #invalid label for mask aug
|
|
# return face_image
|
|
# label = label[2:]
|
|
# params = self.decode_params(label)
|
|
# mask_image_name = np.random.choice(self.mask_image_names, p=self.mask_aug_probs)
|
|
# pos = np.random.uniform(0.33, 0.5)
|
|
# face_image = self.render_mask(face_image, mask_image_name, params, input_is_rgb=input_is_rgb, positions=[0.1, pos, 0.9, 0.7])
|
|
# return face_image
|
|
|
|
@staticmethod
|
|
def encode_params(params):
|
|
p0 = list(params[0])
|
|
p1 = list(params[1])
|
|
p2 = [float(params[2])]
|
|
p3 = list(params[3])
|
|
p4 = list(params[4])
|
|
return p0+p1+p2+p3+p4
|
|
|
|
@staticmethod
|
|
def decode_params(params):
|
|
p0 = params[0:199]
|
|
p0 = np.array(p0, dtype=np.float32).reshape( (-1, 1))
|
|
p1 = params[199:228]
|
|
p1 = np.array(p1, dtype=np.float32).reshape( (-1, 1))
|
|
p2 = params[228]
|
|
p3 = tuple(params[229:232])
|
|
p4 = params[232:235]
|
|
p4 = np.array(p4, dtype=np.float32).reshape( (-1, 1))
|
|
return p0, p1, p2, p3, p4
|
|
|
|
class MaskAugmentation(ImageOnlyTransform):
|
|
|
|
def __init__(
|
|
self,
|
|
mask_names=['mask_white', 'mask_blue', 'mask_black', 'mask_green'],
|
|
mask_probs=[0.4,0.4,0.1,0.1],
|
|
h_low = 0.33,
|
|
h_high = 0.35,
|
|
always_apply=False,
|
|
p=1.0,
|
|
):
|
|
super(MaskAugmentation, self).__init__(always_apply, p)
|
|
self.renderer = MaskRenderer()
|
|
assert len(mask_names)>0
|
|
assert len(mask_names)==len(mask_probs)
|
|
self.mask_names = mask_names
|
|
self.mask_probs = mask_probs
|
|
self.h_low = h_low
|
|
self.h_high = h_high
|
|
#self.hlabel = None
|
|
|
|
|
|
def apply(self, image, hlabel, mask_name, h_pos, **params):
|
|
#print(params.keys())
|
|
#hlabel = params.get('hlabel')
|
|
assert len(hlabel)==237 or len(hlabel)==235, 'make sure the rec dataset includes mask params'
|
|
if len(hlabel)==237:
|
|
if hlabel[1]<0.0:
|
|
return image
|
|
hlabel = hlabel[2:]
|
|
#print(len(hlabel))
|
|
mask_params = self.renderer.decode_params(hlabel)
|
|
image = self.renderer.render_mask(image, mask_name, mask_params, input_is_rgb=True, positions=[0.1, h_pos, 0.9, 0.7])
|
|
return image
|
|
|
|
@property
|
|
def targets_as_params(self):
|
|
return ["image", "hlabel"]
|
|
|
|
def get_params_dependent_on_targets(self, params):
|
|
hlabel = params['hlabel']
|
|
mask_name = np.random.choice(self.mask_names, p=self.mask_probs)
|
|
h_pos = np.random.uniform(self.h_low, self.h_high)
|
|
return {'hlabel': hlabel, 'mask_name': mask_name, 'h_pos': h_pos}
|
|
|
|
def get_transform_init_args_names(self):
|
|
#return ("hlabel", 'mask_names', 'mask_probs', 'h_low', 'h_high')
|
|
return ('mask_names', 'mask_probs', 'h_low', 'h_high')
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tool = MaskRenderer('antelope')
|
|
tool.prepare(det_size=(128,128))
|
|
image = cv2.imread("Tom_Hanks_54745.png")
|
|
params = tool.build_params(image)
|
|
#out = tool.draw_lmk(image)
|
|
#cv2.imwrite('output_lmk.jpg', out)
|
|
#mask_image = cv2.imread("masks/mask1.jpg")
|
|
#mask_image = cv2.imread("masks/black-mask.png")
|
|
#mask_image = cv2.imread("masks/mask2.jpg")
|
|
mask_out = tool.render_mask(image, 'mask_blue', params)# use single thread to test the time cost
|
|
|
|
cv2.imwrite('output_mask.jpg', mask_out)
|
|
|
|
|