Update app.py

app.py

@@ -1,248 +1,248 @@
-import os
-import gradio as gr
-import yaml
-from argparse import ArgumentParser
-from tqdm import tqdm
-
-import numpy as np
-import imageio
-from skimage.transform import resize
-from skimage import img_as_ubyte
-from scipy.spatial import ConvexHull
-import torch
-from sync_batchnorm import DataParallelWithCallback
-import face_alignment
-
-from modules.generator import OcclusionAwareGenerator_SPADE
-from modules.keypoint_detector import KPDetector
-
-
-def normalize_kp(kp_source, kp_driving, kp_driving_initial, adapt_movement_scale=False,
-                 use_relative_movement=False, use_relative_jacobian=False):
-    kp_new = {k: v for k, v in kp_driving.items()}
-
-    if adapt_movement_scale:
-        source_area = ConvexHull(kp_source['value'][0].data.cpu().numpy()).volume
-        driving_area = ConvexHull(kp_driving_initial['value'][0].data.cpu().numpy()).volume
-        adapt_movement_scale = np.sqrt(source_area) / np.sqrt(driving_area)
-        kp_new['value'] = kp_driving['value'] * adapt_movement_scale    # for reenactment demo
-    else:
-        adapt_movement_scale = 1
-
-    if use_relative_movement:
-        kp_value_diff = (kp_driving['value'] - kp_driving_initial['value'])
-        kp_value_diff *= adapt_movement_scale
-        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_checkpoints(config_path, checkpoint_path, cpu=False):
-    with open(config_path) as f:
-        # config = yaml.load(f)
-        config = yaml.load(f, Loader=yaml.FullLoader)
-
-    generator = OcclusionAwareGenerator_SPADE(**config['model_params']['generator_params'],
-                                              **config['model_params']['common_params'])
-    if not cpu:
-        generator.cuda()
-
-    kp_detector = KPDetector(**config['model_params']['kp_detector_params'],
-                             **config['model_params']['common_params'])
-    if not cpu:
-        kp_detector.cuda()
-
-    if cpu:
-        checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu'))
-    else:
-        checkpoint = torch.load(checkpoint_path)
-
-    generator.load_state_dict(checkpoint['generator'])
-    kp_detector.load_state_dict(checkpoint['kp_detector'])
-
-    if not cpu:
-        generator = DataParallelWithCallback(generator)
-        kp_detector = DataParallelWithCallback(kp_detector)
-
-    generator.eval()
-    kp_detector.eval()
-
-    return generator, kp_detector
-
-
-def make_animation(source_image, driving_video, generator, kp_detector, relative=True, adapt_movement_scale=True,
-                   cpu=False):
-    with torch.no_grad():
-        predictions = []
-        source = torch.tensor(source_image[np.newaxis].astype(np.float32)).permute(0, 3, 1, 2)
-        if not cpu:
-            source = source.cuda()
-        driving = torch.tensor(np.array(driving_video)[np.newaxis].astype(np.float32)).permute(0, 4, 1, 2, 3)
-        kp_source = kp_detector(source)
-        kp_driving_initial = kp_detector(driving[:, :, 0])
-
-        for frame_idx in tqdm(range(driving.shape[2])):
-            driving_frame = driving[:, :, frame_idx]
-            if not cpu:
-                driving_frame = driving_frame.cuda()
-            kp_driving = kp_detector(driving_frame)
-            kp_norm = normalize_kp(kp_source=kp_source, kp_driving=kp_driving,
-                                   kp_driving_initial=kp_driving_initial, use_relative_movement=relative,
-                                   use_relative_jacobian=relative, adapt_movement_scale=adapt_movement_scale)
-            out = generator(source, kp_source=kp_source, kp_driving=kp_norm)
-
-            predictions.append(np.transpose(out['prediction'].data.cpu().numpy(), [0, 2, 3, 1])[0])
-    return predictions
-
-
-def find_best_frame_func(source, driving, cpu=False):
-    def normalize_kp_infunc(kp):
-        kp = kp - kp.mean(axis=0, keepdims=True)
-        area = ConvexHull(kp[:, :2]).volume
-        area = np.sqrt(area)
-        kp[:, :2] = kp[:, :2] / area
-        return kp
-
-    fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, flip_input=True,
-                                      device='cpu' if cpu else 'cuda')
-    kp_source = fa.get_landmarks(255 * source)[0]
-    kp_source = normalize_kp_infunc(kp_source)
-    norm = float('inf')
-    frame_num = 0
-    for i, image in tqdm(enumerate(driving)):
-        kp_driving = fa.get_landmarks(255 * image)[0]
-        kp_driving = normalize_kp_infunc(kp_driving)
-        new_norm = (np.abs(kp_source - kp_driving) ** 2).sum()
-        if new_norm < norm:
-            norm = new_norm
-            frame_num = i
-    return frame_num
-
-
-def drive_im(source_image, driving_image, adapt_scale):
-    source_image = resize(source_image, (256, 256))[..., :3]
-    driving_image = [resize(driving_image, (256, 256))[..., :3]]
-
-    prediction = make_animation(source_image, driving_image, generator, kp_detector, relative=False,
-                                adapt_movement_scale=adapt_scale, cpu=cpu)
-    return img_as_ubyte(prediction[0])
-
-
-def drive_vi(source_image, driving_video, mode, find_best_frame, relative, adapt_scale):
-    reader = imageio.get_reader(driving_video)
-    fps = reader.get_meta_data()['fps']
-    driving_video = []
-    try:
-        for im in reader:
-            driving_video.append(im)
-    except RuntimeError:
-        pass
-    reader.close()
-
-    
-    if mode == 'reconstruction':
-        source_image = driving_video[0]
-
-    source_image = resize(source_image, (256, 256))[..., :3]
-    driving_video = [resize(frame, (256, 256))[..., :3] for frame in driving_video]
-
-    if find_best_frame:
-        i = find_best_frame_func(source_image, driving_video, cpu=cpu)
-        print("Best frame: " + str(i))
-        driving_forward = driving_video[i:]
-        driving_backward = driving_video[:(i + 1)][::-1]
-        predictions_forward = make_animation(source_image, driving_forward, generator, kp_detector,
-                                             relative=relative, adapt_movement_scale=adapt_scale, cpu=cpu)
-        predictions_backward = make_animation(source_image, driving_backward, generator, kp_detector,
-                                              relative=relative, adapt_movement_scale=adapt_scale, cpu=cpu)
-        predictions = predictions_backward[::-1] + predictions_forward[1:]
-    else:
-        predictions = make_animation(source_image, driving_video, generator, kp_detector, relative=relative,
-                                     adapt_movement_scale=adapt_scale, cpu=cpu)
-    result_video_path = "result_video.mp4"
-    imageio.mimsave(result_video_path, [img_as_ubyte(frame) for frame in predictions], fps=fps)
-    return result_video_path
-
-
-os.environ['CUDA_VISIBLE_DEVICES'] = "3"
-config = "config/vox-256.yaml"
-checkpoint = "00000099-checkpoint.pth.tar"
-mode = "reenactment"
-relative = True
-adapt_scale = True
-find_best_frame = True
-cpu = False # decided by the deploying environment
-
-description = "We propose a Face Neural Volume Rendering (FNeVR) network for more realistic face animation, by taking the merits of 2D motion warping on facial expression transformation and 3D volume rendering on high-quality image synthesis in a unified framework.<br>[Paper](https://arxiv.org/abs/2209.10340) and [Code](https://github.com/zengbohan0217/FNeVR)"
-im_description = "We can animate a face portrait by a single image in this tab.<br>Please input the origin face and the driving face which provides pose and expression information, then we can obtain the virtual generated face.<br>We can select \"adaptive scale\" parameter for better optic flow estimation using adaptive movement scale based on convex hull of keypoints."
-vi_description = "We can animate a face portrait by a video in this tab.<br>Please input the origin face and the driving video which provides pose and expression information, then we can obtain the virtual generated video.<br>Please select inference mode (reenactment for different identities and reconstruction for the same identities).<br>We can select \"find best frame\" parameter to generate video from the frame that is the most alligned with source image, select \"relative motion\" paramter to use relative keypoint coordinates for preserving global object geometry, and select \"adaptive scale\" parameter for better optic flow estimation using adaptive movement scale based on convex hull of keypoints."
-acknowledgements = "This work was supported by “the Fundamental Research Funds for the Central Universities”, and the National Natural Science Foundation of China under Grant 62076016, Beijing Natural Science Foundation-Xiaomi Innovation Joint Fund L223024. Besides, we gratefully acknowledge the support of [MindSpore](https://www.mindspore.cn), CANN (Compute Architecture for Neural Networks) and Ascend AI processor used for this research.<br>Our FNeVR implementation is inspired by [FOMM](https://github.com/AliaksandrSiarohin/first-order-model) and [DECA](https://github.com/YadiraF/DECA). We appreciate the authors of these papers for making their codes available to the public."
-
-generator, kp_detector = load_checkpoints(config_path=config, checkpoint_path=checkpoint, cpu=cpu)
-
-# iface = gr.Interface(fn=drive_im, 
-#                     inputs=[gr.Image(label="Origin face"), 
-#                             gr.Image(label="Driving face"), 
-#                             gr.CheckboxGroup(label="adapt scale")], 
-#                     outputs=gr.Image(label="Generated face"), examples=[["sup-mat/source.png"], ["sup-mat/driving.png"]], 
-#                     title="Demostration of FNeVR", description=description)
-
-with gr.Blocks(title="Demostration of FNeVR") as demo:
-    gr.Markdown("# <center> Demostration of FNeVR")
-    gr.Markdown(description)
-
-    with gr.Tab("Driving by image"):
-        gr.Markdown(im_description)
-
-        with gr.Row():
-            with gr.Column():
-                gr.Markdown("#### Inputs")
-                inp2 = gr.Image(label="Driving face")
-                inp1 = gr.Image(label="Origin face")
-                
-                gr.Markdown("#### Parameter")
-                inp3 = gr.Checkbox(value=True, label="adaptive scale")
-
-                btn = gr.Button(value="Animate")
-            with gr.Column():
-                gr.Markdown("#### Output")
-                outp = gr.Image(label="Generated face")
-
-        gr.Examples([["sup-mat/driving.png", "sup-mat/source.png"]], [inp2, inp1])
-
-        btn.click(fn=drive_im, inputs=[inp1, inp2, inp3], outputs=outp)
-    with gr.Tab("Driving by video"):
-        gr.Markdown(vi_description)
-
-        with gr.Row():
-            with gr.Column():
-                gr.Markdown("#### Inputs")
-                inp2 = gr.Video(label="Driving video")
-                inp1 = gr.Image(label="Origin face")
-                
-                gr.Markdown("#### Parameters")
-                inp3 = gr.Radio(choices=['reenactment', 'reconstruction'], value="reenactment", label="mode (if \"reconstruction\" selected, origin face is the first frame of driving video)")
-                inp4 = gr.Checkbox(value=True, label="find best frame (more time consumed)")
-                inp5 = gr.Checkbox(value=True, label="relative motion")
-                inp6 = gr.Checkbox(value=True, label="adaptive scale")
-
-                btn = gr.Button(value="Animate")
-            with gr.Column():
-                gr.Markdown("#### Output")
-                outp = gr.Video(label="Generated video")
-
-        gr.Examples([["sup-mat/driving.mp4", "sup-mat/source_for_video.png"]], [inp2, inp1])
-
-        btn.click(fn=drive_vi, inputs=[inp1, inp2, inp3, inp4, inp5, inp6], outputs=outp)
-
-    with gr.Tab("Real time animation"):
-        gr.Markdown("#### Real time animation")
-
-    gr.Markdown("## Acknowledgements")
-    gr.Markdown(acknowledgements)
-
-demo.launch()
+import os
+import gradio as gr
+import yaml
+from argparse import ArgumentParser
+from tqdm import tqdm
+
+import numpy as np
+import imageio
+from skimage.transform import resize
+from skimage import img_as_ubyte
+from scipy.spatial import ConvexHull
+import torch
+from sync_batchnorm import DataParallelWithCallback
+import face_alignment
+
+from modules.generator import OcclusionAwareGenerator_SPADE
+from modules.keypoint_detector import KPDetector
+
+
+def normalize_kp(kp_source, kp_driving, kp_driving_initial, adapt_movement_scale=False,
+                 use_relative_movement=False, use_relative_jacobian=False):
+    kp_new = {k: v for k, v in kp_driving.items()}
+
+    if adapt_movement_scale:
+        source_area = ConvexHull(kp_source['value'][0].data.cpu().numpy()).volume
+        driving_area = ConvexHull(kp_driving_initial['value'][0].data.cpu().numpy()).volume
+        adapt_movement_scale = np.sqrt(source_area) / np.sqrt(driving_area)
+        kp_new['value'] = kp_driving['value'] * adapt_movement_scale    # for reenactment demo
+    else:
+        adapt_movement_scale = 1
+
+    if use_relative_movement:
+        kp_value_diff = (kp_driving['value'] - kp_driving_initial['value'])
+        kp_value_diff *= adapt_movement_scale
+        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_checkpoints(config_path, checkpoint_path, cpu=False):
+    with open(config_path) as f:
+        # config = yaml.load(f)
+        config = yaml.load(f, Loader=yaml.FullLoader)
+
+    generator = OcclusionAwareGenerator_SPADE(**config['model_params']['generator_params'],
+                                              **config['model_params']['common_params'])
+    if not cpu:
+        generator.cuda()
+
+    kp_detector = KPDetector(**config['model_params']['kp_detector_params'],
+                             **config['model_params']['common_params'])
+    if not cpu:
+        kp_detector.cuda()
+
+    if cpu:
+        checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu'))
+    else:
+        checkpoint = torch.load(checkpoint_path)
+
+    generator.load_state_dict(checkpoint['generator'])
+    kp_detector.load_state_dict(checkpoint['kp_detector'])
+
+    if not cpu:
+        generator = DataParallelWithCallback(generator)
+        kp_detector = DataParallelWithCallback(kp_detector)
+
+    generator.eval()
+    kp_detector.eval()
+
+    return generator, kp_detector
+
+
+def make_animation(source_image, driving_video, generator, kp_detector, relative=True, adapt_movement_scale=True,
+                   cpu=False):
+    with torch.no_grad():
+        predictions = []
+        source = torch.tensor(source_image[np.newaxis].astype(np.float32)).permute(0, 3, 1, 2)
+        if not cpu:
+            source = source.cuda()
+        driving = torch.tensor(np.array(driving_video)[np.newaxis].astype(np.float32)).permute(0, 4, 1, 2, 3)
+        kp_source = kp_detector(source)
+        kp_driving_initial = kp_detector(driving[:, :, 0])
+
+        for frame_idx in tqdm(range(driving.shape[2])):
+            driving_frame = driving[:, :, frame_idx]
+            if not cpu:
+                driving_frame = driving_frame.cuda()
+            kp_driving = kp_detector(driving_frame)
+            kp_norm = normalize_kp(kp_source=kp_source, kp_driving=kp_driving,
+                                   kp_driving_initial=kp_driving_initial, use_relative_movement=relative,
+                                   use_relative_jacobian=relative, adapt_movement_scale=adapt_movement_scale)
+            out = generator(source, kp_source=kp_source, kp_driving=kp_norm)
+
+            predictions.append(np.transpose(out['prediction'].data.cpu().numpy(), [0, 2, 3, 1])[0])
+    return predictions
+
+
+def find_best_frame_func(source, driving, cpu=False):
+    def normalize_kp_infunc(kp):
+        kp = kp - kp.mean(axis=0, keepdims=True)
+        area = ConvexHull(kp[:, :2]).volume
+        area = np.sqrt(area)
+        kp[:, :2] = kp[:, :2] / area
+        return kp
+
+    fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, flip_input=True,
+                                      device='cpu' if cpu else 'cuda')
+    kp_source = fa.get_landmarks(255 * source)[0]
+    kp_source = normalize_kp_infunc(kp_source)
+    norm = float('inf')
+    frame_num = 0
+    for i, image in tqdm(enumerate(driving)):
+        kp_driving = fa.get_landmarks(255 * image)[0]
+        kp_driving = normalize_kp_infunc(kp_driving)
+        new_norm = (np.abs(kp_source - kp_driving) ** 2).sum()
+        if new_norm < norm:
+            norm = new_norm
+            frame_num = i
+    return frame_num
+
+
+def drive_im(source_image, driving_image, adapt_scale):
+    source_image = resize(source_image, (256, 256))[..., :3]
+    driving_image = [resize(driving_image, (256, 256))[..., :3]]
+
+    prediction = make_animation(source_image, driving_image, generator, kp_detector, relative=False,
+                                adapt_movement_scale=adapt_scale, cpu=cpu)
+    return img_as_ubyte(prediction[0])
+
+
+def drive_vi(source_image, driving_video, mode, find_best_frame, relative, adapt_scale):
+    reader = imageio.get_reader(driving_video)
+    fps = reader.get_meta_data()['fps']
+    driving_video = []
+    try:
+        for im in reader:
+            driving_video.append(im)
+    except RuntimeError:
+        pass
+    reader.close()
+
+    
+    if mode == 'reconstruction':
+        source_image = driving_video[0]
+
+    source_image = resize(source_image, (256, 256))[..., :3]
+    driving_video = [resize(frame, (256, 256))[..., :3] for frame in driving_video]
+
+    if find_best_frame:
+        i = find_best_frame_func(source_image, driving_video, cpu=cpu)
+        print("Best frame: " + str(i))
+        driving_forward = driving_video[i:]
+        driving_backward = driving_video[:(i + 1)][::-1]
+        predictions_forward = make_animation(source_image, driving_forward, generator, kp_detector,
+                                             relative=relative, adapt_movement_scale=adapt_scale, cpu=cpu)
+        predictions_backward = make_animation(source_image, driving_backward, generator, kp_detector,
+                                              relative=relative, adapt_movement_scale=adapt_scale, cpu=cpu)
+        predictions = predictions_backward[::-1] + predictions_forward[1:]
+    else:
+        predictions = make_animation(source_image, driving_video, generator, kp_detector, relative=relative,
+                                     adapt_movement_scale=adapt_scale, cpu=cpu)
+    result_video_path = "result_video.mp4"
+    imageio.mimsave(result_video_path, [img_as_ubyte(frame) for frame in predictions], fps=fps)
+    return result_video_path
+
+
+os.environ['CUDA_VISIBLE_DEVICES'] = "3"
+config = "config/vox-256.yaml"
+checkpoint = "00000099-checkpoint.pth.tar"
+mode = "reenactment"
+relative = True
+adapt_scale = True
+find_best_frame = True
+cpu = torch.cuda.is_available() # decided by the deploying environment
+
+description = "We propose a Face Neural Volume Rendering (FNeVR) network for more realistic face animation, by taking the merits of 2D motion warping on facial expression transformation and 3D volume rendering on high-quality image synthesis in a unified framework.<br>[Paper](https://arxiv.org/abs/2209.10340) and [Code](https://github.com/zengbohan0217/FNeVR)"
+im_description = "We can animate a face portrait by a single image in this tab.<br>Please input the origin face and the driving face which provides pose and expression information, then we can obtain the virtual generated face.<br>We can select \"adaptive scale\" parameter for better optic flow estimation using adaptive movement scale based on convex hull of keypoints."
+vi_description = "We can animate a face portrait by a video in this tab.<br>Please input the origin face and the driving video which provides pose and expression information, then we can obtain the virtual generated video.<br>Please select inference mode (reenactment for different identities and reconstruction for the same identities).<br>We can select \"find best frame\" parameter to generate video from the frame that is the most alligned with source image, select \"relative motion\" paramter to use relative keypoint coordinates for preserving global object geometry, and select \"adaptive scale\" parameter for better optic flow estimation using adaptive movement scale based on convex hull of keypoints."
+acknowledgements = "This work was supported by “the Fundamental Research Funds for the Central Universities”, and the National Natural Science Foundation of China under Grant 62076016, Beijing Natural Science Foundation-Xiaomi Innovation Joint Fund L223024. Besides, we gratefully acknowledge the support of [MindSpore](https://www.mindspore.cn), CANN (Compute Architecture for Neural Networks) and Ascend AI processor used for this research.<br>Our FNeVR implementation is inspired by [FOMM](https://github.com/AliaksandrSiarohin/first-order-model) and [DECA](https://github.com/YadiraF/DECA). We appreciate the authors of these papers for making their codes available to the public."
+
+generator, kp_detector = load_checkpoints(config_path=config, checkpoint_path=checkpoint, cpu=cpu)
+
+# iface = gr.Interface(fn=drive_im, 
+#                     inputs=[gr.Image(label="Origin face"), 
+#                             gr.Image(label="Driving face"), 
+#                             gr.CheckboxGroup(label="adapt scale")], 
+#                     outputs=gr.Image(label="Generated face"), examples=[["sup-mat/source.png"], ["sup-mat/driving.png"]], 
+#                     title="Demostration of FNeVR", description=description)
+
+with gr.Blocks(title="Demostration of FNeVR") as demo:
+    gr.Markdown("# <center> Demostration of FNeVR")
+    gr.Markdown(description)
+
+    with gr.Tab("Driving by image"):
+        gr.Markdown(im_description)
+
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("#### Inputs")
+                inp2 = gr.Image(label="Driving face")
+                inp1 = gr.Image(label="Origin face")
+                
+                gr.Markdown("#### Parameter")
+                inp3 = gr.Checkbox(value=True, label="adaptive scale")
+
+                btn = gr.Button(value="Animate")
+            with gr.Column():
+                gr.Markdown("#### Output")
+                outp = gr.Image(label="Generated face")
+
+        gr.Examples([["sup-mat/driving.png", "sup-mat/source.png"]], [inp2, inp1])
+
+        btn.click(fn=drive_im, inputs=[inp1, inp2, inp3], outputs=outp)
+    with gr.Tab("Driving by video"):
+        gr.Markdown(vi_description)
+
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("#### Inputs")
+                inp2 = gr.Video(label="Driving video")
+                inp1 = gr.Image(label="Origin face")
+                
+                gr.Markdown("#### Parameters")
+                inp3 = gr.Radio(choices=['reenactment', 'reconstruction'], value="reenactment", label="mode (if \"reconstruction\" selected, origin face is the first frame of driving video)")
+                inp4 = gr.Checkbox(value=True, label="find best frame (more time consumed)")
+                inp5 = gr.Checkbox(value=True, label="relative motion")
+                inp6 = gr.Checkbox(value=True, label="adaptive scale")
+
+                btn = gr.Button(value="Animate")
+            with gr.Column():
+                gr.Markdown("#### Output")
+                outp = gr.Video(label="Generated video")
+
+        gr.Examples([["sup-mat/driving.mp4", "sup-mat/source_for_video.png"]], [inp2, inp1])
+
+        btn.click(fn=drive_vi, inputs=[inp1, inp2, inp3, inp4, inp5, inp6], outputs=outp)
+
+    with gr.Tab("Real time animation"):
+        gr.Markdown("#### Real time animation")
+
+    gr.Markdown("## Acknowledgements")
+    gr.Markdown(acknowledgements)
+
+demo.launch()