update: app.py, gpen for video

app.py

@@ -215,7 +215,7 @@ def swap_image(
     os.makedirs(out_path, exist_ok=True)
     Image.fromarray(result.astype(np.uint8)).save(os.path.join(out_path, name))
     save((result, M, original_target, os.path.join(out_path, "paste_back_" + name), None),
-         trick=trick, use_post=use_gpen)
+         trick=trick, use_gpen=use_gpen)
 
 
 def process_video(
@@ -230,6 +230,7 @@ def process_video(
     frames=9999999,
     use_tddfav2=False,
     landmark_smooth="kalman",
+    use_gpen=False,
 ):
     if isinstance(G, torch.nn.Module):
         G.eval()
@@ -336,7 +337,7 @@ def process_video(
                 target = output[0][0] * 0.5 + 0.5
             else:
                 target = output[0] * 0.5 + 0.5
-        targets.append(np.array(tensor2pil_transform(target)))
+        targets.append(trick.gpen(np.array(tensor2pil_transform(target)), use_gpen=use_gpen))
         Ms.append(M)
         count += 1
         if count > frames:
@@ -373,7 +374,7 @@ def swap_image_gr(img1, img2, use_post=False, use_gpen=False, ):
     return out
 
 
-def swap_video_gr(img1, target_path, frames=9999999):
+def swap_video_gr(img1, target_path, use_gpen=False, frames=9999999):
     root_dir = make_abs_path("./online_data")
     req_id = uuid.uuid1().hex
     data_dir = os.path.join(root_dir, req_id)
@@ -393,6 +394,7 @@ def swap_video_gr(img1, target_path, frames=9999999):
         align_target='ffhq',
         align_source='ffhq',
         use_tddfav2=False,
+        use_gpen=use_gpen,
     )
 
     pool_process = 170
@@ -438,6 +440,13 @@ def swap_video_gr(img1, target_path, frames=9999999):
     return video_save_path
 
 
+css = """
+#warning {background-color: #FFCCCB} 
+.feedback textarea {font-size: 24px !important}
+.run_button {background-color: orange} 
+"""
+
+
 if __name__ == "__main__":
     use_gpu = torch.cuda.is_available()
 
@@ -449,11 +458,11 @@ if __name__ == "__main__":
                 with gr.Column(scale=3):
                     image1_input = gr.Image(label='source')
                     image2_input = gr.Image(label='target')
-                    use_post = gr.Checkbox(label="Post-Process")
-                    use_gpen = gr.Checkbox(label="Super Resolution")
+                    image_use_post = gr.Checkbox(label="Post-Process")
+                    image_use_gpen = gr.Checkbox(label="Super Resolution")
                 with gr.Column(scale=2):
                     image_output = gr.Image()
-                    image_button = gr.Button("Run: Face Swapping")
+                    image_button = gr.Button("Run: Face Swapping", elem_classes="run_button")
         with gr.Tab("Video"):
             with gr.Row():
                 with gr.Column(scale=3):
@@ -461,15 +470,16 @@ if __name__ == "__main__":
                     video_input = gr.Video(label='target')
                 with gr.Column(scale=2):
                     video_output = gr.Video()
-                    video_button = gr.Button("Run: Face Swapping")
+                    video_use_gpen = gr.Checkbox(label="Super Resolution")
+                    video_button = gr.Button("Run: Face Swapping", elem_classes="run_button")
         image_button.click(
             swap_image_gr,
-            inputs=[image1_input, image2_input, use_post, use_gpen],
+            inputs=[image1_input, image2_input, image_use_post, image_use_gpen],
             outputs=image_output,
         )
         video_button.click(
             swap_video_gr,
-            inputs=[image3_input, video_input],
+            inputs=[image3_input, video_input, video_use_gpen],
             outputs=video_output,
         )
 

inference/utils.py

@@ -77,7 +77,7 @@ def get_detector(gpu_mode=False):
     return tddfa, face_boxes
 
 
-def save(x, trick=None, use_post=False):
+def save(x, trick=None, use_gpen=False):
     """ Paste img to ori_img """
     img, mat, ori_img, save_path, img_mask = x
     if mat is None:
@@ -125,7 +125,7 @@ def save(x, trick=None, use_post=False):
     ori_img = ori_img.astype(np.uint8)
 
     if trick is not None:
-        ori_img = trick.gpen(ori_img, use_post)
+        ori_img = trick.gpen(ori_img, use_gpen)
 
     Image.fromarray(ori_img).save(save_path)