finish image and video

app.py

@@ -132,7 +132,7 @@ def drive_im(source_image, driving_image, adapt_scale):
     return img_as_ubyte(prediction[0])
 
 
-def drive_vi(source_image, driving_video, mode, find_best_frame, relative, adapt_scale):
+def drive_vi(source_image, driving_video, mode, find_best_frame, best_frame, relative, adapt_scale):
     reader = imageio.get_reader(driving_video)
     fps = reader.get_meta_data()['fps']
     driving_video = []
@@ -150,8 +150,9 @@ def drive_vi(source_image, driving_video, mode, find_best_frame, relative, adapt
     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)
+    i = 0
+    if find_best_frame != "specific ref frame" or best_frame > 0:
+        i = best_frame if find_best_frame == "specific ref frame" else 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]
@@ -165,7 +166,7 @@ def drive_vi(source_image, driving_video, mode, find_best_frame, relative, adapt
                                      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
+    return result_video_path, i
 
 
 config = "config/vox-256.yaml"
@@ -174,7 +175,7 @@ cpu = True # 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."
+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 \"relative motion\" paramter to use relative keypoint coordinates for preserving global object geometry, select \"adaptive scale\" parameter for better optic flow estimation using adaptive movement scale based on convex hull of keypoints, and select \"find best ref frame\" parameter to generate video from the frame that is the most alligned with source image."
 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)
@@ -202,14 +203,28 @@ with gr.Blocks(title="Demostration of FNeVR") as demo:
                 gr.Markdown("#### Parameter")
                 inp3 = gr.Checkbox(value=True, label="adaptive scale")
 
-                btn = gr.Button(value="Animate")
+                btn1 = gr.Button(value="Animate")
             with gr.Column():
                 gr.Markdown("#### Output")
                 outp = gr.Image(label="Generated face")
 
+                with gr.Row():
+                    with gr.Column():
+                        btn2 = gr.Button(value="Reset")
+                    with gr.Column():
+                        btn3 = gr.Button(value="Cancel")
+
         gr.Examples([["sup-mat/driving.png", "sup-mat/source.png"]], [inp2, inp1])
 
-        btn.click(fn=drive_im, inputs=[inp1, inp2, inp3], outputs=outp)
+        def reset_output():
+            return outp.update(value=None)
+        
+        def reset_all():
+            return inp1.update(value=None), inp2.update(value=None), inp3.update(value=True), outp.update(value=None)
+
+        run = btn1.click(fn=drive_im, inputs=[inp1, inp2, inp3], outputs=outp)
+        btn2.click(fn=reset_all, outputs=[inp1, inp2, inp3, outp])
+        btn3.click(fn=reset_output, outputs=[outp], cancels=[run])
     with gr.Tab("Driving by video"):
         gr.Markdown(vi_description)
 
@@ -220,24 +235,47 @@ with gr.Blocks(title="Demostration of FNeVR") as demo:
                 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")
+                inp3 = gr.Radio(choices=["reenactment", "reconstruction"], value="reenactment", label="mode (if \"reconstruction\" selected, origin face is the first frame of driving video)")
+                inp6 = gr.Checkbox(value=True, label="relative motion")
+                inp7 = gr.Checkbox(value=True, label="adaptive scale")
+                inp4 = gr.Radio(choices=["find best ref frame (more time consumed)", "specific ref frame"], value="find best ref frame (more time consumed)", label="set ref frame (used by relative motion and adaptive scale)")
+                inp5 = gr.Number(label="specific ref frame (default: 0)", value=0, precision=0, visible=False)
+
+                def reset_ref(inp4):
+                    return inp5.update(visible=True) if inp4 == "specific ref frame" else inp5.update(value=0, visible=False)
+
+                inp4.change(fn=reset_ref, inputs=inp4, outputs=inp5)
 
-                btn = gr.Button(value="Animate")
+                btn1 = gr.Button(value="Animate")
             with gr.Column():
                 gr.Markdown("#### Output")
-                outp = gr.Video(label="Generated video")
+                outp1 = gr.Video(label="Generated video")
+                outp2 = gr.Number(label="Ref frame", value=0, precision=0)
+
+                # file = gr.File(value="result_video.mp4", visible=False)
+
+                with gr.Row():
+                    with gr.Column():
+                        btn2 = gr.Button(value="Reset")
+                    with gr.Column():
+                        btn3 = gr.Button(value="Cancel")
+
+        gr.Examples([["sup-mat/driving.mp4", "sup-mat/source_for_video.png", "specific ref frame", 53]], [inp2, inp1, inp4, inp5])
 
-        gr.Examples([["sup-mat/driving.mp4", "sup-mat/source_for_video.png"]], [inp2, inp1])
+        def reset_output():
+            return outp1.update(value=None), outp2.update(value=0)
 
-        btn.click(fn=drive_vi, inputs=[inp1, inp2, inp3, inp4, inp5, inp6], outputs=outp)
+        def reset_all():
+            return inp1.update(value=None), inp2.update(value=None), inp3.update(value="reenactment"), inp4.update(value="find best ref frame (more time consumed)"), inp5.update(value=0), inp6.update(value=True), inp7.update(value=True), outp1.update(value=None), outp2.update(value=0)
 
+        run = btn1.click(fn=drive_vi, inputs=[inp1, inp2, inp3, inp4, inp5, inp6, inp7], outputs=[outp1, outp2])
+        btn2.click(fn=reset_all, outputs=[inp1, inp2, inp3, inp4, inp5, inp6, inp7, outp1, outp2])
+        btn3.click(fn=reset_output, outputs=[outp1, outp2], cancels=[run])
     with gr.Tab("Real time animation"):
         gr.Markdown("#### Real time animation")
 
     gr.Markdown("## Acknowledgements")
     gr.Markdown(acknowledgements)
 
+demo.queue()
 demo.launch()