Update app.py

app.py

@@ -1,51 +1,82 @@
 import argparse
-
+import numpy as np
+import cv2
+import os
 import gradio as gr
+import insightface
+from insightface.app import FaceAnalysis
 from dofaker import FaceSwapper, PoseSwapper
 
-
 def parse_args():
-    parser = argparse.ArgumentParser(description='⏳running face swap')
+    parser = argparse.ArgumentParser(description='⏳ Running face swap')
     parser.add_argument(
         '--inbrowser',
-        help=
-        'whether to automatically launch the interface in a new tab on the default browser.',
+        help='Automatically launch the interface in a new tab on the default browser.',
         dest='inbrowser',
         default=True)
     parser.add_argument(
         '--server_port',
-        help=
-        'will start gradio app on this port (if available). Can be set by environment variable GRADIO_SERVER_PORT. If None, will search for an available port starting at 7860.',
+        help='Start gradio app on this port (if available).',
         dest='server_port',
         type=int,
         default=None)
     return parser.parse_args()
 
-
-def swap_face(input_path, dst_path, src_path, use_enhancer, use_sr, scale,
-              face_sim_thre):
-    faker = FaceSwapper(use_enhancer=use_enhancer,
-                        use_sr=use_sr,
-                        scale=scale,
-                        face_sim_thre=face_sim_thre)
+def swap_face(input_path, dst_path, src_path, use_enhancer, use_sr, scale, face_sim_thre):
+    faker = FaceSwapper(use_enhancer=use_enhancer, use_sr=use_sr, scale=scale, face_sim_thre=face_sim_thre)
     output_path = faker.run(input_path, dst_path, src_path)
     return output_path
 
-
 def swap_pose(input_path, target_path, use_enhancer, use_sr, scale):
     faker = PoseSwapper(use_enhancer=use_enhancer, use_sr=use_sr, scale=scale)
     output_path = faker.run(input_path, target_path)
     return output_path
 
+def add_bbox_padding(bbox, margin=5):
+    return [
+        bbox[0] - margin,
+        bbox[1] - margin,
+        bbox[2] + margin,
+        bbox[3] + margin
+    ]
+
+def select_handler(img, evt):
+    faces = app.get(img)
+    faces = sorted(faces, key=lambda x: x.bbox[0])
+    cropped_image = []
+    face_index = -1
+    sel_face_index = 0
+    for face in faces:
+        box = face.bbox.astype(np.int32)
+        face_index += 1
+        if point_in_box((box[0], box[1]), (box[2], box[3]), (evt.index[0], evt.index[1])):
+            margin = int((box[2] - box[0]) * 0.35)
+            box = add_bbox_padding(box, margin)
+            box = np.clip(box, 0, None)
+            sel_face_index = face_index            
+            cropped_image = img[box[1]:box[3], box[0]:box[2]]
+    return cropped_image, sel_face_index
+
+def point_in_box(bl, tr, p):
+    return bl[0] < p[0] < tr[0] and bl[1] < p[1] < tr[1]
+
+def get_faces(img):
+    faces = app.get(img)
+    faces = sorted(faces, key=lambda x: x.bbox[0])
+    return img, len(faces)
 
 def main():
     args = parse_args()
+    app = FaceAnalysis(name='buffalo_l')
+    app.prepare(ctx_id=0, det_size=(640, 640))
+    swapper = insightface.model_zoo.get_model('inswapper_128.onnx', download=True, download_zip=True)
 
     with gr.Blocks(theme='ysharma/huggingface@=0.0.4') as web_ui:
         gr.Markdown('# 👑 FaceClone')                              
         gr.Markdown('### Deepfake Swap Face and Pose.')
+        
         with gr.Tab('🧩 FaceSwapper'):
-            gr.Markdown(' 🧸 FaceClone:power by Gradio')
+            gr.Markdown(' 🧸 FaceClone: Powered by Gradio')
             with gr.Tab('📥 Face'):
                 with gr.Row():
                     with gr.Column():
