Update app.py

app.py

@@ -7,15 +7,12 @@ import tempfile
 import os
 
 # --- MediaPipe Initialization ---
-# Use try-except block for robustness if mediapipe is not installed correctly
 try:
     mp_face_mesh = mp.solutions.face_mesh
-    # NOTE: refine_landmarks=True gives 478 landmarks. False gives 468.
-    # We will control density by sub-sampling rather than this boolean for more control.
     face_mesh = mp_face_mesh.FaceMesh(
         static_image_mode=True,
         max_num_faces=1,
-        refine_landmarks=True, # Keep this on for the best potential quality
+        refine_landmarks=True,
         min_detection_confidence=0.5
     )
     print("MediaPipe Face Mesh initialized successfully.")
@@ -26,284 +23,145 @@ except (ImportError, AttributeError):
 # --- Helper Functions ---
 
 def get_landmarks(img, landmark_step=1):
-    """
-    Detects face landmarks using MediaPipe Face Mesh.
-    Includes sub-sampling for performance.
-    - landmark_step: Step to sample landmarks. 1 = all, 2 = half, etc.
-    """
     if img is None:
-        print("Warning: Input image is None in get_landmarks.")
         return None
     if face_mesh is None:
-        print("Error: MediaPipe Face Mesh not available.")
         return None
-
     img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
     try:
         results = face_mesh.process(img_rgb)
-    except Exception as e:
-        print(f"Error processing image with MediaPipe: {e}")
+    except Exception:
         return None
-
     if not results.multi_face_landmarks:
-        print("Warning: No face detected.")
         return None
-
     landmarks_mp = results.multi_face_landmarks[0]
     h, w, _ = img.shape
-    
-    # Get all landmarks first
     full_landmarks = np.array([(pt.x * w, pt.y * h) for pt in landmarks_mp.landmark], dtype=np.float32)
-
-    # --- NEW: Sub-sample landmarks for speed ---
-    if landmark_step > 1:
-        # Sample with a step, ensuring correspondence is maintained between faces
-        landmarks = full_landmarks[::landmark_step]
-    else:
-        landmarks = full_landmarks
-
+    landmarks = full_landmarks[::landmark_step] if landmark_step > 1 else full_landmarks
     if not np.all(np.isfinite(landmarks)):
-        print("Warning: Invalid landmark coordinates detected (NaN/inf).")
         return None
+    corners = np.array([[0, 0], [w - 1, 0], [0, h - 1], [w - 1, h - 1]], dtype=np.float32)
+    return np.vstack((landmarks, corners))
 
-    corners = np.array([
-        [0, 0], [w - 1, 0], [0, h - 1], [w - 1, h - 1]
-    ], dtype=np.float32)
-    
-    # Always include corners for stable warping
-    landmarks = np.vstack((landmarks, corners))
-
-    return landmarks
 
 def calculate_delaunay_triangles(rect, points):
-    """Calculates Delaunay triangulation for a set of points. (No changes needed here)"""
     if points is None or len(points) < 3:
         return []
-    if not np.all(np.isfinite(points)):
-        points = points[np.all(np.isfinite(points), axis=1)]
-        if len(points) < 3: return []
-    
     points[:, 0] = np.clip(points[:, 0], rect[0], rect[0] + rect[2] - 1)
     points[:, 1] = np.clip(points[:, 1], rect[1], rect[1] + rect[3] - 1)
-    
     subdiv = cv2.Subdiv2D(rect)
-    point_map = { (int(p[0]), int(p[1])): i for i, p in enumerate(points) }
-    inserted_points_map = {}
-
+    inserted = {}
     for i, p in enumerate(points):
-        point_tuple = (int(p[0]), int(p[1]))
-        if point_tuple not in inserted_points_map:
+        tup = (int(p[0]), int(p[1]))
+        if tup not in inserted:
             try:
-                subdiv.insert(point_tuple)
-                inserted_points_map[point_tuple] = i
+                subdiv.insert(tup)
+                inserted[tup] = i
             except cv2.error:
                 continue
