RabbitRUI/ruispace

1import os2import numpy as np3from PIL import Image4from skimage import io, img_as_float32, transform5import torch6import scipy.io as scio7 8def get_facerender_data(coeff_path, pic_path, first_coeff_path, audio_path, 9                        batch_size, input_yaw_list=None, input_pitch_list=None, input_roll_list=None, 10                        expression_scale=1.0, still_mode = False, preprocess='crop'):11 12    semantic_radius = 1313    video_name = os.path.splitext(os.path.split(coeff_path)[-1])[0]14    txt_path = os.path.splitext(coeff_path)[0]15 16    data={}17 18    img1 = Image.open(pic_path)19    source_image = np.array(img1)20    source_image = img_as_float32(source_image)21    source_image = transform.resize(source_image, (256, 256, 3))22    source_image = source_image.transpose((2, 0, 1))23    source_image_ts = torch.FloatTensor(source_image).unsqueeze(0)24    source_image_ts = source_image_ts.repeat(batch_size, 1, 1, 1)25    data['source_image'] = source_image_ts26 27    source_semantics_dict = scio.loadmat(first_coeff_path)28 29    if preprocess.lower() != 'full':30        source_semantics = source_semantics_dict['coeff_3dmm'][:1,:70]         #1 7031    else:32        source_semantics = source_semantics_dict['coeff_3dmm'][:1,:73]         #1 7033 34    source_semantics_new = transform_semantic_1(source_semantics, semantic_radius)35    source_semantics_ts = torch.FloatTensor(source_semantics_new).unsqueeze(0)36    source_semantics_ts = source_semantics_ts.repeat(batch_size, 1, 1)37    data['source_semantics'] = source_semantics_ts38 39    # target 40    generated_dict = scio.loadmat(coeff_path)41    generated_3dmm = generated_dict['coeff_3dmm']42    generated_3dmm[:, :64] = generated_3dmm[:, :64] * expression_scale43 44    if preprocess.lower() == 'full':45        generated_3dmm = np.concatenate([generated_3dmm, np.repeat(source_semantics[:,70:], generated_3dmm.shape[0], axis=0)], axis=1)46 47    if still_mode:48        generated_3dmm[:, 64:] = np.repeat(source_semantics[:, 64:], generated_3dmm.shape[0], axis=0)49 50    with open(txt_path+'.txt', 'w') as f:51        for coeff in generated_3dmm:52            for i in coeff:53                f.write(str(i)[:7]   + '  '+'\t')54            f.write('\n')55 56    target_semantics_list = [] 57    frame_num = generated_3dmm.shape[0]58    data['frame_num'] = frame_num59    for frame_idx in range(frame_num):60        target_semantics = transform_semantic_target(generated_3dmm, frame_idx, semantic_radius)61        target_semantics_list.append(target_semantics)62 63    remainder = frame_num%batch_size64    if remainder!=0:65        for _ in range(batch_size-remainder):66            target_semantics_list.append(target_semantics)67 68    target_semantics_np = np.array(target_semantics_list)             #frame_num 70 semantic_radius*2+169    target_semantics_np = target_semantics_np.reshape(batch_size, -1, target_semantics_np.shape[-2], target_semantics_np.shape[-1])70    data['target_semantics_list'] = torch.FloatTensor(target_semantics_np)71    data['video_name'] = video_name72    data['audio_path'] = audio_path73    74    if input_yaw_list is not None:75        yaw_c_seq = gen_camera_pose(input_yaw_list, frame_num, batch_size)76        data['yaw_c_seq'] = torch.FloatTensor(yaw_c_seq)77    if input_pitch_list is not None:78        pitch_c_seq = gen_camera_pose(input_pitch_list, frame_num, batch_size)79        data['pitch_c_seq'] = torch.FloatTensor(pitch_c_seq)80    if input_roll_list is not None:81        roll_c_seq = gen_camera_pose(input_roll_list, frame_num, batch_size) 82        data['roll_c_seq'] = torch.FloatTensor(roll_c_seq)83 84    return data85 86def transform_semantic_1(semantic, semantic_radius):87    semantic_list =  [semantic for i in range(0, semantic_radius*2+1)]88    coeff_3dmm = np.concatenate(semantic_list, 0)89    return coeff_3dmm.transpose(1,0)90 91def transform_semantic_target(coeff_3dmm, frame_index, semantic_radius):92    num_frames = coeff_3dmm.shape[0]93    seq = list(range(frame_index- semantic_radius, frame_index + semantic_radius+1))94    index = [ min(max(item, 0), num_frames-1) for item in seq ] 95    coeff_3dmm_g = coeff_3dmm[index, :]96    return coeff_3dmm_g.transpose(1,0)97 98def gen_camera_pose(camera_degree_list, frame_num, batch_size):99 100    new_degree_list = [] 101    if len(camera_degree_list) == 1:102        for _ in range(frame_num):103            new_degree_list.append(camera_degree_list[0]) 104        remainder = frame_num%batch_size105        if remainder!=0:106            for _ in range(batch_size-remainder):107                new_degree_list.append(new_degree_list[-1])108        new_degree_np = np.array(new_degree_list).reshape(batch_size, -1) 109        return new_degree_np110 111    degree_sum = 0.112    for i, degree in enumerate(camera_degree_list[1:]):113        degree_sum += abs(degree-camera_degree_list[i])114    115    degree_per_frame = degree_sum/(frame_num-1)116    for i, degree in enumerate(camera_degree_list[1:]):117        degree_last = camera_degree_list[i]118        degree_step = degree_per_frame * abs(degree-degree_last)/(degree-degree_last)119        new_degree_list =  new_degree_list + list(np.arange(degree_last, degree, degree_step))120    if len(new_degree_list) > frame_num:121        new_degree_list = new_degree_list[:frame_num]122    elif len(new_degree_list) < frame_num:123        for _ in range(frame_num-len(new_degree_list)):124            new_degree_list.append(new_degree_list[-1])125    print(len(new_degree_list))126    print(frame_num)127 128    remainder = frame_num%batch_size129    if remainder!=0:130        for _ in range(batch_size-remainder):131            new_degree_list.append(new_degree_list[-1])132    new_degree_np = np.array(new_degree_list).reshape(batch_size, -1) 133    return new_degree_np134    135