added progress bar

app.py

@@ -13,9 +13,10 @@ wellcomingMessage = """
 
 swapper = FaceSwapper()
 
-def swap_single_photo(src_img, src_idx, dst_img, dst_idx):
+def swap_single_photo(src_img, src_idx, dst_img, dst_idx, progress=gr.Progress(track_tqdm=True)):
     log = ""
     try:
+        progress(0, desc="Preparing files")
         src_path = "SinglePhoto/data_src.jpg"
         dst_path = "SinglePhoto/data_dst.jpg"
         output_path = "SinglePhoto/output_swapped.jpg"
@@ -27,9 +28,11 @@ def swap_single_photo(src_img, src_idx, dst_img, dst_idx):
         cv2.imwrite(src_path, src_img_bgr)
         cv2.imwrite(dst_path, dst_img_bgr)
         log += f"Saved source to {src_path}, destination to {dst_path}\n"
+        progress(0.5, desc="Swapping faces")
         result = swapper.swap_faces(src_path, int(src_idx), dst_path, int(dst_idx))
         cv2.imwrite(output_path, result)
         log += f"Swapped and saved result to {output_path}\n"
+        progress(0.8, desc="Cleaning up")
         try:
             if os.path.exists(src_path):
                 os.remove(src_path)
@@ -38,12 +41,14 @@ def swap_single_photo(src_img, src_idx, dst_img, dst_idx):
             log += "Cleaned up temp files.\n"
         except Exception as cleanup_error:
             log += f"Cleanup error: {cleanup_error}\n"
+        progress(1, desc="Done")
         return output_path, log
     except Exception as e:
         log += f"Error: {e}\n"
+        progress(1, desc="Error")
         return None, log
 
-def swap_video(src_img, src_idx, video, dst_idx):
+def swap_video(src_img, src_idx, video, dst_idx, progress=gr.Progress(track_tqdm=True)):
     log = ""
     src_path = "VideoSwapping/data_src.jpg"
     dst_video_path = "VideoSwapping/data_dst.mp4"
@@ -60,6 +65,7 @@ def swap_video(src_img, src_idx, video, dst_idx):
     src_img_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
     cv2.imwrite(src_path, src_img_bgr)
     log += f"Saved source image to {src_path}\n"
+    progress(0.05, desc="Saved source image")
 
     if isinstance(video, str) and os.path.exists(video):
         shutil.copy(video, dst_video_path)
@@ -70,6 +76,8 @@ def swap_video(src_img, src_idx, video, dst_idx):
     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.15, desc="Extracted frames")
+
     for idx, frame_path in enumerate(frame_paths):
         out_path = os.path.join(swapped_dir, f"swapped_{idx:05d}.jpg")
         try:
@@ -79,11 +87,13 @@ def swap_video(src_img, src_idx, video, dst_idx):
         except Exception as e:
             cv2.imwrite(out_path, cv2.imread(frame_path))
             log += f"Failed to swap frame {idx}: {e}\n"
+        progress(0.15 + 0.6 * (idx + 1) / len(frame_paths), desc=f"Swapping frames ({idx+1}/{len(frame_paths)})")
     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.8, desc="Muxing audio")
 
     # Add audio from original video
     ok, audio_log = add_audio_to_video(dst_video_path, output_video_path, final_output_path)
@@ -105,7 +115,7 @@ def swap_video(src_img, src_idx, video, dst_idx):
         log += "Cleaned up temp files and folders.\n"
     except Exception as cleanup_error:
         log += f"Cleanup error: {cleanup_error}\n"
-
+    progress(1, desc="Done")
     return final_output_path, log
 
 def add_audio_to_video(original_video_path, video_no_audio_path, output_path):
@@ -130,7 +140,7 @@ def add_audio_to_video(original_video_path, video_no_audio_path, output_path):
     except subprocess.CalledProcessError as e:
         return False, e.stderr.decode()
 
-def swap_multi_src_single_dst(src_imgs, dst_img, dst_idx):
+def swap_multi_src_single_dst(src_imgs, dst_img, dst_idx, progress=gr.Progress(track_tqdm=True)):
     log = ""
     results = []
     src_dir = "MultiSrcSingleDst/src"
@@ -146,6 +156,7 @@ def swap_multi_src_single_dst(src_imgs, dst_img, dst_idx):
     dst_path = os.path.join(dst_dir, "data_dst.jpg")
     cv2.imwrite(dst_path, dst_img_bgr)
     log += f"Saved destination image to {dst_path}\n"
+    progress(0.05, desc="Saved destination image")
 
     for i, src_img in enumerate(src_imgs):
         if isinstance(src_img, tuple):
@@ -163,9 +174,11 @@ def swap_multi_src_single_dst(src_imgs, dst_img, dst_idx):
         except Exception as e:
             results.append(f"Error: {e}")
             log += f"Error swapping source {i}: {e}\n"
+        progress((i + 1) / len(src_imgs), desc=f"Swapping source {i+1}/{len(src_imgs)}")
+    progress(1, desc="Done")
     return results, log
 
-def swap_multi_src_multi_dst(src_imgs, dst_imgs, dst_indices):
+def swap_multi_src_multi_dst(src_imgs, dst_imgs, dst_indices, progress=gr.Progress(track_tqdm=True)):
     log = ""
     results = []
     src_dir = "MultiSrcMultiDst/src"
@@ -180,6 +193,8 @@ def swap_multi_src_multi_dst(src_imgs, dst_imgs, dst_indices):
     else:
         dst_indices_list = [int(idx) for idx in dst_indices]
 
+    total = max(1, len(src_imgs) * len(dst_imgs))
+    count = 0
     for i, src_img in enumerate(src_imgs):
         if isinstance(src_img, tuple):
             src_img = src_img[0]
@@ -228,9 +243,12 @@ def swap_multi_src_multi_dst(src_imgs, dst_imgs, dst_indices):
             except Exception as e:
                 results.append(f"Error: {e}")
                 log += f"Error swapping src {i} with dst {j}: {e}\n"
+            count += 1
+            progress(count / total, desc=f"Swapping ({count}/{total})")
+    progress(1, desc="Done")
     return results, log
 
-def swap_single_src_multi_dst(src_img, dst_imgs, dst_indices):
+def swap_single_src_multi_dst(src_img, dst_imgs, dst_indices, progress=gr.Progress(track_tqdm=True)):
     log = ""
     results = []
     src_dir = "SingleSrcMultiDst/src"
@@ -246,6 +264,7 @@ def swap_single_src_multi_dst(src_img, dst_imgs, dst_indices):
     src_img_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
     cv2.imwrite(src_path, src_img_bgr)
     log += f"Saved source image to {src_path}\n"
+    progress(0.05, desc="Saved source image")
 
     if isinstance(dst_indices, str):
         dst_indices_list = [int(idx.strip()) for idx in dst_indices.split(",") if idx.strip().isdigit()]
@@ -269,6 +288,8 @@ def swap_single_src_multi_dst(src_img, dst_imgs, dst_indices):
         except Exception as e:
             results.append(f"Error: {e}")
             log += f"Error swapping with destination {j}: {e}\n"
+        progress((j + 1) / len(dst_imgs), desc=f"Swapping destination {j+1}/{len(dst_imgs)}")
+    progress(1, desc="Done")
     return results, log
 
 with gr.Blocks() as demo: