''' Combines 3D face mode + rendering https://github.com/cleardusk/3DDFA https://github.com/YadiraF/face3d ''' import click from app.settings import types from app.utils import click_utils from app.settings import app_cfg as cfg @click.command() @click.option('-i', '--input', 'opt_fp_in', default=None, required=True, help='Image filepath') @click.option('-o', '--output', 'opt_fp_out', default=None, help='GIF output path') @click.option('--size', 'opt_size', type=(int, int), default=(300, 300), help='Output image size') @click.option('-g', '--gpu', 'opt_gpu', default=0, help='GPU index') @click.option('-f', '--force', 'opt_force', is_flag=True, help='Force overwrite file') @click.option('--bbox-init', 'opt_bbox_init', is_flag=True, help='Use landmarks for ROI instead of BBox') @click.option('--size', 'opt_render_dim', type=(int, int), default=(512, 512), help='2.5D render image size') @click.option('--display/--no-display', 'opt_display', is_flag=True, default=False, help='Display detections to debug') @click.pass_context def cli(ctx, opt_fp_in, opt_fp_out, opt_gpu, opt_bbox_init, opt_size, opt_render_dim, opt_force, opt_display): """3D face demo""" import sys import os from os.path import join from pathlib import Path import time from tqdm import tqdm import numpy as np import pandas as pd import cv2 as cv import dlib from app.models.bbox import BBox from app.utils import logger_utils, file_utils, im_utils, display_utils, draw_utils from app.utils import plot_utils from app.processors import face_detector, face_landmarks from app.models.data_store import DataStore import torch import torchvision.transforms as transforms import torch.backends.cudnn as cudnn import scipy.io as sio sys.path.append(join(Path.cwd().parent, '3rdparty')) # change name of 3DDFA to d3DDFA because can't start with number from d3DDFA import mobilenet_v1 from d3DDFA.utils.ddfa import ToTensorGjz, NormalizeGjz, str2bool from d3DDFA.utils import inference as d3dfa_utils from d3DDFA.utils.inference import parse_roi_box_from_landmark, crop_img, predict_68pts from d3DDFA.utils.inference import dump_to_ply, dump_vertex, draw_landmarks from d3DDFA.utils.inference import predict_dense, parse_roi_box_from_bbox, get_colors from d3DDFA.utils.inference import write_obj_with_colors from d3DDFA.utils.estimate_pose import parse_pose from d3DDFA.utils.render import get_depths_image, cget_depths_image, cpncc from d3DDFA.utils import paf as d3dfa_paf_utils # https://github.com/YadiraF/face3d # compile cython module in face3d/mesh/cython/ python setup.py build_ext -i from face3d.face3d import mesh as face3d_mesh log = logger_utils.Logger.getLogger() # ------------------------------------------------- # load image fpp_in = Path(opt_fp_in) im = cv.imread(opt_fp_in) #im = im_utils.resize(im_orig, width=opt_size[0], height=opt_size[1]) # im = im_orig.copy() # ---------------------------------------------------------------------------- # detect face face_detector = face_detector.DetectorDLIBCNN(gpu=opt_gpu) # -1 for CPU bboxes = face_detector.detect(im, largest=True) bbox = bboxes[0] dim = im.shape[:2][::-1] bbox_dim = bbox.to_dim(dim) if not bbox: log.error('no face detected') return else: log.info(f'face detected: {bbox_dim.to_xyxy()}') # ------------------------------------------------------------------------- # landmarks landmark_predictor = face_landmarks.Dlib2D_68() lanmarks = landmark_predictor.landmarks(im, bbox_dim) # ------------------------------------------------------------------------- # 3ddfa STD_SIZE = 120 # load pre-tained model fp_ckpt = join(cfg.DIR_MODELS_PYTORCH, '3ddfa', 'phase1_wpdc_vdc_v2.pth.tar') arch = 'mobilenet_1' checkpoint = torch.load(fp_ckpt, map_location=lambda storage, loc: storage)['state_dict'] model = getattr(mobilenet_v1, arch)(num_classes=62) # 62 = 12(pose) + 40(shape) +10(expression) model_dict = model.state_dict() # because the model is trained by multiple gpus, prefix module should be removed for k in checkpoint.keys(): model_dict[k.replace('module.', '')] = checkpoint[k] model.load_state_dict(model_dict, strict=False) if opt_gpu > -1: cudnn.benchmark = True model = model.cuda() model.eval() # forward st = time.time() fp_tri = join(cfg.DIR_MODELS_PYTORCH, '3ddfa', 'tri.mat') triangles = sio.loadmat(fp_tri)['tri'] transform = transforms.Compose([ToTensorGjz(), NormalizeGjz(mean=127.5, std=128)]) pts_res = [] Ps = [] # Camera matrix collection poses = [] # pose collection, [todo: validate it] vertices_lst = [] # store multiple face vertices # use landmark as roi pts = np.array(lanmarks).T # roi_box = d3dfa_utils.parse_roi_box_from_landmark(pts) roi_box = parse_roi_box_from_bbox(bbox_dim.to_xyxy()) im_crop = d3dfa_utils.crop_img(im, roi_box) im_crop = cv.resize(im_crop, dsize=(STD_SIZE, STD_SIZE), interpolation=cv.INTER_LINEAR) # forward torch_input = transform(im_crop).unsqueeze(0) with torch.no_grad(): if opt_gpu > -1: torch_input = torch_input.cuda() param = model(torch_input) param = param.squeeze().cpu().numpy().flatten().astype(np.float32) # 68 pts pts68 = d3dfa_utils.predict_68pts(param, roi_box) pts_res.append(pts68) P, pose = parse_pose(param) Ps.append(P) poses.append(pose) # dense face 3d vertices vertices = d3dfa_utils.predict_dense(param, roi_box) vertices_lst.append(vertices) log.info(f'generated 3d data in: {(time.time() - st):.2f}s') # filepath helper function def to_fp(fpp, ext, suffix=None): if suffix: fp = join(fpp.parent, f'{fpp.stem}_{suffix}.{ext}') else: fp = join(fpp.parent, f'{fpp.stem}.{ext}') return fp # save .mat colors = d3dfa_utils.get_colors(im, vertices) vertices_orig = vertices.copy() fp_mat_3df = to_fp(fpp_in, 'mat', suffix='face3d') sio.savemat(fp_mat_3df, {'vertices': vertices, 'colors': colors, 'triangles': triangles}) # save PAF #fp_paf = to_fp(fpp_in, 'jpg', suffix='paf') #opt_paf_size = 3 # PAF feature kernel size #im_paf = d3dfa_paf_utils.gen_img_paf(img_crop=im_crop, param=param, kernel_size=opt_paf_size) #cv.imwrite(fp_paf, im_paf) # save pose image # P, pose = parse_pose(param) # Camera matrix (without scale), and pose (yaw, pitch, roll, to verify) img_pose = draw_utils.plot_pose_box(im, Ps, pts_res) fp_pose = to_fp(fpp_in, 'jpg', suffix='pose') cv.imwrite(fp_pose, img_pose) # save depth image fp_depth = to_fp(fpp_in, 'png', suffix='depth') # depths_img = get_depths_image(im, vertices_lst, tri-1) # python version im_depth = cget_depths_image(im, vertices_lst, triangles - 1) # cython version cv.imwrite(fp_depth, im_depth) # save pncc image fp_pose = to_fp(fpp_in, 'png', suffix='pncc') pncc_feature = cpncc(im, vertices_lst, triangles - 1) # cython version cv.imwrite(fp_pose, pncc_feature[:, :, ::-1]) # cv.imwrite will swap RGB -> BGR # save .ply #fp_ply = to_fp(fpp_in, 'ply') #dump_to_ply(vertices, triangles, fp_ply) # skip: save .mat (3ddfa default not compatible with face3d utils) #fp_mat = to_fp(fpp_in, 'mat') #d3dfa_utils.dump_vertex(vertices, fp_mat) # save 68 points #fp_txt = to_fp(fpp_in, 'txt', suffix='68') #np.savetxt(to_fp(fpp_in, 'txt'), pts68, fmt='%.3f') # save roi #fp_txt = to_fp(fpp_in, 'txt', suffix='roi') #np.savetxt(fp_txt, roi_box, fmt='%.3f')a # save crop #fp_crop = to_fp(fpp_in, 'jpg', suffix='crop') #cv.imwrite(fp_crop, im_crop) # save obj colors = d3dfa_utils.get_colors(im, vertices_orig) fp_obj = to_fp(fpp_in, 'obj') write_obj_with_colors(fp_obj, vertices_orig, triangles, colors) #fp_landmarks = to_fp(fpp_in, 'jpg', suffix='3DDFA') # show_flg? #d3dfa_utils.draw_landmarks(im, pts_res, wfp=fp_landmarks, show_flg=False) # ------------------------------------------------------------------------- # face3d # create 3D mesh photo face # if loading file # TODO find where vertices is being changed vertices = vertices_orig # vertices changes somewhere, so keep copy # preprocess 3D data from 3DDFA for face3d rendering vertices = vertices.transpose() triangles = triangles.transpose() vertices = vertices.astype(np.float64) # change data type # subtract 1 from triangle vertex indices (depends on your .mat file) triangles = np.array([np.array([t[0]-1, t[1]-1, t[2]-1]).astype(np.int32) for t in triangles]) vertices -= np.array([abs(np.min(vertices[:,0])), np.min(abs(vertices[:,1])), np.min(abs(vertices[:,2]))]) vertices -= np.array([np.mean(vertices[:,0]), np.mean(vertices[:,1]), np.mean(vertices[:,2])]) # colors = np.array([c[::-1] for c in colors]) # BGR --> RGB colors = colors/np.max(colors) # normalize color range # set max render size (about 75% of canvas size) max_render_size = int(max(opt_render_dim) * .75) s = max_render_size/(np.max(vertices[:,1]) - np.min(vertices[:,1])) # rotation matrix R = face3d_mesh.transform.angle2matrix([-180, -20, 0]) # no translation. center of obj:[0,0] t = [0, 0, 0] vertices_trans = face3d_mesh.transform.similarity_transform(vertices, s, R, t) # lighting: add point lights, positions are defined in world space light_pos = np.array([[-128, -128, 512]]) light_clr_amt = np.array([[1, 1, 1]]) colors_lit = face3d_mesh.light.add_light(vertices_trans, triangles, colors, light_pos, light_clr_amt) # transform from world space to camera space (what the world is in the eye of observer) vertices_cam = face3d_mesh.transform.lookat_camera(vertices_trans, eye = [0, 0, 0], at = np.array([0, 0, 1]), up = None) # project from 3d world space into 2d image plane. orthographic or perspective projection vertices_proj = face3d_mesh.transform.orthographic_project(vertices_cam) # ------------------------------------------------------------------------- # render 2D image w = h = max(opt_render_dim) vertices_im = face3d_mesh.transform.to_image(vertices_proj, h, w) rendering = face3d_mesh.render.render_colors(vertices_im, triangles, colors_lit, h, w) cv.imshow('', rendering) display_utils.handle_keyboard() # ---------------------------------------------------------------------------- # save if opt_fp_out: # save pose only fpp_out = Path(opt_fp_out) fp_out = join(fpp_out.parent, f'{fpp_out.stem}_real{fpp_out.suffix}') cv.imwrite(fp_out, im_age_real) fp_out = join(fpp_out.parent, f'{fpp_out.stem}_apparent{fpp_out.suffix}') cv.imwrite(fp_out, im_age_apparent) fp_out = join(fpp_out.parent, f'{fpp_out.stem}_gender{fpp_out.suffix}') cv.imwrite(fp_out, im_age_apparent) # ---------------------------------------------------------------------------- # display if opt_display: # show all images here cv.imshow('real', im_age_real) cv.imshow('apparent', im_age_apparent) cv.imshow('gender', im_gender) display_utils.handle_keyboard()