@@ -53,40 +84,21 @@ def main():
                         image_input = gr.Image(type='filepath')
                         with gr.Row():
                             with gr.Column():
-                                gr.Markdown(
-                                    '🎭 Target source included with Face')
+                                gr.Markdown('🎭 Target source included with Face')
                                 dst_face_image = gr.Image(type='filepath')
                             with gr.Column():
-                                gr.Markdown(
-                                    '🧑‍🦳 Source Face to Replace Target Face')
+                                gr.Markdown('🧑‍🦳 Source Face to Replace Target Face')
                                 src_face_image = gr.Image(type='filepath')
-
                     with gr.Column():
                         output_image = gr.Image(type='filepath')
-                        use_enhancer = gr.Checkbox(
-                            label="👑 Face Enhance",
-                            info="🧿Whether use face enhance model.")
+                        use_enhancer = gr.Checkbox(label="👑 Face Enhance", info="🧿 Whether use face enhance model.")
                         with gr.Row():
-                            use_sr = gr.Checkbox(
-                                label="🛠️ Super resolution",
-                                info="⚒️ Whether use image resolution model.")
-                            scale = gr.Number(
-                                value=1, label='📐image super resolution scale')
+                            use_sr = gr.Checkbox(label="🛠️ Super resolution", info="⚒️ Whether use image resolution model.")
+                            scale = gr.Number(value=1, label='📐 Image super resolution scale')
                         with gr.Row():
-                            face_sim_thre = gr.Number(
-                                value=0.6,
-                                label='🧲 face similarity threshold',
-                                minimum=0.0,
-                                maximum=1.0)
-                        convert_button = gr.Button('🔍Swap')
-                        convert_button.click(fn=swap_face,
-                                             inputs=[
-                                                 image_input, dst_face_image,
-                                                 src_face_image, use_enhancer,
-                                                 use_sr, scale, face_sim_thre
-                                             ],
-                                             outputs=[output_image],
-                                             api_name='🔎Image swap')
+                            face_sim_thre = gr.Number(value=0.6, label='🧲 Face similarity threshold', minimum=0.0, maximum=1.0)
+                        convert_button = gr.Button('🔍 Swap')
+                        convert_button.click(fn=swap_face, inputs=[image_input, dst_face_image, src_face_image, use_enhancer, use_sr, scale, face_sim_thre], outputs=[output_image])
 
             with gr.Tab('🖥️ Video'):
                 with gr.Row():
@@ -95,101 +107,42 @@ def main():
                         video_input = gr.Video()
                         with gr.Row():
                             with gr.Column():
-                                gr.Markdown(
-                                    '👳 Target Face')
+                                gr.Markdown('👳 Target Face')
                                 dst_face_image = gr.Image(type='filepath')
                             with gr.Column():
-                                gr.Markdown(
-                                    '🫅 Source Face')
+                                gr.Markdown('🫅 Source Face')
                                 src_face_image = gr.Image(type='filepath')
-
                     with gr.Column():
                         output_video = gr.Video()
-                        use_enhancer = gr.Checkbox(
-                            label="📸 Face Enhance",
-                            info="📷 Whether use face enhance model.")
+                        use_enhancer = gr.Checkbox(label="📸 Face Enhance", info="📷 Whether use face enhance model.")
                         with gr.Row():
-                            use_sr = gr.Checkbox(
-                                label="🔋 Super Resolution",
-                                info="🔦Whether use image resolution model.")
-                            scale = gr.Number(
-                                value=1, label='🗞️ Super Resolution Image')
+                            use_sr = gr.Checkbox(label="🔋 Super Resolution", info="🔦 Whether use image resolution model.")
+                            scale = gr.Number(value=1, label='🗞️ Super Resolution Image')
                         with gr.Row():
