Support swapping multiple faces from a single source image into multiple faces in a target image

Author: ashleykleynhans

.gitignore

@@ -164,3 +164,5 @@ checkpoints
 CodeFormer
 test_input.json
 *.jpg
+*.jpeg
+*.png

rp_handler.py

@@ -71,20 +71,16 @@ def get_many_faces(face_analyser,
 
 
 def swap_face(face_swapper,
-              face_analyser,
-              source_img,
+              source_faces,
               target_faces,
               source_index,
               target_index,
               temp_frame):
     """
     paste source_face on target image
     """
+    source_face = source_faces[source_index]
     target_face = target_faces[target_index]
-    source_face = get_one_face(face_analyser, cv2.cvtColor(np.array(source_img[source_index]), cv2.COLOR_RGB2BGR))
-    if source_face is None:
-        logger.error('No source face found')
-        raise Exception('No source face found!')
 
     return face_swapper.get(temp_frame, target_face, source_face, paste_back=True)
 
@@ -106,34 +102,45 @@ def process(source_img: Union[Image.Image, List],
 
     # detect faces that will be replaced in target_img
     target_faces = get_many_faces(face_analyser, target_img)
+
     if target_faces is not None:
         temp_frame = copy.deepcopy(target_img)
         if isinstance(source_img, list) and len(source_img) == len(target_faces):
             logger.info('Replacing the faces in the target image from left to right by order')
             for i in range(len(target_faces)):
+                source_faces = get_many_faces(face_analyser, cv2.cvtColor(np.array(source_img[i]), cv2.COLOR_RGB2BGR))
                 source_index = i
                 target_index = i
 
                 temp_frame = swap_face(
                     face_swapper,
-                    face_analyser,
-                    source_img,
+                    source_faces,
                     target_faces,
                     source_index,
                     target_index,
                     temp_frame
                 )
-        else:
+        elif len(source_img) == 1:
+            # detect source faces that will be replaced into the target image
+            source_faces = get_many_faces(face_analyser, cv2.cvtColor(np.array(source_img[0]), cv2.COLOR_RGB2BGR))
+
             if target_index == -1:
-                logger.info('Replacing all faces in the target image to the face in the source image')
+                if len(source_faces) > 1 and len(source_faces) != len(target_faces):
+                    logger.error('Number of faces in the source image and target image must match')
+                    raise Exception('Number of faces in the source image and target image must match')
+
+                if len(source_faces) == 1:
+                    logger.info('Replacing all faces in target image with the same face from the source image')
+                else:
+                    print('Replacing all faces in the target image with the faces from the source image')
+
                 for i in range(len(target_faces)):
-                    source_index = 0
+                    source_index = 0 if len(source_faces) == 1 else i
                     target_index = i
 
                     temp_frame = swap_face(
                         face_swapper,
-                        face_analyser,
-                        source_img,
+                        source_faces,
                         target_faces,
                         source_index,
                         target_index,
@@ -145,13 +152,15 @@ def process(source_img: Union[Image.Image, List],
 
                 temp_frame = swap_face(
                     face_swapper,
-                    face_analyser,
-                    source_img,
+                    source_faces,
                     target_faces,
                     source_index,
                     target_index,
                     temp_frame
                 )
+        else:
+            logger.error('Unsupported face configuration')
+            raise Exception('Unsupported face configuration')
         result = temp_frame
     else:
         logger.error('No target faces found')