Update run.py

run.py

@@ -2,34 +2,32 @@ import os
 from moviepy.editor import VideoFileClip
 import insightface
 import cv2
+import numpy as np
 
-# Load InsightFace model
+# Load the InsightFace model
 model = insightface.app.FaceAnalysis(name="buffalo_l", providers=["CPUExecutionProvider"])
 model.prepare(ctx_id=0)
 
 def swap_face(source_path, target_path, output_path):
     source_img = cv2.imread(source_path)
-    faces = model.get(source_img)
-    if len(faces) == 0:
+    source_faces = model.get(source_img)
+    if len(source_faces) == 0:
         raise Exception("No face found in source.")
-    source_face = faces[0]
+    source_face = source_faces[0]
+    source_crop = source_img[
+        int(source_face.bbox[1]):int(source_face.bbox[3]),
+        int(source_face.bbox[0]):int(source_face.bbox[2])
+    ]
 
     if target_path.endswith(".jpg"):
         target_img = cv2.imread(target_path)
-        faces = model.get(target_img)
-        if len(faces) == 0:
-            raise Exception("No face found in target.")
-
-        # Crop and resize source face
-        sx1, sy1, sx2, sy2 = source_face.bbox.astype(int)
-        source_crop = source_img[sy1:sy2, sx1:sx2]
-
-        tx1, ty1, tx2, ty2 = faces[0].bbox.astype(int)
-        target_w, target_h = tx2 - tx1, ty2 - ty1
-        source_resized = cv2.resize(source_crop, (target_w, target_h))
-
-        # Paste on target
-        target_img[ty1:ty2, tx1:tx2] = source_resized
+        target_faces = model.get(target_img)
+        if len(target_faces) == 0:
+            raise Exception("No face found in target image.")
+        target_face = target_faces[0]
+        x1, y1, x2, y2 = map(int, target_face.bbox)
+        resized_crop = cv2.resize(source_crop, (x2 - x1, y2 - y1))
+        target_img[y1:y2, x1:x2] = resized_crop
         cv2.imwrite(output_path, target_img)
 
     else:
@@ -39,8 +37,16 @@ def swap_face(source_path, target_path, output_path):
             frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
             faces = model.get(frame_bgr)
             if len(faces) > 0:
-                tx1, ty1, tx2, ty2 = faces[0].bbox.astype(int)
-                target_w, target_h = tx2 - tx1, ty2 - ty1
-
-                sx1, sy1, sx2, sy2 = source_face.bbox.astype(int)
-                source_crop = source_img
+                face = faces[0]
+                x1, y1, x2, y2 = map(int, face.bbox)
+                resized_crop = cv2.resize(source_crop, (x2 - x1, y2 - y1))
+                frame_bgr[y1:y2, x1:x2] = resized_crop
+            return cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
+
+        new_clip = clip.fl_image(process_frame)
+        new_clip.write_videofile(output_path, audio=True)
+
+def run_faceswap(source_path, target_path):
+    output_path = "result.jpg" if target_path.endswith(".jpg") else "result.mp4"
+    swap_face(source_path, target_path, output_path)
+    return output_path