Update app.py

app.py

@@ -5,7 +5,6 @@ import numpy as np
 import os
 
 UPLOAD_FOLDER = 'uploads'
-ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'mp4', 'avi'}
 
 detector = dlib.get_frontal_face_detector()
 
@@ -31,18 +30,23 @@ def swap_faces(img, face1, face2):
 
     return img
 
-def process_file(image):
-    if image is None:
-        return "Please upload an image."
+def process_file(target_image, source_image):
+    if target_image is None or source_image is None:
+        return "Please upload both a target image and a source image."
     
-    filename = os.path.join(UPLOAD_FOLDER, 'uploaded_image.jpg')
-    cv2.imwrite(filename, image)
-
-    faces, img = detect_faces(filename)
-    if len(faces) < 2:
-        return "Need at least two faces to swap."
+    target_filename = os.path.join(UPLOAD_FOLDER, 'target_image.jpg')
+    source_filename = os.path.join(UPLOAD_FOLDER, 'source_image.jpg')
+    
+    cv2.imwrite(target_filename, target_image)
+    cv2.imwrite(source_filename, source_image)
 
-    swapped_img = swap_faces(img, faces[0], faces[1])
+    faces_target, img_target = detect_faces(target_filename)
+    faces_source, img_source = detect_faces(source_filename)
+    
+    if len(faces_target) < 1 or len(faces_source) < 1:
+        return "Both images must have at least one face."
+    
+    swapped_img = swap_faces(img_target, faces_target[0], faces_source[0])
     
     result_filename = os.path.join(UPLOAD_FOLDER, 'result_image.jpg')
     cv2.imwrite(result_filename, swapped_img)
@@ -53,10 +57,13 @@ if not os.path.exists(UPLOAD_FOLDER):
 
 iface = gr.Interface(
     fn=process_file,
-    inputs=gr.Image(type="numpy", label="Upload an Image"),  # Gradio handles the file upload
+    inputs=[
+        gr.Image(type="numpy", label="Upload Target Image"),  # Target image where the face will be swapped
+        gr.Image(type="numpy", label="Upload Source Face")  # Source image to extract the face from
+    ],
     outputs=gr.Image(type="numpy", label="Swapped Faces"),
     title="Face Swap",
-    description="Upload an image with at least two faces, and this tool will swap the faces."
+    description="Upload a target image and a source face image, and this tool will swap the face from the source onto the target."
 )
 
 if __name__ == "__main__":