mirror of
https://github.com/deepinsight/insightface.git
synced 2026-05-16 13:46:15 +00:00
171 lines
6.9 KiB
Python
171 lines
6.9 KiB
Python
# Copyright (c) 2020, Baris Gecer. All rights reserved.
|
|
#
|
|
# This work is made available under the CC BY-NC-SA 4.0.
|
|
# To view a copy of this license, see LICENSE
|
|
|
|
import PIL.Image
|
|
from skimage.morphology import binary_dilation, disk
|
|
from skimage.filters import gaussian
|
|
from scipy.interpolate import NearestNDInterpolator
|
|
import cv2 as cv
|
|
from utils.shading import lambertian_shading
|
|
from menpo3d.rasterize import rasterize_mesh
|
|
from menpo.shape import TriMesh, TexturedTriMesh, ColouredTriMesh
|
|
from menpo.image import Image
|
|
from utils.image_rasterization import *
|
|
import menpo.io as mio
|
|
import os
|
|
|
|
gold = (244/256,178/256,88/256)
|
|
green = (26/256,140/256,57/256)
|
|
blue = (24/256,148/256,187/256)
|
|
red = (255/256,51/256,51/256)
|
|
purple = (102/256,0/256,204/256)
|
|
pink = (255/256,21/256,255/256)
|
|
|
|
def perspective_camera(img_shape, focal_length_mult=2,
|
|
pose_angle_deg=[0,0,0], translation=4):
|
|
f = np.array(img_shape).max() * focal_length_mult
|
|
rot_z = rotation_z(180 + pose_angle_deg[2])
|
|
rot_y = rotation_y(180 + pose_angle_deg[1])
|
|
rot_x = rotation_x(pose_angle_deg[0])
|
|
rotation = rot_z.compose_before(rot_y).compose_before(rot_x)
|
|
|
|
translation = Translation([0, 0, translation])
|
|
projection = PerspectiveProjection(f, img_shape)
|
|
return PerspectiveCamera(rotation, translation, projection)
|
|
|
|
def rasterize_mesh_at_template(mesh, img_shape=(1024,1024), pose_angle_deg=[0, 0, 0], translation=4, shaded=True, **kwargs):
|
|
camera = perspective_camera(img_shape, pose_angle_deg=pose_angle_deg, translation=translation)
|
|
if shaded:
|
|
mesh = lambertian_shading(mesh, **kwargs)
|
|
mesh.points[mesh.points == 0] = None
|
|
return rasterize_mesh(camera.apply(mesh), img_shape)
|
|
|
|
def export_tmesh(tmesh, path, pose_angle_deg=[0, 30, 0], translation=4):
|
|
render = rasterize_mesh_at_template(tmesh, pose_angle_deg=pose_angle_deg, shaded=False, translation=translation)
|
|
mio.export_image(render, path, overwrite=True)
|
|
|
|
def export_shape(tmesh, path, pose_angle_deg=[0, 30, 0]):
|
|
mesh = ColouredTriMesh(tmesh.points, tmesh.trilist,
|
|
np.tile(np.array(blue), [len(tmesh.points), 1]))
|
|
render = rasterize_mesh_at_template(mesh, pose_angle_deg=pose_angle_deg, shaded=True)
|
|
mio.export_image(render, path, overwrite=True)
|
|
|
|
def uv_color_normalize(img_uv_src, angle_uv_src, img_uv, angle_uv):
|
|
intersection = (angle_uv_src.pixels[0] < 0.6) & (angle_uv_src.pixels[0] > 0.3) & (angle_uv.pixels[0] > 0.7)
|
|
if np.any(intersection):
|
|
target_dist = img_uv_src.pixels[:, intersection]
|
|
target_mean = np.mean(target_dist, 1).reshape([3, -1])
|
|
targed_std = np.std(target_dist, 1).reshape([3, -1])
|
|
|
|
source_dist = img_uv.pixels[:, intersection]
|
|
source_mean = np.mean(source_dist, 1).reshape([3, -1])
|
|
source_std = np.std(source_dist, 1).reshape([3, -1])
|
|
|
|
temp = ((img_uv.pixels.reshape([3, -1]) - source_mean) / source_std) * targed_std + target_mean
|
|
return Image(temp.reshape(img_uv.pixels.shape))
|
|
else:
|
|
return img_uv
|
|
|
|
def uv_stiching(img_uv_list, angle_uv_list, smoothing_sigma = 10, seamless_clone=False):
|
|
max_ind = np.argmax(np.array(angle_uv_list).mean(axis=1), axis=0)
|
|
mask_out_all = np.max(np.array(angle_uv_list).mean(axis=1),axis=0) ==-1
|
|
max_ind[mask_out_all] = len(img_uv_list) -1
|
|
all_uvs = np.clip(np.array(img_uv_list), 0, 1)
|
|
|
|
max_ind_one_hot = np.zeros((max_ind.size, len(img_uv_list)))
|
|
max_ind_one_hot[np.arange(max_ind.size), max_ind.flatten()] = 1
|
|
max_ind_one_hot = max_ind_one_hot.reshape(max_ind.shape + (-1,))
|
|
|
|
if smoothing_sigma>0:
|
|
max_ind_one_hot_g = gaussian(max_ind_one_hot, sigma=smoothing_sigma, multichannel=True, mode='reflect')
|
|
else:
|
|
max_ind_one_hot_g = max_ind_one_hot
|
|
|
|
max_ind_one_hot = np.tile(max_ind_one_hot, [3, 1, 1, 1])
|
|
max_ind_one_hot = np.transpose(max_ind_one_hot, [3, 0, 1, 2])
|
|
|
|
max_ind_one_hot_g = np.tile(max_ind_one_hot_g, [3, 1, 1, 1])
|
|
max_ind_one_hot_g = np.transpose(max_ind_one_hot_g, [3, 0, 1, 2])
|
|
|
|
if seamless_clone:
|
|
all_uvs_uint8 = (np.transpose(all_uvs,[0,2,3,1])*255).astype(np.uint8)
|
|
max_ind_one_uint8 = (np.transpose(1 - max_ind_one_hot,[0,2,3,1])*255).astype(np.uint8)
|
|
max_ind_one_uint8_g = (np.transpose(1 - max_ind_one_hot_g,[0,2,3,1])*255).astype(np.uint8)
|
|
|
|
dst = (Image(np.sum(all_uvs * max_ind_one_hot_g, axis=0)).pixels_with_channels_at_back()*255).astype(np.uint8)
|
|
for i, (src, mask, mask_g) in enumerate(zip(all_uvs_uint8[1:], max_ind_one_uint8[1:], max_ind_one_uint8_g[1:])):
|
|
mask_inv = 255-mask.copy()
|
|
mask_inv[:,0:5] = 255
|
|
mask_inv[:, -6:-1] = 255
|
|
mask_inv[0:5,:] = 255
|
|
mask_inv[-6:-1,:] = 255
|
|
mask_inv_dilated = binary_dilation(mask_inv.astype(np.bool)[:,:,0], disk(50)).astype(np.uint8)*255
|
|
mask_inv_dilated = np.transpose(np.tile(mask_inv_dilated,[3,1,1]),[1,2,0])
|
|
|
|
mask_inv_dilated_g = gaussian(mask_inv_dilated/255, sigma=100, multichannel=True, mode='reflect')
|
|
mixed = (src * mask_inv_dilated_g + dst * (1-mask_inv_dilated_g)).astype(np.uint8)
|
|
|
|
im_clone = cv.seamlessClone(dst, mixed, mask_inv, (512,512), cv.NORMAL_CLONE)
|
|
dst = (im_clone * (1-mask_g/255) + dst* (mask_g/255)).astype(np.uint8)
|
|
final_uv = Image(np.transpose(dst.astype(np.float)/255,[2,0,1]))
|
|
else:
|
|
final_uv = Image(np.sum(all_uvs * max_ind_one_hot_g, axis=0))
|
|
|
|
final_uv.pixels = np.clip(final_uv.pixels,0.0, 1.0)
|
|
return final_uv, max_ind_one_hot
|
|
|
|
def im_menpo2PIL(menpo_im):
|
|
return PIL.Image.fromarray((menpo_im.pixels_with_channels_at_back() * 255).astype(np.uint8))
|
|
|
|
def im_PIL2menpo(pil_im):
|
|
return Image.init_from_channels_at_back(np.array(pil_im).astype(np.float) / 255)
|
|
|
|
def fill_UV(UV):
|
|
mask = np.sum(UV.pixels, 0) == 0
|
|
xx, yy = np.meshgrid(np.arange(UV.shape[1]), np.arange(UV.shape[0]))
|
|
xym = np.vstack((np.ravel(xx[~mask]), np.ravel(yy[~mask]))).T
|
|
data = UV.pixels[:, ~mask]
|
|
for i in range(3):
|
|
interp = NearestNDInterpolator(xym, data[i])
|
|
result = interp(np.ravel(xx[mask]), np.ravel(yy[mask]))
|
|
UV.pixels[i, mask] = result
|
|
return UV
|
|
|
|
from keras.utils import get_file
|
|
import bz2
|
|
def unpack_bz2(src_path):
|
|
data = bz2.BZ2File(src_path).read()
|
|
dst_path = src_path[:-4]
|
|
with open(dst_path, 'wb') as fp:
|
|
fp.write(data)
|
|
return dst_path
|
|
|
|
|
|
def fix_obj(fp):
|
|
os.path.dirname(fp)
|
|
template = """# Produced by Dimensional Imaging OBJ exporter
|
|
# http://www.di3d.com
|
|
#
|
|
#
|
|
newmtl merged_material
|
|
Ka 0.5 0.5 0.5
|
|
Kd 0.5 0.5 0.5
|
|
Ks 0.47 0.47 0.47
|
|
d 1
|
|
Ns 0
|
|
illum 2
|
|
map_Kd {}.png
|
|
#
|
|
#
|
|
# EOF""".format(os.path.splitext(os.path.basename(fp))[0])
|
|
with open(os.path.join(os.path.dirname(fp), os.path.splitext(os.path.basename(fp))[0] + '.mtl'), 'w') as f:
|
|
f.write(template)
|
|
|
|
with open(fp, 'r+') as f:
|
|
content = f.read()
|
|
f.seek(0, 0)
|
|
f.write('mtllib ' + os.path.splitext(os.path.basename(fp))[0] + '.mtl' + '\n' + content)
|
|
|