data_gen/runs/binarizer_nerf.py

import os
import numpy as np
import math
import json
import imageio
import torch
import tqdm
import cv2

from data_util.face3d_helper import Face3DHelper
from utils.commons.euler2rot import euler_trans_2_c2w, c2w_to_euler_trans
from data_gen.utils.process_video.euler2quaterion import euler2quaterion, quaterion2euler
from deep_3drecon.deep_3drecon_models.bfm import ParametricFaceModel


def euler2rot(euler_angle):
    batch_size = euler_angle.shape[0]
    theta = euler_angle[:, 0].reshape(-1, 1, 1)
    phi = euler_angle[:, 1].reshape(-1, 1, 1)
    psi = euler_angle[:, 2].reshape(-1, 1, 1)
    one = torch.ones(batch_size, 1, 1).to(euler_angle.device)
    zero = torch.zeros(batch_size, 1, 1).to(euler_angle.device)
    rot_x = torch.cat((
        torch.cat((one, zero, zero), 1),
        torch.cat((zero, theta.cos(), theta.sin()), 1),
        torch.cat((zero, -theta.sin(), theta.cos()), 1),
    ), 2)
    rot_y = torch.cat((
        torch.cat((phi.cos(), zero, -phi.sin()), 1),
        torch.cat((zero, one, zero), 1),
        torch.cat((phi.sin(), zero, phi.cos()), 1),
    ), 2)
    rot_z = torch.cat((
        torch.cat((psi.cos(), -psi.sin(), zero), 1),
        torch.cat((psi.sin(), psi.cos(), zero), 1),
        torch.cat((zero, zero, one), 1)
    ), 2)
    return torch.bmm(rot_x, torch.bmm(rot_y, rot_z))


def rot2euler(rot_mat):
    batch_size = len(rot_mat)
    # we assert that y in in [-0.5pi, 0.5pi]
    cos_y = torch.sqrt(rot_mat[:, 1, 2] * rot_mat[:, 1, 2] + rot_mat[:, 2, 2] * rot_mat[:, 2, 2])
    theta_x = torch.atan2(-rot_mat[:, 1, 2], rot_mat[:, 2, 2])
    theta_y = torch.atan2(rot_mat[:, 2, 0], cos_y)
    theta_z = torch.atan2(rot_mat[:, 0, 1], rot_mat[:, 0, 0])
    euler_angles = torch.zeros([batch_size, 3])
    euler_angles[:, 0] = theta_x
    euler_angles[:, 1] = theta_y
    euler_angles[:, 2] = theta_z
    return euler_angles

index_lm68_from_lm468 = [127,234,93,132,58,136,150,176,152,400,379,365,288,361,323,454,356,70,63,105,66,107,336,296,334,293,300,168,197,5,4,75,97,2,326,305,
                         33,160,158,133,153,144,362,385,387,263,373,380,61,40,37,0,267,270,291,321,314,17,84,91,78,81,13,311,308,402,14,178]

def plot_lm2d(lm2d):
    WH = 512
    img = np.ones([WH, WH, 3], dtype=np.uint8) * 255
    
    for i in range(len(lm2d)):
        x, y = lm2d[i]
        color = (255,0,0)
        img = cv2.circle(img, center=(int(x),int(y)), radius=3, color=color, thickness=-1)
        font = cv2.FONT_HERSHEY_SIMPLEX
    for i in range(len(lm2d)):
        x, y = lm2d[i]
        img = cv2.putText(img, f"{i}", org=(int(x),int(y)), fontFace=font, fontScale=0.3, color=(255,0,0))
    return img

def get_face_rect(lms, h, w):
    """
    lms: [68, 2]
    h, w: int
    return: [4,]
    """
    assert len(lms) == 68
    # min_x, max_x = np.min(lms, 0)[0], np.max(lms, 0)[0]
    min_x, max_x = np.min(lms[:, 0]), np.max(lms[:, 0])
    cx = int((min_x+max_x)/2.0)
    cy = int(lms[27, 1])
    h_w = int((max_x-cx)*1.5)
    h_h = int((lms[8, 1]-cy)*1.15)
    rect_x = cx - h_w
    rect_y = cy - h_h
    if rect_x < 0:
        rect_x = 0
    if rect_y < 0:
        rect_y = 0
    rect_w = min(w-1-rect_x, 2*h_w)
    rect_h = min(h-1-rect_y, 2*h_h)
    # rect = np.array((rect_x, rect_y, rect_w, rect_h), dtype=np.int32)
    # rect = [rect_x, rect_y, rect_w, rect_h]
    rect = [rect_x, rect_x + rect_w, rect_y, rect_y + rect_h] # min_j,  max_j, min_i, max_i
    return rect # this x is width, y is height

def get_lip_rect(lms, h, w):
    """
    lms: [68, 2]
    h, w: int
    return: [4,]
    """
    # this x is width, y is height
    # for lms, lms[:, 0] is width, lms[:, 1] is height
    assert len(lms) == 68
    lips = slice(48, 60)
    lms = lms[lips]
    min_x, max_x = np.min(lms[:, 0]), np.max(lms[:, 0])
    min_y, max_y = np.min(lms[:, 1]), np.max(lms[:, 1])
    cx = int((min_x+max_x)/2.0)
    cy = int((min_y+max_y)/2.0)
    h_w = int((max_x-cx)*1.2)
    h_h = int((max_y-cy)*1.2)
    
    h_w = max(h_w, h_h)
    h_h = h_w

    rect_x = cx - h_w
    rect_y = cy - h_h
    rect_w = 2*h_w
    rect_h = 2*h_h
    if rect_x < 0:
        rect_x = 0
    if rect_y < 0:
        rect_y = 0
    
    if rect_x + rect_w > w:
        rect_x = w - rect_w
    if rect_y + rect_h > h:
        rect_y = h - rect_h

    rect = [rect_x, rect_x + rect_w, rect_y, rect_y + rect_h] # min_j,  max_j, min_i, max_i
    return rect # this x is width, y is height


# def get_lip_rect(lms, h, w):
#     """
#     lms: [68, 2]
#     h, w: int
#     return: [4,]
#     """
#     assert len(lms) == 68
#     lips = slice(48, 60)
#     # this x is width, y is height
#     xmin, xmax = int(lms[lips, 1].min()), int(lms[lips, 1].max())
#     ymin, ymax = int(lms[lips, 0].min()), int(lms[lips, 0].max())
#     # padding to H == W
#     cx = (xmin + xmax) // 2
#     cy = (ymin + ymax) // 2
#     l = max(xmax - xmin, ymax - ymin) // 2
#     xmin = max(0, cx - l)
#     xmax = min(h, cx + l)
#     ymin = max(0, cy - l)
#     ymax = min(w, cy + l)
#     lip_rect = [xmin, xmax, ymin, ymax]
#     return lip_rect

def get_win_conds(conds, idx, smo_win_size=8, pad_option='zero'):
    """
    conds: [b, t=16, h=29]
    idx: long, time index of the selected frame
    """
    idx = max(0, idx)
    idx = min(idx, conds.shape[0]-1)
    smo_half_win_size = smo_win_size//2
    left_i = idx - smo_half_win_size
    right_i = idx + (smo_win_size - smo_half_win_size)
    pad_left, pad_right = 0, 0
    if left_i < 0:
        pad_left = -left_i
        left_i = 0
    if right_i > conds.shape[0]:
        pad_right = right_i - conds.shape[0]
        right_i = conds.shape[0]
    conds_win = conds[left_i:right_i]
    if pad_left > 0:
        if pad_option == 'zero':
            conds_win = np.concatenate([np.zeros_like(conds_win)[:pad_left], conds_win], axis=0)
        elif pad_option == 'edge':
            edge_value = conds[0][np.newaxis, ...]
            conds_win = np.concatenate([edge_value] * pad_left + [conds_win], axis=0)
        else: 
            raise NotImplementedError
    if pad_right > 0:
        if pad_option == 'zero':
            conds_win = np.concatenate([conds_win, np.zeros_like(conds_win)[:pad_right]], axis=0)
        elif pad_option == 'edge':
            edge_value = conds[-1][np.newaxis, ...]
            conds_win = np.concatenate([conds_win] + [edge_value] * pad_right , axis=0)
        else: 
            raise NotImplementedError
    assert conds_win.shape[0] == smo_win_size
    return conds_win


