add 3d render

1 files changed, 314 insertions, 0 deletions

diff --git a/megapixels/commands/demo/face_3ddfa.py b/megapixels/commands/demo/face_3ddfa.py
new file mode 100644
index 00000000..6182aeb6
--- /dev/null
+++ b/megapixels/commands/demo/face_3ddfa.py
@@ -0,0 +1,314 @@
+'''
+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()
+