Update app.py

app.py

@@ -1,4 +1,3 @@
-###
 import streamlit as st
 import numpy as np
 import cv2
@@ -7,114 +6,134 @@ from insightface.app import FaceAnalysis
 import tempfile
 import os
 
-# Initialize face analysis and load model
-app = FaceAnalysis(name='buffalo_l')
-app.prepare(ctx_id=0, det_size=(640, 640))
+st.set_page_config(page_title="Face Swapper", layout="centered")
 
-# Load the face swapper model
-swapper = insightface.model_zoo.get_model('inswapper_128.onnx', download=False, download_zip=False)
+# ------------------------------
+# Sidebar options
+# ------------------------------
+st.sidebar.title("⚙️ Settings")
 
+# CPU / GPU selection
+device_option = st.sidebar.radio("Choose Device", ["CPU", "GPU"], index=0)
+ctx_id = 0 if device_option == "GPU" else -1
+
+# ------------------------------
+# Load models
+# ------------------------------
+@st.cache_resource
+def load_models(ctx_id):
+    app = FaceAnalysis(name='buffalo_l')
+    app.prepare(ctx_id=ctx_id, det_size=(640, 640))
+
+    swapper = insightface.model_zoo.get_model(
+        'inswapper_128.onnx', download=False, download_zip=False
+    )
+    return app, swapper
+
+app, swapper = load_models(ctx_id)
+
+# ------------------------------
+# Face swapping function
+# ------------------------------
 def swap_faces_in_video(image, video, progress):
-    """
-    Swaps faces from a source image with faces detected in a video and returns the path to the output video file.
-    
-    image: Source image (as an array)
-    video: Path to the input video file
-    progress: Streamlit progress object
-    """
     source_faces = app.get(image)
-
     if len(source_faces) == 0:
-        st.error("No face detected in the source image.")
+        st.error("❌ No face detected in the source image.")
         return None
 
     source_face = source_faces[0]
 
-    # Create a temporary file to save the output video
-    output_path = tempfile.mktemp(suffix='.avi')
+    # Temporary output file (MP4 for better compatibility)
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp_out:
+        output_path = tmp_out.name
 
-    # Open the video file
     cap = cv2.VideoCapture(video)
 
-    # Get video properties for output
     frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
     frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
     frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
     fps = cap.get(cv2.CAP_PROP_FPS)
 
-    # Define the codec and create a VideoWriter object
-    fourcc = cv2.VideoWriter_fourcc(*'XVID')
+    # MP4 writer
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
     out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height))
 
-    for i in range(frame_count):
+    i = 0
+    while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
-            break  # Exit if the video is finished
+            break
 
-        # Detect faces in the current frame
         target_faces = app.get(frame)
-
-        # Create a copy of the frame for the result
         result_frame = frame.copy()
 
-        # Swap faces for each detected face in the video frame
         for target_face in target_faces:
-            result_frame = swapper.get(result_frame, target_face, source_face, paste_back=True)
+            # depending on insightface version, frame or result_frame may be needed
+            try:
+                result_frame = swapper.get(frame, target_face, source_face, paste_back=True)
+            except:
+                result_frame = swapper.get(result_frame, target_face, source_face, paste_back=True)
 
-        # Write the result frame to the output video
         out.write(result_frame)
 
-        # Update progress bar
-        progress.progress((i + 1) / frame_count)
+        i += 1
+        if frame_count > 0:
+            progress.progress(min(1.0, i / frame_count))
 
-    # Release resources
     cap.release()
     out.release()
-    
     return output_path
 
+# ------------------------------
 # Streamlit UI
-st.title("Face Swapper in Video")
-st.write("Upload an image and a video to swap faces.")
+# ------------------------------
+st.title("🎭 Face Swapper in Video")
+st.write("Upload a **source image** and a **target video**, preview them, then swap faces.")
 
-# File uploader for the source image
+# Upload files
 image_file = st.file_uploader("Upload Source Image", type=["jpg", "jpeg", "png"])
+video_file = st.file_uploader("Upload Target Video", type=["mp4", "avi"])
+
+# Preview uploaded files
+if image_file:
+    st.subheader("📷 Source Image Preview")
+    st.image(image_file, caption="Source Image", use_column_width=True)
+
+if video_file:
+    st.subheader("🎬 Target Video Preview")
+    st.video(video_file)
 
-# File uploader for the video
-video_file = st.file_uploader("Upload Video", type=["mp4", "avi"])
+# Button to process
+if st.button("🚀 Start Face Swap"):
+    if image_file and video_file:
+        source_image = cv2.imdecode(
+            np.frombuffer(image_file.read(), np.uint8), cv2.IMREAD_COLOR
+        )
 
-if st.button("Swap Faces"):
-    if image_file is not None and video_file is not None:
-        # Read the source image
-        source_image = cv2.imdecode(np.frombuffer(image_file.read(), np.uint8), cv2.IMREAD_COLOR)
-        
-        # Save the uploaded video temporarily
         with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp_video:
             tmp_video.write(video_file.read())
             tmp_video_path = tmp_video.name
 
-        # Show a spinner and a progress bar while processing
-        with st.spinner("Processing video..."):
+        with st.spinner("Processing video... Please wait ⏳"):
             progress_bar = st.progress(0)
             output_video_path = swap_faces_in_video(source_image, tmp_video_path, progress_bar)
 
         if output_video_path:
-            st.success("Face swapping completed!")
-            # Play the processed video in Streamlit
+            st.success("✅ Face swapping completed!")
+            st.subheader("📺 Output Video Preview")
             st.video(output_video_path)
 
-            # Provide an option to download the processed video
             with open(output_video_path, "rb") as f:
                 st.download_button(
-                    label="Download Processed Video",
+                    label="⬇️ Download Processed Video",
                     data=f,
-                    file_name="output_swapped_video.avi",
-                    mime="video/x-msvideo"
+                    file_name="output_swapped_video.mp4",
+                    mime="video/mp4"
                 )
 
-            # Clean up temporary files
-            os.remove(tmp_video_path)  # Clean up temporary video file
-            # Optionally, keep the output video after displaying
-            # os.remove(output_video_path)  # Uncomment to delete after displaying
+            # Cleanup
+            os.remove(tmp_video_path)
+            # Optionally delete output file after preview
+            # os.remove(output_video_path)
     else:
-        st.error("Please upload both an image and a video.")
+        st.error("⚠️ Please upload both a source image and a video.")