-    
-    triangle_list = subdiv.getTriangleList()
-    delaunay_triangles = []
-    for t in triangle_list:
-        pts_coords = [(int(t[0]), int(t[1])), (int(t[2]), int(t[3])), (int(t[4]), int(t[5]))]
-        if all(rect[0] <= p[0] < rect[0] + rect[2] and rect[1] <= p[1] < rect[1] + rect[3] for p in pts_coords):
-            indices = [inserted_points_map.get(coord) for coord in pts_coords]
-            if all(idx is not None for idx in indices) and len(set(indices)) == 3:
-                delaunay_triangles.append(indices)
-    return delaunay_triangles
+    triangles = subdiv.getTriangleList()
+    delaunay = []
+    for t in triangles:
+        coords = [(int(t[0]), int(t[1])), (int(t[2]), int(t[3])), (int(t[4]), int(t[5]))]
+        if all(rect[0] <= x < rect[0] + rect[2] and rect[1] <= y < rect[1] + rect[3] for x, y in coords):
+            idxs = [inserted.get(c) for c in coords]
+            if all(i is not None for i in idxs) and len(set(idxs)) == 3:
+                delaunay.append(idxs)
+    return delaunay
 
 
 def warp_triangle(img1, img2, t1, t2):
-    """Warps a triangle from img1 to img2. (No changes needed here)"""
-    if len(t1) != 3 or len(t2) != 3 or not np.all(np.isfinite(t1)) or not np.all(np.isfinite(t2)):
+    if len(t1) != 3 or len(t2) != 3:
         return
-    try:
-        r1 = cv2.boundingRect(np.float32([t1]))
-        r2 = cv2.boundingRect(np.float32([t2]))
-
-        if r1[2] <= 0 or r1[3] <= 0 or r2[2] <= 0 or r2[3] <= 0: return
-
-        t1_rect = [(t1[i][0] - r1[0], t1[i][1] - r1[1]) for i in range(3)]
-        t2_rect = [(t2[i][0] - r2[0], t2[i][1] - r2[1]) for i in range(3)]
-
-        mask = np.zeros((r2[3], r2[2], 3), dtype=np.float32)
-        cv2.fillConvexPoly(mask, np.int32(t2_rect), (1.0, 1.0, 1.0), 16, 0)
-        
-        img1_rect = img1[r1[1]:r1[1] + r1[3], r1[0]:r1[0] + r1[2]]
-        if img1_rect.size == 0: return
-
-        size = (r2[2], r2[3])
-        warp_mat = cv2.getAffineTransform(np.float32(t1_rect), np.float32(t2_rect))
-        img2_rect = cv2.warpAffine(img1_rect, warp_mat, size, None, flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT_101)
-
-        img2_rect *= mask
-        
-        y_start, y_end = r2[1], r2[1] + r2[3]
-        x_start, x_end = r2[0], r2[0] + r2[2]
-        
-        h_img2, w_img2, _ = img2.shape
-        if y_start >= h_img2 or x_start >= w_img2: return
+    r1 = cv2.boundingRect(np.float32([t1]))
+    r2 = cv2.boundingRect(np.float32([t2]))
+    if r1[2] == 0 or r1[3] == 0 or r2[2] == 0 or r2[3] == 0:
+        return
+    t1_rect = [(t1[i][0] - r1[0], t1[i][1] - r1[1]) for i in range(3)]
+    t2_rect = [(t2[i][0] - r2[0], t2[i][1] - r2[1]) for i in range(3)]
+    mask = np.zeros((r2[3], r2[2], 3), dtype=np.float32)
+    cv2.fillConvexPoly(mask, np.int32(t2_rect), (1.0, 1.0, 1.0), 16, 0)
+    img1_rect = img1[r1[1]:r1[1]+r1[3], r1[0]:r1[0]+r1[2]]
+    if img1_rect.size == 0:
+        return
+    warp_mat = cv2.getAffineTransform(np.float32(t1_rect), np.float32(t2_rect))
+    img2_rect = cv2.warpAffine(img1_rect, warp_mat, (r2[2], r2[3]), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT_101)
+    img2_rect *= mask
+    y1, y2 = r2[1], r2[1] + r2[3]
+    x1, x2 = r2[0], r2[0] + r2[2]
+    img2[y1:y2, x1:x2] = img2[y1:y2, x1:x2] * (1 - mask) + img2_rect
 
-        img2[y_start:y_end, x_start:x_end] = img2[y_start:y_end, x_start:x_end] * (1.0 - mask) + img2_rect
-    except (cv2.error, IndexError):
-        pass # Ignore degenerate triangles or slicing errors
 
-# --- Main Morphing Function (Modified) ---
 def morph_faces(img1_orig, img2_orig, alpha, resize_dim, landmark_step):
-    """
-    Morphs two faces with a seamless blending strategy to avoid artifacts.
-    """
-    start_time = time.time()
     if img1_orig is None or img2_orig is None:
         return np.zeros((resize_dim, resize_dim, 3), dtype=np.uint8)
