Upload 6 files

demoworking.py

@@ -0,0 +1,470 @@
+#@title demo.py with fixed paths
+
+import os
+import numpy as np
+import torch
+import yaml
+from modules.generator import OcclusionAwareSPADEGeneratorEam
+from modules.keypoint_detector import KPDetector, HEEstimator
+import argparse
+import imageio
+from modules.transformer import Audio2kpTransformerBBoxQDeepPrompt as Audio2kpTransformer
+from modules.prompt import EmotionDeepPrompt, EmotionalDeformationTransformer
+from scipy.io import wavfile
+
+from modules.model_transformer import get_rotation_matrix, keypoint_transformation
+from skimage import io, img_as_float32
+from skimage.transform import resize
+import torchaudio
+import soundfile as sf
+from scipy.spatial import ConvexHull
+
+import torch.nn.functional as F
+import glob
+from tqdm import tqdm
+import gzip
+
+emo_label = ['ang',  'con',  'dis',  'fea',  'hap',  'neu',  'sad',  'sur']
+emo_label_full = ['angry',  'contempt',  'disgusted',  'fear',  'happy',  'neutral',  'sad',  'surprised']
+latent_dim = 16
+
+MEL_PARAMS_25 = {
+    "n_mels": 80,
+    "n_fft": 2048,
+    "win_length": 640,
+    "hop_length": 640
+}
+
+to_melspec = torchaudio.transforms.MelSpectrogram(**MEL_PARAMS_25)
+mean, std = -4, 4
+
+expU = torch.from_numpy(np.load('/content/EAT_code/expPCAnorm_fin/U_mead.npy')[:,:32])
+expmean = torch.from_numpy(np.load('/content/EAT_code/expPCAnorm_fin/mean_mead.npy'))
+
+root_wav = '/content/EAT_code/demo/video_processed/bo_1resized'
+def normalize_kp(kp_source, kp_driving, kp_driving_initial,
+                 use_relative_movement=True, use_relative_jacobian=True):
+
+    kp_new = {k: v for k, v in kp_driving.items()}
+    if use_relative_movement:
+        kp_value_diff = (kp_driving['value'] - kp_driving_initial['value'])
+        kp_new['value'] = kp_value_diff + kp_source['value']
+
+        if use_relative_jacobian:
+            jacobian_diff = torch.matmul(kp_driving['jacobian'], torch.inverse(kp_driving_initial['jacobian']))
+            kp_new['jacobian'] = torch.matmul(jacobian_diff, kp_source['jacobian'])
+
+    return kp_new
+
+def _load_tensor(data):
+    wave_path = data
+    wave, sr = sf.read(wave_path)
+    wave_tensor = torch.from_numpy(wave).float()
+    return wave_tensor
+
+def build_model(config, device_ids=[0]):
+    generator = OcclusionAwareSPADEGeneratorEam(**config['model_params']['generator_params'],
+                                                 **config['model_params']['common_params'])
+    if torch.cuda.is_available():
+        print('cuda is available')
+        generator.to(device_ids[0])
+
+    kp_detector = KPDetector(**config['model_params']['kp_detector_params'],
+                             **config['model_params']['common_params'])
+
+    if torch.cuda.is_available():
+        kp_detector.to(device_ids[0])
+
+
+    audio2kptransformer = Audio2kpTransformer(**config['model_params']['audio2kp_params'], face_ea=True)
+
+    if torch.cuda.is_available():
+        audio2kptransformer.to(device_ids[0])
+
+    sidetuning = EmotionalDeformationTransformer(**config['model_params']['audio2kp_params'])
+
+    if torch.cuda.is_available():
+        sidetuning.to(device_ids[0])
+
+    emotionprompt = EmotionDeepPrompt()
+
+    if torch.cuda.is_available():
+        emotionprompt.to(device_ids[0])
+
+    return generator, kp_detector, audio2kptransformer, sidetuning, emotionprompt
+
+
+def prepare_test_data(img_path, audio_path, opt, emotype, use_otherimg=True):
+    # sr,_ = wavfile.read(audio_path)
+
+    if use_otherimg:
+        source_latent = np.load(img_path.replace('cropped', 'latent')[:-4]+'.npy', allow_pickle=True)
+    else:
+        source_latent = np.load(img_path.replace('images', 'latent')[:-9]+'.npy', allow_pickle=True)
+    he_source = {}
+    for k in source_latent[1].keys():
+        he_source[k] = torch.from_numpy(source_latent[1][k][0]).unsqueeze(0).cuda()
+
+    # source images
+    source_img = img_as_float32(io.imread(img_path)).transpose((2, 0, 1))
+    asp = os.path.basename(audio_path)[:-4]
+
+    # latent code
+    y_trg = emo_label.index(emotype)
+    z_trg = torch.randn(latent_dim)
+
+    # driving latent
+    latent_path_driving = f'{root_wav}/latent_evp_25/{asp}.npy'
+    pose_gz = gzip.GzipFile(f'{root_wav}/poseimg/{asp}.npy.gz', 'r')
+    poseimg = np.load(pose_gz)
+    deepfeature = np.load(f'{root_wav}/deepfeature32/{asp}.npy')
+    driving_latent = np.load(latent_path_driving[:-4]+'.npy', allow_pickle=True)
+    he_driving = driving_latent[1]
+
+    # gt frame number
+    frames = glob.glob(f'{root_wav}/images_evp_25/cropped/*.jpg')
+    num_frames = len(frames)
+
+    wave_tensor = _load_tensor(audio_path)
+    if len(wave_tensor.shape) > 1:
+        wave_tensor = wave_tensor[:, 0]
+    mel_tensor = to_melspec(wave_tensor)
+    mel_tensor = (torch.log(1e-5 + mel_tensor) - mean) / std
+    name_len = min(mel_tensor.shape[1], poseimg.shape[0], deepfeature.shape[0])
+
+    audio_frames = []
+    poseimgs = []
+    deep_feature = []
+
+    pad, deep_pad = np.load('/content/EAT_code/pad.npy', allow_pickle=True)
+
+    if name_len < num_frames:
+        diff = num_frames - name_len
+        if diff > 2:
+            print(f"Attention: the frames are {diff} more than name_len, we will use name_len to replace num_frames")
+            num_frames=name_len
+            for k in he_driving.keys():
+                he_driving[k] = he_driving[k][:name_len, :]
+    for rid in range(0, num_frames):
+        audio = []
+        poses = []
+        deeps = []
+        for i in range(rid - opt['num_w'], rid + opt['num_w'] + 1):
+            if i < 0:
+                audio.append(pad)
+                poses.append(poseimg[0])
+                deeps.append(deep_pad)
+            elif i >= name_len:
+                audio.append(pad)
+                poses.append(poseimg[-1])
+                deeps.append(deep_pad)
+            else:
+                audio.append(mel_tensor[:, i])
+                poses.append(poseimg[i])
+                deeps.append(deepfeature[i])
+
+        audio_frames.append(torch.stack(audio, dim=1))
+        poseimgs.append(poses)
+        deep_feature.append(deeps)
+    audio_frames = torch.stack(audio_frames, dim=0)
+    poseimgs = torch.from_numpy(np.array(poseimgs))
+    deep_feature = torch.from_numpy(np.array(deep_feature)).to(torch.float)
+    return audio_frames, poseimgs, deep_feature, source_img, he_source, he_driving, num_frames, y_trg, z_trg, latent_path_driving
+
+def load_ckpt(ckpt, kp_detector, generator, audio2kptransformer, sidetuning, emotionprompt):
+    checkpoint = torch.load(ckpt, map_location=torch.device('cpu'))
+    if audio2kptransformer is not None:
+        audio2kptransformer.load_state_dict(checkpoint['audio2kptransformer'])
+    if generator is not None:
+        generator.load_state_dict(checkpoint['generator'])
+    if kp_detector is not None:
+        kp_detector.load_state_dict(checkpoint['kp_detector'])
+    if sidetuning is not None:
+        sidetuning.load_state_dict(checkpoint['sidetuning'])
+    if emotionprompt is not None:
+        emotionprompt.load_state_dict(checkpoint['emotionprompt'])
+
+import cv2
+import dlib
+from tqdm import tqdm
+from skimage import transform as tf
+detector = dlib.get_frontal_face_detector()
+predictor = dlib.shape_predictor('/content/EAT_code/demo/shape_predictor_68_face_landmarks.dat')
+
+def shape_to_np(shape, dtype="int"):
+    # initialize the list of (x, y)-coordinates
+    coords = np.zeros((shape.num_parts, 2), dtype=dtype)
+
+    # loop over all facial landmarks and convert them
+    # to a 2-tuple of (x, y)-coordinates
+    for i in range(0, shape.num_parts):
+        coords[i] = (shape.part(i).x, shape.part(i).y)
+
+    # return the list of (x, y)-coordinates
+    return coords
+
+def crop_image(image_path, out_path):
+    template = np.load('/content/EAT_code/demo/bo_1resized_template.npy')
+    image = cv2.imread(image_path)
+    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    rects = detector(gray, 1)  #detect human face
+    if len(rects) != 1:
+        return 0
+    for (j, rect) in enumerate(rects):
+        shape = predictor(gray, rect) #detect 68 points
+        shape = shape_to_np(shape)
+
+    pts2 = np.float32(template[:47,:])
+    pts1 = np.float32(shape[:47,:]) #eye and nose
+    tform = tf.SimilarityTransform()
+    tform.estimate( pts2, pts1) #Set the transformation matrix with the explicit parameters.
+
+    dst = tf.warp(image, tform, output_shape=(256, 256))
+
+    dst = np.array(dst * 255, dtype=np.uint8)
+
+    cv2.imwrite(out_path, dst)
+
+def preprocess_imgs(allimgs, tmp_allimgs_cropped):
+    name_cropped = []
+    for path in tmp_allimgs_cropped:
+        name_cropped.append(os.path.basename(path))
+    for path in allimgs:
+        if os.path.basename(path) in name_cropped:
+            continue
+        else:
+            out_path = path.replace('imgs1/', 'imgs_cropped1/')
+            crop_image(path, out_path)
+
+from sync_batchnorm import DataParallelWithCallback
+def load_checkpoints_extractor(config_path, checkpoint_path, cpu=False):
+
+    with open(config_path) as f:
+        config = yaml.load(f, Loader=yaml.FullLoader)
+
+    kp_detector = KPDetector(**config['model_params']['kp_detector_params'],
+                             **config['model_params']['common_params'])
+    if not cpu:
+        kp_detector.cuda()
+
+    he_estimator = HEEstimator(**config['model_params']['he_estimator_params'],
+                               **config['model_params']['common_params'])
+    if not cpu:
+        he_estimator.cuda()
+
+    if cpu:
+        checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu'))
+    else:
+        checkpoint = torch.load(checkpoint_path)
+
+    kp_detector.load_state_dict(checkpoint['kp_detector'])
+    he_estimator.load_state_dict(checkpoint['he_estimator'])
+
+    if not cpu:
+        kp_detector = DataParallelWithCallback(kp_detector)
+        he_estimator = DataParallelWithCallback(he_estimator)
+
+    kp_detector.eval()
+    he_estimator.eval()
+
+    return kp_detector, he_estimator
+
+def estimate_latent(driving_video, kp_detector, he_estimator):
+    with torch.no_grad():
+        predictions = []
+        driving = torch.tensor(np.array(driving_video)[np.newaxis].astype(np.float32)).permute(0, 4, 1, 2, 3).cuda()
+        kp_canonical = kp_detector(driving[:, :, 0])
+        he_drivings = {'yaw': [], 'pitch': [], 'roll': [], 't': [], 'exp': []}
+
+        for frame_idx in range(driving.shape[2]):
+            driving_frame = driving[:, :, frame_idx]
+            he_driving = he_estimator(driving_frame)
+            for k in he_drivings.keys():
+                he_drivings[k].append(he_driving[k])
+    return [kp_canonical, he_drivings]
+
+def extract_keypoints(extract_list):
+    kp_detector, he_estimator = load_checkpoints_extractor(config_path='/content/EAT_code/config/vox-256-spade.yaml', checkpoint_path='/content/EAT_code/ckpt/pretrain_new_274.pth.tar')
+    if not os.path.exists('./demo/imgs_latent/'):
+        os.makedirs('./demo/imgs_latent/')
+    for imgname in tqdm(extract_list):
+        path_frames = [imgname]
+        filesname=os.path.basename(imgname)[:-4]
+        if os.path.exists(f'./demo/imgs_latent/'+filesname+'.npy'):
+            continue
+        driving_frames = []
+        for im in path_frames:
+            driving_frames.append(imageio.imread(im))
+        driving_video = [resize(frame, (256, 256))[..., :3] for frame in driving_frames]
+
+        kc, he = estimate_latent(driving_video, kp_detector, he_estimator)
+        kc = kc['value'].cpu().numpy()
+        for k in he:
+            he[k] = torch.cat(he[k]).cpu().numpy()
+        np.save('./demo/imgs_latent/'+filesname, [kc, he])
+
+def preprocess_cropped_imgs(allimgs_cropped):
+    extract_list = []
+    for img_path in allimgs_cropped:
+        if not os.path.exists(img_path.replace('cropped', 'latent')[:-4]+'.npy'):
+            extract_list.append(img_path)
+    if len(extract_list) > 0:
+        print('=========', "Extract latent keypoints from New image", '======')
+        extract_keypoints(extract_list)
+
+def test(ckpt, emotype, save_dir=" "):
+    # with open("config/vox-transformer2.yaml") as f:
+    with open("/content/EAT_code/config/deepprompt_eam3d_st_tanh_304_3090_all.yaml") as f:
+        config = yaml.load(f, Loader=yaml.FullLoader)
+    cur_path = os.getcwd()
+    generator, kp_detector, audio2kptransformer, sidetuning, emotionprompt = build_model(config)
+    load_ckpt(ckpt, kp_detector=kp_detector, generator=generator, audio2kptransformer=audio2kptransformer, sidetuning=sidetuning, emotionprompt=emotionprompt)
+
+    audio2kptransformer.eval()
+    generator.eval()
+    kp_detector.eval()
+    sidetuning.eval()
+    emotionprompt.eval()
+
+    all_wavs2 = [f'{root_wav}/{os.path.basename(root_wav)}.wav']
+    allimg = glob.glob('/content/EAT_code/demo/imgs1/*.jpg')
+    tmp_allimg_cropped = glob.glob('/content/EAT_code/demo/imgs_cropped1/*.jpg')
+    preprocess_imgs(allimg, tmp_allimg_cropped) # crop and align images
+
+    allimg_cropped = glob.glob('/content/EAT_code/demo/imgs_cropped1/*.jpg')
+    preprocess_cropped_imgs(allimg_cropped) # extract latent keypoints if necessary
+
+    for ind in tqdm(range(len(all_wavs2))):
+        for img_path in tqdm(allimg_cropped):
+            audio_path = all_wavs2[ind]
+            # read in data
+            audio_frames, poseimgs, deep_feature, source_img, he_source, he_driving, num_frames, y_trg, z_trg, latent_path_driving = prepare_test_data(img_path, audio_path, config['model_params']['audio2kp_params'], emotype)
+
+
+            with torch.no_grad():
+                source_img = torch.from_numpy(source_img).unsqueeze(0).cuda()
+                kp_canonical = kp_detector(source_img, with_feature=True)     # {'value': value, 'jacobian': jacobian}
+                kp_cano = kp_canonical['value']
+
+                x = {}
+                x['mel'] = audio_frames.unsqueeze(1).unsqueeze(0).cuda()
+                x['z_trg'] = z_trg.unsqueeze(0).cuda()
+                x['y_trg'] = torch.tensor(y_trg, dtype=torch.long).cuda().reshape(1)
+                x['pose'] = poseimgs.cuda()
+                x['deep'] = deep_feature.cuda().unsqueeze(0)
+                x['he_driving'] = {'yaw': torch.from_numpy(he_driving['yaw']).cuda().unsqueeze(0),
+                                'pitch': torch.from_numpy(he_driving['pitch']).cuda().unsqueeze(0),
+                                'roll': torch.from_numpy(he_driving['roll']).cuda().unsqueeze(0),
+                                't': torch.from_numpy(he_driving['t']).cuda().unsqueeze(0),
+                                }
+
+                ### emotion prompt
+                emoprompt, deepprompt = emotionprompt(x)
+                a2kp_exps = []
+                emo_exps = []
+                T = 5
+                if T == 1:
+                    for i in range(x['mel'].shape[1]):
+                        xi = {}
+                        xi['mel'] = x['mel'][:,i,:,:,:].unsqueeze(1)
+                        xi['z_trg'] = x['z_trg']
+                        xi['y_trg'] = x['y_trg']
+                        xi['pose'] = x['pose'][i,:,:,:,:].unsqueeze(0)
+                        xi['deep'] = x['deep'][:,i,:,:,:].unsqueeze(1)
+                        xi['he_driving'] = {'yaw': x['he_driving']['yaw'][:,i,:].unsqueeze(0),
+                                    'pitch': x['he_driving']['pitch'][:,i,:].unsqueeze(0),
+                                    'roll': x['he_driving']['roll'][:,i,:].unsqueeze(0),
+                                    't': x['he_driving']['t'][:,i,:].unsqueeze(0),
+                                    }
+                        he_driving_emo_xi, input_st_xi = audio2kptransformer(xi, kp_canonical, emoprompt=emoprompt, deepprompt=deepprompt, side=True)           # {'yaw': yaw, 'pitch': pitch, 'roll': roll, 't': t, 'exp': exp}
+                        emo_exp = sidetuning(input_st_xi, emoprompt, deepprompt)
+                        a2kp_exps.append(he_driving_emo_xi['emo'])
+                        emo_exps.append(emo_exp)
+                elif T is not None:
+                    for i in range(x['mel'].shape[1]//T+1):
+                        if i*T >= x['mel'].shape[1]:
+                            break
+                        xi = {}
+                        xi['mel'] = x['mel'][:,i*T:(i+1)*T,:,:,:]
+                        xi['z_trg'] = x['z_trg']
+                        xi['y_trg'] = x['y_trg']
+                        xi['pose'] = x['pose'][i*T:(i+1)*T,:,:,:,:]
+                        xi['deep'] = x['deep'][:,i*T:(i+1)*T,:,:,:]
+                        xi['he_driving'] = {'yaw': x['he_driving']['yaw'][:,i*T:(i+1)*T,:],
+                                    'pitch': x['he_driving']['pitch'][:,i*T:(i+1)*T,:],
+                                    'roll': x['he_driving']['roll'][:,i*T:(i+1)*T,:],
+                                    't': x['he_driving']['t'][:,i*T:(i+1)*T,:],
+                                    }
+                        he_driving_emo_xi, input_st_xi = audio2kptransformer(xi, kp_canonical, emoprompt=emoprompt, deepprompt=deepprompt, side=True)           # {'yaw': yaw, 'pitch': pitch, 'roll': roll, 't': t, 'exp': exp}
+                        emo_exp = sidetuning(input_st_xi, emoprompt, deepprompt)
+                        a2kp_exps.append(he_driving_emo_xi['emo'])
+                        emo_exps.append(emo_exp)
+
+                if T is None:
+                    he_driving_emo, input_st = audio2kptransformer(x, kp_canonical, emoprompt=emoprompt, deepprompt=deepprompt, side=True)           # {'yaw': yaw, 'pitch': pitch, 'roll': roll, 't': t, 'exp': exp}
+                    emo_exps = sidetuning(input_st, emoprompt, deepprompt).reshape(-1, 45)
+                else:
+                    he_driving_emo = {}
+                    he_driving_emo['emo'] = torch.cat(a2kp_exps, dim=0)
+                    emo_exps = torch.cat(emo_exps, dim=0).reshape(-1, 45)
+
+                exp = he_driving_emo['emo']
+                device = exp.get_device()
+                exp = torch.mm(exp, expU.t().to(device))
+                exp = exp + expmean.expand_as(exp).to(device)
+                exp = exp + emo_exps
+
+
+                source_area = ConvexHull(kp_cano[0].cpu().numpy()).volume
+                exp = exp * source_area
+
+                he_new_driving = {'yaw': torch.from_numpy(he_driving['yaw']).cuda(),
+                                'pitch': torch.from_numpy(he_driving['pitch']).cuda(),
+                                'roll': torch.from_numpy(he_driving['roll']).cuda(),
+                                't': torch.from_numpy(he_driving['t']).cuda(),
+                                'exp': exp}
+                he_driving['exp'] = torch.from_numpy(he_driving['exp']).cuda()
+
+                kp_source = keypoint_transformation(kp_canonical, he_source, False)
+                mean_source = torch.mean(kp_source['value'], dim=1)[0]
+                kp_driving = keypoint_transformation(kp_canonical, he_new_driving, False)
+                mean_driving = torch.mean(torch.mean(kp_driving['value'], dim=1), dim=0)
+                kp_driving['value'] = kp_driving['value']+(mean_source-mean_driving).unsqueeze(0).unsqueeze(0)
+                bs = kp_source['value'].shape[0]
+                predictions_gen = []
+                for i in tqdm(range(num_frames)):
+                    kp_si = {}
+                    kp_si['value'] = kp_source['value'][0].unsqueeze(0)
+                    kp_di = {}
+                    kp_di['value'] = kp_driving['value'][i].unsqueeze(0)
+                    generated = generator(source_img, kp_source=kp_si, kp_driving=kp_di, prompt=emoprompt)
+                    predictions_gen.append(
+                        (np.transpose(generated['prediction'].data.cpu().numpy(), [0, 2, 3, 1])[0] * 255).astype(np.uint8))
+
+            log_dir = save_dir
+            os.makedirs(os.path.join(log_dir, "temp"), exist_ok=True)
+
+            f_name = os.path.basename(img_path[:-4]) + "_" + emotype + "_" + os.path.basename(latent_path_driving)[:-4] + ".mp4"
+            video_path = os.path.join(log_dir, "temp", f_name)
+            imageio.mimsave(video_path, predictions_gen, fps=25.0)
+
+            save_video = os.path.join(log_dir, f_name)
+            cmd = r'ffmpeg -loglevel error -y -i "%s" -i "%s" -vcodec copy -shortest "%s"' % (video_path, audio_path, save_video)
+            os.system(cmd)
+            os.remove(video_path)
+
+if __name__ == '__main__':
+    argparser = argparse.ArgumentParser()
+    argparser.add_argument("--save_dir", type=str, default="/content/EAT_code/Results ", help="path of the output video")
+    argparser.add_argument("--name", type=str, default="deepprompt_eam3d_all_final_313", help="path of the output video")
+    argparser.add_argument("--emo", type=str, default="hap", help="emotion type ('ang',  'con',  'dis',  'fea',  'hap',  'neu',  'sad',  'sur')")
+    argparser.add_argument("--root_wav", type=str, default='./demo/video_processed/M003_neu_1_001', help="emotion type ('ang',  'con',  'dis',  'fea',  'hap',  'neu',  'sad',  'sur')")
+    args = argparser.parse_args()
+
+    root_wav=args.root_wav
+
+    if len(args.name) > 1:
+        name = args.name
+        print(name)
+    test(f'/content/EAT_code/ckpt/deepprompt_eam3d_all_final_313.pth.tar', args.emo, save_dir=f'./demo/output/{name}/')
+