def load_processed_data(processed_dir):
    # load necessary files
    background_img_name = os.path.join(processed_dir, "bg.jpg")
    assert os.path.exists(background_img_name)
    head_img_dir = os.path.join(processed_dir, "head_imgs")
    torso_img_dir = os.path.join(processed_dir, "inpaint_torso_imgs")
    gt_img_dir = os.path.join(processed_dir, "gt_imgs")

    hubert_npy_name = os.path.join(processed_dir, "aud_hubert.npy")
    mel_f0_npy_name = os.path.join(processed_dir, "aud_mel_f0.npy")
    coeff_npy_name = os.path.join(processed_dir, "coeff_fit_mp.npy")
    lm2d_npy_name = os.path.join(processed_dir, "lms_2d.npy")
    
    ret_dict = {}

    ret_dict['bg_img'] = imageio.imread(background_img_name)
    ret_dict['H'], ret_dict['W'] = ret_dict['bg_img'].shape[:2]
    ret_dict['focal'], ret_dict['cx'], ret_dict['cy'] = face_model.focal, face_model.center, face_model.center

    print("loading lm2d coeff ...")
    lm2d_arr = np.load(lm2d_npy_name)
    face_rect_lst = []
    lip_rect_lst = []
    for lm2d in lm2d_arr:
        if len(lm2d) in [468, 478]:
            lm2d = lm2d[index_lm68_from_lm468]
        face_rect = get_face_rect(lm2d, ret_dict['H'], ret_dict['W'])
        lip_rect = get_lip_rect(lm2d, ret_dict['H'], ret_dict['W'])
        face_rect_lst.append(face_rect)
        lip_rect_lst.append(lip_rect)
    face_rects = np.stack(face_rect_lst, axis=0) # [T, 4]

    print("loading fitted 3dmm coeff ...")
    coeff_dict = np.load(coeff_npy_name, allow_pickle=True).tolist()
    identity_arr = coeff_dict['id']
    exp_arr = coeff_dict['exp']
    ret_dict['id'] = identity_arr
    ret_dict['exp'] = exp_arr
    euler_arr = ret_dict['euler'] = coeff_dict['euler']
    trans_arr = ret_dict['trans'] = coeff_dict['trans']
    print("calculating lm3d ...")
    idexp_lm3d_arr = face3d_helper.reconstruct_idexp_lm3d(torch.from_numpy(identity_arr), torch.from_numpy(exp_arr)).cpu().numpy().reshape([-1, 68*3])
    len_motion = len(idexp_lm3d_arr)
    video_idexp_lm3d_mean = idexp_lm3d_arr.mean(axis=0)
    video_idexp_lm3d_std = idexp_lm3d_arr.std(axis=0)
    ret_dict['idexp_lm3d'] = idexp_lm3d_arr
    ret_dict['idexp_lm3d_mean'] = video_idexp_lm3d_mean
    ret_dict['idexp_lm3d_std'] = video_idexp_lm3d_std
    
    # now we convert the euler_trans from deep3d convention to adnerf convention
    eulers = torch.FloatTensor(euler_arr)
    trans = torch.FloatTensor(trans_arr)
    rots = face_model.compute_rotation(eulers) # rotation matrix is a better intermediate for convention-transplan than euler

    # handle the camera pose to geneface's convention
    trans[:, 2] = 10 - trans[:, 2] # 抵消fit阶段的to_camera操作，即trans[...,2] = 10 - trans[...,2]
    rots = rots.permute(0, 2, 1)
    trans[:, 2] = - trans[:,2] # 因为intrinsic proj不同
    # below is the NeRF camera preprocessing strategy, see `save_transforms` in data_util/process.py 
    trans = trans / 10.0
    rots_inv = rots.permute(0, 2, 1)
    trans_inv = - torch.bmm(rots_inv, trans.unsqueeze(2))

    pose = torch.eye(4, dtype=torch.float32).unsqueeze(0).repeat([len_motion, 1, 1]) # [T, 4, 4]
    pose[:, :3, :3] = rots_inv
    pose[:, :3, 3] = trans_inv[:, :, 0]
    c2w_transform_matrices = pose.numpy()

    # process the audio features used for postnet training
    print("loading hubert ...")
    hubert_features = np.load(hubert_npy_name)
    print("loading Mel and F0 ...")
    mel_f0_features = np.load(mel_f0_npy_name, allow_pickle=True).tolist()

    ret_dict['hubert'] = hubert_features
    ret_dict['mel'] = mel_f0_features['mel']
    ret_dict['f0'] = mel_f0_features['f0']

    # obtaining train samples
    frame_indices = list(range(len_motion))
    num_train = len_motion // 11 * 10
    train_indices = frame_indices[:num_train]
    val_indices = frame_indices[num_train:]

    for split in ['train', 'val']:
        if split == 'train':
            indices = train_indices
            samples = []
            ret_dict['train_samples'] = samples
        elif split == 'val':
            indices = val_indices
            samples = []
            ret_dict['val_samples'] = samples
        
        for idx in indices:
            sample = {}
            sample['idx'] = idx
            sample['head_img_fname'] = os.path.join(head_img_dir,f"{idx:08d}.png")
            sample['torso_img_fname'] = os.path.join(torso_img_dir,f"{idx:08d}.png")
            sample['gt_img_fname'] = os.path.join(gt_img_dir,f"{idx:08d}.jpg")
            # assert os.path.exists(sample['head_img_fname']) and os.path.exists(sample['torso_img_fname']) and os.path.exists(sample['gt_img_fname'])
            sample['face_rect'] = face_rects[idx]
            sample['lip_rect'] = lip_rect_lst[idx]
            sample['c2w'] = c2w_transform_matrices[idx]
            samples.append(sample)
    return ret_dict


