added swap face custom in single photo and video

.gitignore

@@ -5,6 +5,8 @@ VideoSwappingAllFaces/
 MultiSrcSingleDst/
 MultiSrcMultiDst/
 SingleSrcMultiDst/
+CustomSwap/
+CustomVideoSwap/
 __pycache__/*
 .gradio/
 

app.py

@@ -16,6 +16,7 @@ swapper = FaceSwapper()
 
 def swap_single_photo(src_img, src_idx, dst_img, dst_idx, progress=gr.Progress(track_tqdm=True)):
     log = ""
+    start_time = time.time()
     try:
         progress(0, desc="Preparing files")
         src_path = "SinglePhoto/data_src.jpg"
@@ -43,14 +44,19 @@ def swap_single_photo(src_img, src_idx, dst_img, dst_idx, progress=gr.Progress(t
         except Exception as cleanup_error:
             log += f"Cleanup error: {cleanup_error}\n"
         progress(1, desc="Done")
+        elapsed = time.time() - start_time
+        log += f"Elapsed time: {elapsed:.2f} seconds\n"
         return output_path, log
     except Exception as e:
         log += f"Error: {e}\n"
         progress(1, desc="Error")
+        elapsed = time.time() - start_time
+        log += f"Elapsed time: {elapsed:.2f} seconds\n"
         return None, log
 
 def swap_video(src_img, src_idx, video, dst_idx, progress=gr.Progress()):
     log = ""
+    start_time = time.time()
     src_path = "VideoSwapping/data_src.jpg"
     dst_video_path = "VideoSwapping/data_dst.mp4"
     frames_dir = "VideoSwapping/video_frames"
@@ -83,14 +89,14 @@ def swap_video(src_img, src_idx, video, dst_idx, progress=gr.Progress()):
 
     # Prepare swapped frames list and resume if possible
     swapped_files = set(os.listdir(swapped_dir))
-    start_time = time.time()
     total_frames = len(frame_paths)
+    start_loop_time = time.time()
     for idx, frame_path in enumerate(frame_paths):
         swapped_name = f"swapped_{idx:05d}.jpg"
         out_path = os.path.join(swapped_dir, swapped_name)
         if swapped_name in swapped_files and os.path.exists(out_path):
             log += f"Frame {idx}: already swapped, skipping.\n"
-            elapsed = time.time() - start_time
+            elapsed = time.time() - start_loop_time
             avg_time = elapsed / (idx + 1) if idx + 1 > 0 else 0
             remaining = avg_time * (total_frames - (idx + 1))
             mins, secs = divmod(int(remaining), 60)
@@ -110,7 +116,7 @@ def swap_video(src_img, src_idx, video, dst_idx, progress=gr.Progress()):
         except Exception as e:
             cv2.imwrite(out_path, cv2.imread(frame_path))
             log += f"Failed to swap frame {idx}: {e}\n"
-        elapsed = time.time() - start_time
+        elapsed = time.time() - start_loop_time
         avg_time = elapsed / (idx + 1) if idx + 1 > 0 else 0
         remaining = avg_time * (total_frames - (idx + 1))
         mins, secs = divmod(int(remaining), 60)
@@ -143,6 +149,8 @@ def swap_video(src_img, src_idx, video, dst_idx, progress=gr.Progress()):
     except Exception as cleanup_error:
         log += f"Cleanup error: {cleanup_error}\n"
     progress(1, desc="Done")
+    elapsed = time.time() - start_time
+    log += f"Elapsed time: {elapsed:.2f} seconds\n"
     return final_output_path, log
 
 def add_audio_to_video(original_video_path, video_no_audio_path, output_path):
@@ -325,6 +333,7 @@ def swap_video_all_faces(src_img, video, num_faces_to_swap, progress=gr.Progress
     For each frame, swaps one by one, always using the latest swapped image for the next face.
     """
     log = ""
+    start_time = time.time()
     src_path = "VideoSwappingAllFaces/data_src.jpg"
     dst_video_path = "VideoSwappingAllFaces/data_dst.mp4"
     frames_dir = "VideoSwappingAllFaces/video_frames"
@@ -357,8 +366,8 @@ def swap_video_all_faces(src_img, video, num_faces_to_swap, progress=gr.Progress
     swapped_files = set(os.listdir(swapped_dir))
     temp_dir = os.path.join(swapped_dir, "temp_swap")
     os.makedirs(temp_dir, exist_ok=True)
-    start_time = time.time()
     total_frames = len(frame_paths)
+    start_loop_time = time.time()
 
     for idx, frame_path in enumerate(frame_paths):
         swapped_name = f"swapped_{idx:05d}.jpg"
@@ -366,7 +375,7 @@ def swap_video_all_faces(src_img, video, num_faces_to_swap, progress=gr.Progress
         temp_frame_path = os.path.join(temp_dir, "temp.jpg")
         if swapped_name in swapped_files and os.path.exists(out_path):
             log += f"Frame {idx}: already swapped, skipping.\n"
-            elapsed = time.time() - start_time
+            elapsed = time.time() - start_loop_time
             avg_time = elapsed / (idx + 1) if idx + 1 > 0 else 0
             remaining = avg_time * (total_frames - (idx + 1))
             mins, secs = divmod(int(remaining), 60)
@@ -387,7 +396,7 @@ def swap_video_all_faces(src_img, video, num_faces_to_swap, progress=gr.Progress
         except Exception as e:
             cv2.imwrite(out_path, cv2.imread(frame_path))
             log += f"Failed to swap frame {idx}: {e}\n"
-        elapsed = time.time() - start_time
+        elapsed = time.time() - start_loop_time
         avg_time = elapsed / (idx + 1) if idx + 1 > 0 else 0
         remaining = avg_time * (total_frames - (idx + 1))
         mins, secs = divmod(int(remaining), 60)
@@ -422,8 +431,223 @@ def swap_video_all_faces(src_img, video, num_faces_to_swap, progress=gr.Progress
     except Exception as cleanup_error:
         log += f"Cleanup error: {cleanup_error}\n"
     progress(1, desc="Done")
+    elapsed = time.time() - start_time
+    log += f"Elapsed time: {elapsed:.2f} seconds\n"
+    return final_output_path, log
+
+def swap_faces_custom(src_imgs, dst_img, mapping_str, progress=gr.Progress(track_tqdm=True)):
+    """
+    src_imgs: list of source images (numpy arrays)
+    dst_img: destination image (numpy array)
+    mapping_str: comma-separated string, e.g. "2,1,3"
+    """
+    log = ""
+    start_time = time.time()
+    dst_path = "CustomSwap/data_dst.jpg"
+    output_path = "CustomSwap/output_swapped.jpg"
+    src_dir = "CustomSwap/src"
+    temp_dir = "CustomSwap/temp"
+    os.makedirs(src_dir, exist_ok=True)
+    os.makedirs(temp_dir, exist_ok=True)
+    os.makedirs(os.path.dirname(dst_path), exist_ok=True)
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+
+    # Save destination image
+    dst_img_bgr = cv2.cvtColor(dst_img, cv2.COLOR_RGB2BGR)
+    cv2.imwrite(dst_path, dst_img_bgr)
+    log += f"Saved destination image to {dst_path}\n"
+
+    # Save all source images
+    src_paths = []
+    for i, src_img in enumerate(src_imgs):
+        src_path = os.path.join(src_dir, f"data_src_{i+1}.jpg")
+        if isinstance(src_img, tuple):
+            src_img = src_img[0]
+        if src_img is None:
+            log += f"Source image {i+1} is None, skipping.\n"
+            continue
+        if isinstance(src_img, np.ndarray):
+            src_img_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
+            cv2.imwrite(src_path, src_img_bgr)
+            src_paths.append(src_path)
+            log += f"Saved source image {i+1} to {src_path}\n"
+        elif isinstance(src_img, str) and os.path.exists(src_img):
+            shutil.copy(src_img, src_path)
+            src_paths.append(src_path)
+            log += f"Copied source image {i+1} from {src_img} to {src_path}\n"
+        else:
+            log += f"Source image {i+1} is not a valid image, skipping.\n"
+
+    # Parse mapping
+    try:
+        mapping = [int(x.strip()) for x in mapping_str.split(",") if x.strip().isdigit()]
+    except Exception as e:
+        log += f"Error parsing mapping: {e}\n"
+        elapsed = time.time() - start_time
+        log += f"Elapsed time: {elapsed:.2f} seconds\n"
+        return None, log
+
+    # Use a temp file for iterative swapping
+    temp_dst_path = os.path.join(temp_dir, "temp_dst.jpg")
+    shutil.copy(dst_path, temp_dst_path)
+
+    for face_idx, src_idx in enumerate(mapping, start=1):
+        if src_idx < 1 or src_idx > len(src_paths):
+            log += f"Invalid source index {src_idx} for face {face_idx}, skipping.\n"
+            continue
+        try:
+            swapped_img = swapper.swap_faces(src_paths[src_idx-1], 1, temp_dst_path, face_idx)
+            cv2.imwrite(temp_dst_path, swapped_img)
+            log += f"Swapped face {face_idx} in destination with source {src_idx}\n"
+        except Exception as e:
+            log += f"Failed to swap face {face_idx} with source {src_idx}: {e}\n"
+
+    shutil.copy(temp_dst_path, output_path)
+    log += f"Saved swapped image to {output_path}\n"
+    if os.path.exists(temp_dst_path):
+        os.remove(temp_dst_path)
+    elapsed = time.time() - start_time
+    log += f"Elapsed time: {elapsed:.2f} seconds\n"
+    return output_path, log
+
+def swap_video_custom_mapping(src_imgs, video, mapping_str, progress=gr.Progress()):
+    """
+    Swaps faces in each frame of the video according to the user mapping.
+    src_imgs: list of source images (numpy arrays)
+    video: path or numpy array of the video
+    mapping_str: comma-separated string, e.g. "2,1,3"
+    """
+    log = ""
+    start_time = time.time()
+    src_dir = "CustomVideoSwap/src"
+    temp_dir = "CustomVideoSwap/temp"
+    frames_dir = "CustomVideoSwap/frames"
+    swapped_dir = "CustomVideoSwap/swapped_frames"
+    output_video_path = "CustomVideoSwap/output_tmp_output_video.mp4"
+    final_output_path = "CustomVideoSwap/output_with_audio.mp4"
+    dst_video_path = "CustomVideoSwap/data_dst.mp4"
+
+    os.makedirs(src_dir, exist_ok=True)
+    os.makedirs(temp_dir, exist_ok=True)
+    os.makedirs(frames_dir, exist_ok=True)
+    os.makedirs(swapped_dir, exist_ok=True)
+
+    # Save all source images
+    src_paths = []
+    for i, src_img in enumerate(src_imgs):
+        src_path = os.path.join(src_dir, f"data_src_{i+1}.jpg")
+        if isinstance(src_img, tuple):
+            src_img = src_img[0]
+        if src_img is None:
+            log += f"Source image {i+1} is None, skipping.\n"
+            continue
+        if isinstance(src_img, np.ndarray):
+            src_img_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
+            cv2.imwrite(src_path, src_img_bgr)
+            src_paths.append(src_path)
+            log += f"Saved source image {i+1} to {src_path}\n"
+        elif isinstance(src_img, str) and os.path.exists(src_img):
+            shutil.copy(src_img, src_path)
+            src_paths.append(src_path)
+            log += f"Copied source image {i+1} from {src_img} to {src_path}\n"
+        else:
+            log += f"Source image {i+1} is not a valid image, skipping.\n"
+
+    # Parse mapping
+    try:
+        mapping = [int(x.strip()) for x in mapping_str.split(",") if x.strip().isdigit()]
+    except Exception as e:
+        log += f"Error parsing mapping: {e}\n"
+        elapsed = time.time() - start_time
+        log += f"Elapsed time: {elapsed:.2f} seconds\n"
+        return None, log
+
+    # Prepare video
+    if isinstance(video, str) and os.path.exists(video):
+        shutil.copy(video, dst_video_path)
+        log += f"Copied video to {dst_video_path}\n"
+    else:
+        dst_video_path = video
+
+    from VideoSwapping import extract_frames, frames_to_video
+
+    frame_paths = extract_frames(dst_video_path, frames_dir)
+    log += f"Extracted {len(frame_paths)} frames to {frames_dir}\n"
+    progress(0.1, desc="Extracted frames")
+
+    swapped_files = set(os.listdir(swapped_dir))
+    temp_frame_path = os.path.join(temp_dir, "temp.jpg")
+    total_frames = len(frame_paths)
+    start_loop_time = time.time()
+
+    for idx, frame_path in enumerate(frame_paths):
+        swapped_name = f"swapped_{idx:05d}.jpg"
+        out_path = os.path.join(swapped_dir, swapped_name)
+        if swapped_name in swapped_files and os.path.exists(out_path):
+            log += f"Frame {idx}: already swapped, skipping.\n"
+            elapsed = time.time() - start_loop_time
+            avg_time = elapsed / (idx + 1) if idx + 1 > 0 else 0
+            remaining = avg_time * (total_frames - (idx + 1))
+            mins, secs = divmod(int(remaining), 60)
+            progress(0.1 + 0.7 * (idx + 1) / total_frames, desc=f"Swapping {idx+1}/{total_frames} | {mins:02d}:{secs:02d} left")
+            continue
+        try:
+            shutil.copy(frame_path, temp_frame_path)
+            for face_idx, src_idx in enumerate(mapping, start=1):
+                if src_idx < 1 or src_idx > len(src_paths):
+                    log += f"Invalid source index {src_idx} for face {face_idx} in frame {idx}, skipping.\n"
+                    continue
+                try:
+                    swapped_img = swapper.swap_faces(src_paths[src_idx-1], 1, temp_frame_path, face_idx)
+                    cv2.imwrite(temp_frame_path, swapped_img)
+                    log += f"Frame {idx}: Swapped face {face_idx} with source {src_idx}\n"
+                except Exception as e:
+                    log += f"Frame {idx}: Failed to swap face {face_idx} with source {src_idx}: {e}\n"
+            shutil.copy(temp_frame_path, out_path)
+            log += f"Swapped all faces in frame {idx} and saved to {out_path}\n"
+            if os.path.exists(temp_frame_path):
+                os.remove(temp_frame_path)
+        except Exception as e:
+            cv2.imwrite(out_path, cv2.imread(frame_path))
+            log += f"Failed to swap frame {idx}: {e}\n"
+        elapsed = time.time() - start_loop_time
+        avg_time = elapsed / (idx + 1) if idx + 1 > 0 else 0
+        remaining = avg_time * (total_frames - (idx + 1))
+        mins, secs = divmod(int(remaining), 60)
+        progress(0.1 + 0.7 * (idx + 1) / total_frames, desc=f"Swapping {idx+1}/{total_frames} | {mins:02d}:{secs:02d} left")
+    if os.path.exists(temp_dir):
+        shutil.rmtree(temp_dir)
+
+    cap = cv2.VideoCapture(dst_video_path)
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    cap.release()
+    frames_to_video(swapped_dir, output_video_path, fps)
+    log += f"Combined swapped frames into video {output_video_path}\n"
+    progress(0.9, desc="Muxing audio")
+
+    ok, audio_log = add_audio_to_video(dst_video_path, output_video_path, final_output_path)
+    if ok:
+        log += f"Added audio to {final_output_path}\n"
+    else:
+        log += f"Audio muxing failed: {audio_log}\n"
+        final_output_path = output_video_path
+
+    try:
+        if os.path.exists(dst_video_path):
+            os.remove(dst_video_path)
+        if os.path.exists(frames_dir):
+            shutil.rmtree(frames_dir)
+        if os.path.exists(swapped_dir):
+            shutil.rmtree(swapped_dir)
+        log += "Cleaned up temp files and folders.\n"
+    except Exception as cleanup_error:
+        log += f"Cleanup error: {cleanup_error}\n"
+    progress(1, desc="Done")
+    elapsed = time.time() - start_time
+    log += f"Elapsed time: {elapsed:.2f} seconds\n"
     return final_output_path, log
 
+# Add this to your Gradio UI:
 with gr.Blocks() as demo:
     gr.Markdown(wellcomingMessage)
     with gr.Tab("Single Photo Swapping"):
@@ -440,38 +664,38 @@ with gr.Blocks() as demo:
                 gr.Textbox(label="Log Output", lines=8, interactive=False)
             ],
         )
-    with gr.Tab("Video Swapping"):
+    
+    with gr.Tab("SingleSrc MultiDst"):
         gr.Interface(
-            fn=swap_video,
+            fn=swap_single_src_multi_dst,
             inputs=[
                 gr.Image(label="Source Image"),
-                gr.Number(value=1, label="Source Face Index"),
-                gr.Video(label="Target Video"),
-                gr.Number(value=1, label="Destination Face Index"),
+                gr.Gallery(label="Destination Images", type="numpy", columns=3),
+                gr.Textbox(label="Destination Face Indices (comma-separated, e.g. 1,1,2)"),
             ],
             outputs=[
-                gr.Video(label="Swapped Video"),
+                gr.Gallery(label="Swapped Images"),
                 gr.Textbox(label="Log Output", lines=8, interactive=False)
             ],
         )
-    with gr.Tab("Video All Faces"):
+    with gr.Tab("MultiSrc SingleDst"):
         gr.Interface(
-            fn=swap_video_all_faces,
+            fn=swap_multi_src_single_dst,
             inputs=[
-                gr.Image(label="Source Image"),
-                gr.Video(label="Target Video"),
-                gr.Number(value=1, label="Number of Faces to Swap"),
+                gr.Gallery(label="Source Images", type="numpy", columns=3),
+                gr.Image(label="Destination Image"),
+                gr.Number(value=1, label="Destination Face Index"),
             ],
             outputs=[
-                gr.Video(label="Swapped Video"),
+                gr.Gallery(label="Swapped Images"),
                 gr.Textbox(label="Log Output", lines=8, interactive=False)
             ],
         )
-    with gr.Tab("SingleSrc MultiDst"):
+    with gr.Tab("MultiSrc MultiDst"):
         gr.Interface(
-            fn=swap_single_src_multi_dst,
+            fn=swap_multi_src_multi_dst,
             inputs=[
-                gr.Image(label="Source Image"),
+                gr.Gallery(label="Source Images", type="numpy", columns=3),
                 gr.Gallery(label="Destination Images", type="numpy", columns=3),
                 gr.Textbox(label="Destination Face Indices (comma-separated, e.g. 1,1,2)"),
             ],
@@ -480,29 +704,56 @@ with gr.Blocks() as demo:
                 gr.Textbox(label="Log Output", lines=8, interactive=False)
             ],
         )
-    with gr.Tab("MultiSrc SingleDst"):
+    with gr.Tab("Custom Face Mapping"):
         gr.Interface(
-            fn=swap_multi_src_single_dst,
+            fn=swap_faces_custom,
             inputs=[
                 gr.Gallery(label="Source Images", type="numpy", columns=3),
                 gr.Image(label="Destination Image"),
+                gr.Textbox(label="Mapping (comma-separated, e.g. 2,1,3)"),
+            ],
+            outputs=[
+                gr.Image(label="Swapped Image"),
+                gr.Textbox(label="Log Output", lines=8, interactive=False)
+            ],
+        )
+    with gr.Tab("Video Swapping"):
+        gr.Interface(
+            fn=swap_video,
+            inputs=[
+                gr.Image(label="Source Image"),
+                gr.Number(value=1, label="Source Face Index"),
+                gr.Video(label="Target Video"),
                 gr.Number(value=1, label="Destination Face Index"),
             ],
             outputs=[
-                gr.Gallery(label="Swapped Images"),
+                gr.Video(label="Swapped Video"),
                 gr.Textbox(label="Log Output", lines=8, interactive=False)
             ],
         )
-    with gr.Tab("MultiSrc MultiDst"):
+    with gr.Tab("Video All Faces"):
         gr.Interface(
-            fn=swap_multi_src_multi_dst,
+            fn=swap_video_all_faces,
+            inputs=[
+                gr.Image(label="Source Image"),
+                gr.Video(label="Target Video"),
+                gr.Number(value=1, label="Number of Faces to Swap"),
+            ],
+            outputs=[
+                gr.Video(label="Swapped Video"),
+                gr.Textbox(label="Log Output", lines=8, interactive=False)
+            ],
+        )
+    with gr.Tab("Custom Video Face Mapping"):
+        gr.Interface(
+            fn=swap_video_custom_mapping,
             inputs=[
                 gr.Gallery(label="Source Images", type="numpy", columns=3),
-                gr.Gallery(label="Destination Images", type="numpy", columns=3),
-                gr.Textbox(label="Destination Face Indices (comma-separated, e.g. 1,1,2)"),
+                gr.Video(label="Target Video"),
+                gr.Textbox(label="Mapping (comma-separated, e.g. 2,1,3)"),
             ],
             outputs=[
-                gr.Gallery(label="Swapped Images"),
+                gr.Video(label="Swapped Video"),
                 gr.Textbox(label="Log Output", lines=8, interactive=False)
             ],
         )

requirements.txt

@@ -5,4 +5,5 @@ scikit-learn
 scikit-image
 tqdm
 mxnet
-insightface
+insightface
+gradio