-
-    # --- Preprocessing with dynamic resize_dim ---
-    try:
-        img1 = cv2.resize(img1_orig, (resize_dim, resize_dim), interpolation=cv2.INTER_LINEAR)
-        img2 = cv2.resize(img2_orig, (resize_dim, resize_dim), interpolation=cv2.INTER_LINEAR)
-    except cv2.error:
-        return np.zeros((resize_dim, resize_dim, 3), dtype=np.uint8)
-
-    h, w, _ = img1.shape
-    rect = (0, 0, w, h)
-
-    # --- Landmark Detection with dynamic landmark_step ---
+    img1 = cv2.resize(img1_orig, (resize_dim, resize_dim))
+    img2 = cv2.resize(img2_orig, (resize_dim, resize_dim))
     landmarks1 = get_landmarks(img1, landmark_step)
     landmarks2 = get_landmarks(img2, landmark_step)
-
     if landmarks1 is None or landmarks2 is None or landmarks1.shape != landmarks2.shape:
-        print("Landmark error. Falling back to simple alpha blend.")
-        return cv2.addWeighted(img1, 1 - alpha, img2, alpha, 0)
-
-    # --- Landmark Interpolation (determines the shape of the output face) ---
-    landmarks_morphed = (1 - alpha) * landmarks1 + alpha * landmarks2
-    
-    # --- Triangulation (based on the final morphed shape) ---
-    try:
-        triangles_indices = calculate_delaunay_triangles(rect, landmarks_morphed.copy())
-        if not triangles_indices:
-            print("Triangulation failed. Falling back to simple alpha blend.")
-            return cv2.addWeighted(img1, 1 - alpha, img2, alpha, 0)
-    except Exception as e:
-        print(f"Error during triangulation: {e}. Falling back to simple alpha blend.")
-        return cv2.addWeighted(img1, 1 - alpha, img2, alpha, 0)
-
-
-    # --- SEAMLESS WARPING AND BLENDING ---
-    
-    # 1. Create two empty canvases for the fully warped images
-    img1_float = img1.astype(np.float32) / 255.0
-    img2_float = img2.astype(np.float32) / 255.0
-    img1_warped = np.zeros_like(img1_float)
-    img2_warped = np.zeros_like(img2_float)
-
-    # 2. Warp triangles from each source to their morphed positions on the respective canvases
-    for indices in triangles_indices:
-        if any(idx >= len(landmarks1) for idx in indices): continue # Safety check
-        
-        # Get triangle vertices for source 1, source 2, and the morphed shape
-        t1 = landmarks1[indices]
-        t2 = landmarks2[indices]
-        t_morphed = landmarks_morphed[indices]
-        
-        # Warp the triangle from img1 to the morphed position on the img1_warped canvas
-        warp_triangle(img1_float, img1_warped, t1, t_morphed)
-        
-        # Warp the triangle from img2 to the morphed position on the img2_warped canvas
-        warp_triangle(img2_float, img2_warped, t2, t_morphed)
+        return cv2.addWeighted(img1, 1-alpha, img2, alpha, 0)
+    morphed_pts = (1-alpha)*landmarks1 + alpha*landmarks2
+    rect = (0, 0, resize_dim, resize_dim)
+    tris = calculate_delaunay_triangles(rect, morphed_pts)
+    if not tris:
+        return cv2.addWeighted(img1, 1-alpha, img2, alpha, 0)
+    img1_f = img1.astype(np.float32)/255.0
+    img2_f = img2.astype(np.float32)/255.0
+    w1 = np.zeros_like(img1_f)
+    w2 = np.zeros_like(img2_f)
+    for ids in tris:
+        t1 = landmarks1[ids]; t2 = landmarks2[ids]; tm = morphed_pts[ids]
+        warp_triangle(img1_f, w1, t1, tm)
+        warp_triangle(img2_f, w2, t2, tm)
+    morph = (1-alpha)*w1 + alpha*w2
+    return (morph*255).astype(np.uint8)
 
-    # 3. Perform a single, final alpha blend of the two completed warped images
-    morphed_img_float = (1.0 - alpha) * img1_warped + alpha * img2_warped
-    
-    # --- Final Conversion ---
-    morphed_img = (morphed_img_float * 255.0).clip(0, 255).astype(np.uint8)
-    end_time = time.time()
-    print(f"Frame morph ({w}x{h}, {len(landmarks1)} landmarks) took: {end_time - start_time:.4f}s")
-    return morphed_img
 
-# --- Video Processing Function (Modified) ---
 def process_video(video_path, target_img, transition_level, resolution, landmark_sampling):
