extract face depth of all images in a directory with 3DDFA_V2

depth.py

"""
3DDFA_V2: https://github.com/cleardusk/3DDFA_V2
most of the code comes from their example Google Colab project, some of it
was slightly altered
ran this on the 3DDFA_V2 directory, so you might have to change some paths
check original code to see other options
"""
import cv2
import yaml

from FaceBoxes import FaceBoxes
from TDDFA import TDDFA
from utils.render import render
from utils.depth import depth
from utils.pncc import pncc
from utils.uv import uv_tex
from utils.pose import viz_pose
from utils.serialization import ser_to_ply, ser_to_obj
from utils.functions import draw_landmarks, get_suffix

import matplotlib.pyplot as plt
from skimage import io

import sys
import os

def save_depth(img_fp, target_fp):
    print(f"Getting image from path {img_fp}")
    img = cv2.imread(img_fp)
    # plt.imshow(img[..., ::-1])
    # face detection
    boxes = face_boxes(img)
    if len(boxes) > 1:
        print(f"Detected {len(boxes)} faces @ {img_fp}!")
        print(boxes)
    # regress 3DMM params
    param_lst, roi_box_lst = tddfa(img, boxes)
    # reconstruct vertices and render depth
    ver_lst = tddfa.recon_vers(param_lst, roi_box_lst, dense_flag=True)
    return depth(
            img, ver_lst, tddfa.tri,
            show_flag=False, wfp=target_fp,
            with_bg_flag=False
            )


# load config
cfg = yaml.load(open('configs/mb1_120x120.yml'), Loader=yaml.SafeLoader)

# Init FaceBoxes and TDDFA, recommend using onnx flag
onnx_flag = True  # or True to use ONNX to speed up
if onnx_flag:
    # !pip install onnxruntime

    import os
    os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
    os.environ['OMP_NUM_THREADS'] = '4'
    from FaceBoxes.FaceBoxes_ONNX import FaceBoxes_ONNX
    from TDDFA_ONNX import TDDFA_ONNX

    face_boxes = FaceBoxes_ONNX()
    tddfa = TDDFA_ONNX(**cfg)
else:
    face_boxes = FaceBoxes()
    tddfa = TDDFA(gpu_mode=False, **cfg)

# python3 depth.py sourcedir targetdir
sources = [s + '/' if (s[-1] != '/') else s for s in sys.argv[1:-1]]
target = sys.argv[-1]
if target[-1] != '/':
    target += '/'

for s in sources:
    for f in os.listdir(s):
        img_fp = s + f
        target_fp = target + f
        save_depth(img_fp, target_fp)