-                            face_sim_thre = gr.Number(
-                                value=0.6,
-                                label='💌face similarity threshold',
-                                minimum=0.0,
-                                maximum=1.0)
+                            face_sim_thre = gr.Number(value=0.6, label='💌 Face similarity threshold', minimum=0.0, maximum=1.0)
                         convert_button = gr.Button('🎯 Swap')
-                        convert_button.click(fn=swap_face,
-                                             inputs=[
-                                                 video_input, dst_face_image,
-                                                 src_face_image, use_enhancer,
-                                                 use_sr, scale, face_sim_thre
-                                             ],
-                                             outputs=[output_video],
-                                             api_name='💽 Video Swap 💽')
+                        convert_button.click(fn=swap_face, inputs=[video_input, dst_face_image, src_face_image, use_enhancer, use_sr, scale, face_sim_thre], outputs=[output_video])
 
         with gr.Tab('🪩 PoseSwapper'):
-            gr.Markdown(' 🧸FaceClone:Deepfake power by Gradio')
+            gr.Markdown(' 🧸 FaceClone: Deepfake powered by Gradio')
             with gr.Tab('📷 Image'):
                 with gr.Row():
                     with gr.Column():
-                        gr.Markdown('🧑‍🦳The source image to be swapped')
+                        gr.Markdown('🧑‍🦳 The source image to be swapped')
                         image_input = gr.Image(type='filepath')
                         gr.Markdown('🗳️ The target image with pose')
                         target = gr.Image(type='filepath')
-
                     with gr.Column():
                         output_image = gr.Image(type='filepath')
-                        use_enhancer = gr.Checkbox(
-                            label="👑 Face Enhance",
-                            info="🧾Whether use face enhance model.")
+                        use_enhancer = gr.Checkbox(label="👑 Face Enhance", info="🧾 Whether use face enhance model.")
                         with gr.Row():
-                            use_sr = gr.Checkbox(
-                                label="🛠️ Super Resolution",
-                                info="⌛Whether use image resolution model.")
-                            scale = gr.Number(
-                                value=1, label='📸 Super resolution scale')
+                            use_sr = gr.Checkbox(label="🛠️ Super Resolution", info="⌛ Whether use image resolution model.")
+                            scale = gr.Number(value=1, label='📸 Super resolution scale')
                         convert_button = gr.Button('🎯 Swap')
-                        convert_button.click(fn=swap_pose,
-                                             inputs=[
-                                                 image_input, target,
-                                                 use_enhancer, use_sr, scale
-                                             ],
-                                             outputs=[output_image],
-                                             api_name='🎭 Image slSwap 🎭')
-
-            # with gr.Tab('Video'):
-            #     with gr.Row():
-            #         with gr.Column():
-            #             gr.Markdown('The source video to be swapped')
-            #             video_input = gr.Image(type='filepath')
-            #             gr.Markdown('The target image with pose')
-            #             target = gr.Video()
-
-            #         with gr.Column():
-            #             output_video = gr.Video()
-            #             use_enhancer = gr.Checkbox(
-            #                 label="face enhance",
-            #                 info="Whether use face enhance model.")
-            #             with gr.Row():
-            #                 use_sr = gr.Checkbox(
-            #                     label="super resolution",
-            #                     info="Whether use image resolution model.")
-            #                 scale = gr.Number(value=1,
-            #                                 label='image super resolution scale')
-            #             convert_button = gr.Button('Swap')
-            #             convert_button.click(fn=swap_pose,
-            #                                 inputs=[
-            #                                     video_input, target, use_enhancer,
-            #                                     use_sr, scale
-            #                                 ],
-            #                                 outputs=[output_video],
-            #                                 api_name='video swap')
+                        convert_button.click(fn=swap_pose, inputs=[image_input, target, use_enhancer, use_sr, scale], outputs=[output_image])
 
     web_ui.launch(inbrowser=args.inbrowser, server_port=args.server_port)
 
-
 if __name__ == '__main__':
     main()
+