-    """
-    Callback function that now receives resolution and landmark settings from the UI.
-    """
-
-    target_img = cv2.cvtColor(target_img, cv2.COLOR_RGB2BGR)
-    
     if video_path is None or target_img is None:
-        # Create a dummy video to avoid Gradio errors on empty inputs
-        dummy_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name
-        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-        out = cv2.VideoWriter(dummy_path, fourcc, 24, (resolution, resolution))
+        dummy = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name
+        out = cv2.VideoWriter(dummy, cv2.VideoWriter_fourcc(*'mp4v'), 24, (resolution, resolution))
         out.release()
-        return dummy_path
-
-    alpha = (transition_level + 1.0) / 2.0
-    alpha = float(np.clip(alpha, 0.0, 1.0))
-
+        return dummy
+    target_bgr = cv2.cvtColor(target_img, cv2.COLOR_RGB2BGR)
+    alpha = float(np.clip((transition_level+1)/2,0,1))
     cap = cv2.VideoCapture(video_path)
-    if not cap.isOpened():
-        raise IOError(f"Cannot open video file: {video_path}")
-    
     fps = cap.get(cv2.CAP_PROP_FPS) or 24
-    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-    tmp_out = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
-    
-    # --- Use dynamic resolution for the output video ---
-    out = cv2.VideoWriter(tmp_out.name, fourcc, fps, (resolution, resolution))
-
-    frame_count = 0
+    out_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name
+    out = cv2.VideoWriter(out_file, cv2.VideoWriter_fourcc(*'mp4v'), fps, (resolution, resolution))
     while True:
         ret, frame = cap.read()
-        if not ret:
-            break
-        
-        # Pass the new parameters to the morphing function
-        morphed = morph_faces(frame, target_img, alpha, resolution, landmark_sampling)
-        out.write(morphed)
-        frame_count += 1
-    
-    print(f"Processed {frame_count} frames.")
-    cap.release()
-    out.release()
-    return tmp_out.name
+        if not ret: break
+        mor = morph_faces(frame, target_bgr, alpha, resolution, landmark_sampling)
+        out.write(mor)
+    cap.release(); out.release()
+    return out_file
 
-# --- Gradio App (Modified) ---
+# --- Gradio App ---
 css = """video, img { object-fit: contain !important; }"""
 with gr.Blocks(css=css) as iface:
     gr.Markdown("# Real-Time Video Face Morph 🚀")
-    gr.Markdown("Adjust resolution and landmark density for a trade-off between speed and quality.")
+    gr.Markdown("Use the button below to generate and show a progress bar during processing.")
     with gr.Row():
         video_input = gr.Video(label="Input Video")
         img_input = gr.Image(type="numpy", label="Target Face Image")
-    
     with gr.Row():
-        # --- NEW: UI controls for performance ---
-        resolution_slider = gr.Dropdown(
-            [256, 384, 512, 768], 
-            value=512, 
-            label="Processing Resolution",
-            info="Lower resolution means much faster processing."
-        )
-        landmark_slider = gr.Slider(
-            1, 4, 
-            value=1, 
-            step=1, 
-            label="Landmark Sub-sampling",
-            info="1=Max Quality (~478 landmarks), 4=Max Speed (~120 landmarks)"
-        )
-
-    slider = gr.Slider(-1.0, 1.0, value=0.0, step=0.05, label="Transition Level (-1 = Video, 1 = Image)")
+        resolution_slider = gr.Dropdown([256,384,512,768], value=512, label="Resolution")
+        landmark_slider = gr.Slider(1,4,value=1,step=1, label="Landmark Sub-sampling")
+    transition_slider = gr.Slider(-1.0,1.0,value=0.0,step=0.05, label="Transition Level")
+    generate_btn = gr.Button("Generate Morph 🚀", variant="primary")
+    progress_bar = gr.Progress()
     video_output = gr.Video(label="Morphed Video")
-
-    # Gather all input components
-    inputs = [video_input, img_input, slider, resolution_slider, landmark_slider]
-
-    # Trigger processing on any input change
-    for component in inputs:
-        component.change(
-            fn=process_video,
-            inputs=inputs,
-            outputs=video_output,
-            show_progress="full"
-        )
+    
+    generate_btn.click(
+        fn=process_video,
+        inputs=[video_input, img_input, transition_slider, resolution_slider, landmark_slider],
+        outputs=video_output,
+        show_progress=True
+    )
+    
     gr.Markdown("---\n*Built with Gradio, OpenCV & MediaPipe.*")
 
 if __name__ == "__main__":
-    iface.launch(debug=True)
+    iface.launch(debug=True)