| | import os, sys |
| | import cv2 |
| | import numpy as np |
| | from time import time |
| | from scipy.io import savemat,loadmat |
| | import argparse |
| | from tqdm import tqdm |
| | import torch |
| | import mmcv |
| |
|
| | sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) |
| | from decalib.deca import DECA |
| | from decalib.datasets import datasets |
| | from decalib.utils import util |
| | from decalib.utils.config import get_cfg_defaults |
| | from decalib.utils.tensor_cropper import transform_points |
| | from loguru import logger |
| |
|
| | @logger.catch |
| | def api_multi_deca( |
| | savefolder, |
| | visfolder, |
| | inputpath, |
| | |
| | device='cuda', |
| | saveVis=True, |
| | saveMat=True, |
| | rasterizer_type='pytorch3d', |
| | face_detector='mtcnn', |
| | ): |
| | deca_cfg=get_cfg_defaults() |
| | os.makedirs(savefolder, exist_ok=True) |
| | os.makedirs(visfolder, exist_ok=True) |
| | print(f'-- please check the results in {savefolder}') |
| | testdata = datasets.TestData(inputpath, iscrop=True, face_detector=face_detector) |
| | deca_cfg.model.use_tex = False |
| | deca_cfg.rasterizer_type = rasterizer_type |
| | deca = DECA(config = deca_cfg, device=device) |
| | |
| | for batch_list in tqdm(testdata): |
| | if isinstance(batch_list,dict) and batch_list.get('is_missing',None): |
| | name=batch_list['name'] |
| | open(os.path.join(savefolder, f'{name}.pkl.empty'), 'a').close() |
| | print('no face detected! skip： {name}') |
| | continue |
| |
|
| | original_image=None |
| | tobj_returns_list=[] |
| | for index,tobj in enumerate(batch_list): |
| | name = tobj['imagename'] |
| | images = tobj['dst_image'].to(device)[None,...] |
| |
|
| | with torch.no_grad(): |
| | codedict = deca.encode(images) |
| | opdict, visdict = deca.decode(codedict) |
| | |
| | |
| | tform = tobj['tform'][None, ...] |
| | tform = torch.inverse(tform).transpose(1,2).to(device) |
| | |
| | if original_image is None: |
| | original_image = tobj['original_image'][None, ...].to(device) |
| | elif orig_visdict['shape_images'] is not None: |
| | original_image=orig_visdict['shape_images'] |
| | |
| | _, orig_visdict = deca.decode(codedict, render_orig=True, original_image=original_image, tform=tform) |
| |
|
| | tobj_returns = util.dict_tensor2npy( |
| | { |
| | 'bbox':tobj['bbox'], |
| | 'tform': tobj['tform'], |
| | 'imagename': tobj['imagename'], |
| | 'landmarks2d':opdict['landmarks2d'], |
| | 'shape_params':codedict['shape'], |
| | 'expression_params':codedict['exp'], |
| | 'pose_params':codedict['pose'] |
| | }) |
| | tobj_returns_list.append(tobj_returns) |
| | |
| | |
| | if saveVis: |
| | cv2.imwrite( |
| | os.path.join(visfolder, name + '.jpg'), |
| | deca.visualize({ |
| | "shape_images":orig_visdict['shape_images'], |
| | })) |
| | |
| | if saveMat: |
| | mmcv.dump(tobj_returns_list,os.path.join(savefolder, name + '.pkl')) |
| |
|
| | print('deca ended') |
| | |
| | |
| | |