class Binarizer:
    def __init__(self):
        self.data_dir = 'data/'
        
    def parse(self, video_id):
        processed_dir = os.path.join(self.data_dir, 'processed/videos', video_id)
        binary_dir = os.path.join(self.data_dir, 'binary/videos', video_id)
        out_fname = os.path.join(binary_dir, "trainval_dataset.npy")
        os.makedirs(binary_dir, exist_ok=True)
        ret = load_processed_data(processed_dir)
        mel_name = os.path.join(processed_dir, 'aud_mel_f0.npy')
        mel_f0_dict = np.load(mel_name, allow_pickle=True).tolist()
        ret.update(mel_f0_dict)
        np.save(out_fname, ret, allow_pickle=True)



if __name__ == '__main__':
    from argparse import ArgumentParser
    parser = ArgumentParser()
    parser.add_argument('--video_id', type=str, default='May', help='')
    args = parser.parse_args()
    ### Process Single Long Audio for NeRF dataset
    video_id = args.video_id
    face_model = ParametricFaceModel(bfm_folder='deep_3drecon/BFM', 
                camera_distance=10, focal=1015)
    face_model.to("cpu")
    face3d_helper = Face3DHelper()

    binarizer = Binarizer()
    binarizer.parse(video_id)
    print(f"Binarization for {video_id} Done!")