Update app.py

app.py

@@ -6,263 +6,288 @@ import time
 import tempfile
 import os
 
-# --- MediaPipe Initialization ---
+# --- MediaPipe Initialization for Video ---
+# Optimized for video by setting static_image_mode to False
+# This enables tracking and significantly improves speed and stability.
 try:
     mp_face_mesh = mp.solutions.face_mesh
     face_mesh = mp_face_mesh.FaceMesh(
-        static_image_mode=False,
+        static_image_mode=False, # <-- The most important change for video
         max_num_faces=1,
         refine_landmarks=True,
-        min_detection_confidence=0.5
+        min_detection_confidence=0.5,
+        min_tracking_confidence=0.5 # Confidence for tracking across frames
     )
-    print("MediaPipe Face Mesh initialized successfully.")
+    print("MediaPipe Face Mesh initialized for VIDEO successfully.")
 except (ImportError, AttributeError):
     print("Error: Could not initialize MediaPipe Face Mesh. Is mediapipe installed correctly?")
     face_mesh = None
 
-# --- Helper Functions ---
+# --- Helper Functions (Refactored for Efficiency) ---
 
+def get_landmarks_from_result(results, img_shape):
+    """Extracts landmarks from a MediaPipe results object."""
+    if not results or not results.multi_face_landmarks:
+        return None
+    h, w = img_shape[:2]
+    # Note: Using a NumPy array directly is faster than list comprehensions for this.
+    landmarks = np.array([(lm.x * w, lm.y * h) for lm in results.multi_face_landmarks[0].landmark], dtype=np.float32)
+
+    # Add image corners to landmarks for robust triangulation
+    corners = np.array([[0, 0], [w - 1, 0], [0, h - 1], [w - 1, h - 1]], dtype=np.float32)
+    return np.vstack((landmarks, corners))
 
 
-def get_face_mask_box(img, feather, padding=0):
-    h, w = img.shape[:2]
+def get_face_mask_box_from_landmarks(landmarks, img_shape, feather, padding=0):
+    """Generates a face mask and bounding box from pre-computed landmarks."""
+    h, w = img_shape[:2]
+    # We only need the facial landmarks, not the corners we added
+    face_landmarks = landmarks[:-4]
+    
+    hull = cv2.convexHull(face_landmarks.astype(np.int32))
     mask = np.zeros((h, w), dtype=np.uint8)
-    results = face_mesh.process(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
-    if not results.multi_face_landmarks:
-        return None, None
-    pts = np.array([(int(p.x * w), int(p.y * h)) for p in results.multi_face_landmarks[0].landmark], np.int32)
-    hull = cv2.convexHull(pts)
     cv2.fillConvexPoly(mask, hull, 255)
+
     x, y, bw, bh = cv2.boundingRect(hull)
-    # apply padding
+    
+    # Apply padding
     x_pad = max(x - padding, 0)
     y_pad = max(y - padding, 0)
-    x2 = min(x + bw + padding, w)
-    y2 = min(y + bh + padding, h)
-    mask_roi = mask[y_pad:y2, x_pad:x2]
-    # inside feather
-    if feather > 0 and mask_roi.size > 0:
-        k = int(feather)
-        mask_roi = cv2.GaussianBlur(mask_roi, (k*2+1, k*2+1), 0)
-    return mask_roi, (x_pad, y_pad, x2 - x_pad, y2 - y_pad)
+    x2_pad = min(x + bw + padding, w)
+    y2_pad = min(y + bh + padding, h)
 
+    # Feather the mask for smoother blending
+    if feather > 0:
+        k = int(feather)
+        if k % 2 == 0: k += 1 # Kernel size must be odd
+        mask = cv2.GaussianBlur(mask, (k, k), 0)
 
-def cut_and_feather(img, feather):
-    h, w = img.shape[:2]
-    mask = np.zeros((h, w), dtype=np.uint8)
-    results = face_mesh.process(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
-    if not results.multi_face_landmarks:
-        return np.zeros_like(img), None, None
-    pts = np.array([(int(p.x * w), int(p.y * h)) for p in results.multi_face_landmarks[0].landmark], np.int32)
-    hull = cv2.convexHull(pts)
-    cv2.fillConvexPoly(mask, hull, 255)
-    # bounding box
-    x, y, bw, bh = cv2.boundingRect(hull)
-    # feather mask
-    k = int(feather)
-    if k > 0:
-        mask = cv2.GaussianBlur(mask, (k*2+1, k*2+1), 0)
-    # extract face ROI
-    face_roi = img[y:y+bh, x:x+bw]
-    mask_roi = mask[y:y+bh, x:x+bw]
-    # apply mask
-    fg = cv2.bitwise_and(face_roi, face_roi, mask=mask_roi)
-    # prepare alpha
-    alpha = mask_roi.astype(np.float32) / 255.0
-    # composite onto transparent background same size
-    out = (fg.astype(np.float32) * alpha[..., None]).astype(np.uint8)
-    return out, mask_roi, (x, y, bw, bh)
-
-def get_landmarks(img, landmark_step=1):
-    if img is None or face_mesh is None:
-        return None
-    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-    try:
-        results = face_mesh.process(img_rgb)
-    except Exception:
-        return None
-    if not results.multi_face_landmarks:
-        return None
-    landmarks_mp = results.multi_face_landmarks[0]
-    h, w, _ = img.shape
-    pts = np.array([(pt.x * w, pt.y * h) for pt in landmarks_mp.landmark], dtype=np.float32)
-    landmarks = pts[::landmark_step] if landmark_step > 1 else pts
-    if not np.all(np.isfinite(landmarks)):
-        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))
-
+    box = (x_pad, y_pad, x2_pad - x_pad, y2_pad - y_pad)
+    mask_roi = mask[y_pad:y2_pad, x_pad:x2_pad]
+    
+    return mask_roi, box
 
 def calculate_delaunay_triangles(rect, points):
-    if points is None or len(points)<3:
+    """Calculates Delaunay triangles for a set of points."""
+    if points is None or 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)
-    inserted = {}
-    for i,p in enumerate(points):
-        key = (int(p[0]), int(p[1]))
-        if key not in inserted:
-            try:
-                subdiv.insert(key)
-                inserted[key]=i
-            except cv2.error:
-                continue
-    tris = subdiv.getTriangleList()
-    delaunay=[]
-    for t in tris:
-        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):
-    if len(t1)!=3 or len(t2)!=3:
-        return
-    r1=cv2.boundingRect(np.float32([t1]))
-    r2=cv2.boundingRect(np.float32([t2]))
+    # Using a dictionary is faster for checking existing points
+    point_map = { (int(p[0]), int(p[1])): i for i, p in enumerate(points) }
+    for p in point_map.keys():
+        subdiv.insert(p)
+
+    triangle_list = subdiv.getTriangleList()
+    delaunay_triangles = []
     
-    if r1[2] <= 0 or r1[3] <= 0 or r2[2] <= 0 or r2[3] <= 0:
-        return
+    for t in triangle_list:
+        pt1 = (int(t[0]), int(t[1]))
+        pt2 = (int(t[2]), int(t[3]))
+        pt3 = (int(t[4]), int(t[5]))
+        
+        # Check if all points are within the image boundaries
+        if rect[0] <= pt1[0] < rect[0] + rect[2] and rect[1] <= pt1[1] < rect[1] + rect[3] and \
+           rect[0] <= pt2[0] < rect[0] + rect[2] and rect[1] <= pt2[1] < rect[1] + rect[3] and \
+           rect[0] <= pt3[0] < rect[0] + rect[2] and rect[1] <= pt3[1] < rect[1] + rect[3]:
+            
+            # Get the indices from our original point list
+            idx1 = point_map.get(pt1)
+            idx2 = point_map.get(pt2)
+            idx3 = point_map.get(pt3)
 
-    img1_rect = img1[r1[1]:r1[1]+r1[3], r1[0]:r1[0]+r1[2]]
-    if img1_rect.size == 0: 
-        return
+            if idx1 is not None and idx2 is not None and idx3 is not None:
+                delaunay_triangles.append([idx1, idx2, idx3])
+                
+    return delaunay_triangles
 
-    t1r=[(t1[i][0]-r1[0],t1[i][1]-r1[1]) for i in range(3)]
-    t2r=[(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(t2r),(1,1,1),16)
+
+def warp_triangle(img1, img2, t1, t2):
+    """Warps a triangle from img1 to img2."""
+    r1 = cv2.boundingRect(t1)
+    r2 = cv2.boundingRect(t2)
+
+    # Crop triangle ROI
+    t1_rect = t1 - r1[:2]
+    t2_rect = t2 - r2[:2]
     
-    src=img1[r1[1]:r1[1]+r1[3],r1[0]:r1[0]+r1[2]]
-    M=cv2.getAffineTransform(np.float32(t1r),np.float32(t2r))
-    warped=cv2.warpAffine(src,M,(r2[2],r2[3]),flags=cv2.INTER_LINEAR,borderMode=cv2.BORDER_REFLECT_101)
-    warped*=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)+warped
-
-
-def morph_faces(img1, img2, alpha, dim, step):
-    if img1 is None or img2 is None:
-        return np.zeros((dim,dim,3),dtype=np.uint8)
-    a=cv2.resize(img1,(dim,dim)); b=cv2.resize(img2,(dim,dim))
-    l1=get_landmarks(a,step); l2=get_landmarks(b,step)
-    if l1 is None or l2 is None or l1.shape!=l2.shape:
-        return cv2.addWeighted(a,1-alpha,b,alpha,0)
-    m=(1-alpha)*l1+alpha*l2
-    tris=calculate_delaunay_triangles((0,0,dim,dim),m)
-    if not tris:
-        return cv2.addWeighted(a,1-alpha,b,alpha,0)
-    A=a.astype(np.float32)/255; B=b.astype(np.float32)/255
-    Wa=np.zeros_like(A); Wb=np.zeros_like(B)
+    img1_cropped = img1[r1[1]:r1[1]+r1[3], r1[0]:r1[0]+r1[2]]
+
+    # Get affine transform
+    warp_mat = cv2.getAffineTransform(t1_rect, t2_rect)
+    img2_cropped = cv2.warpAffine(img1_cropped, warp_mat, (r2[2], r2[3]), None,
+                                  flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT_101)
+
+    # Create mask for blending
+    mask = np.zeros((r2[3], r2[2]), dtype=np.uint8)
+    cv2.fillConvexPoly(mask, t2_rect.astype(np.int32), 255)
+
+    # Blend the warped triangle
+    img2_rect = img2[r2[1]:r2[1]+r2[3], r2[0]:r2[0]+r2[2]]
+    img2_rect[mask > 0] = img2_cropped[mask > 0]
+
+
+def morph_faces(img1, img2, l1, l2, alpha, dim):
+    """
+    Morphs two faces using pre-computed landmarks.
+    This function no longer performs landmark detection itself.
+    """
+    # Create intermediate landmarks
+    morphed_landmarks = (1 - alpha) * l1 + alpha * l2
+
+    # Triangulate the intermediate mesh
+    tris = calculate_delaunay_triangles((0, 0, dim, dim), morphed_landmarks)
+    if not tris: # If triangulation fails, just cross-fade
+        return cv2.addWeighted(img1, 1 - alpha, img2, alpha, 0)
+
+    morphed_img = np.zeros_like(img1, dtype=np.float32)
+    img1_f = img1.astype(np.float32)
+    img2_f = img2.astype(np.float32)
+
     for ids in tris:
-        warp_triangle(A,Wa,l1[ids],m[ids]); warp_triangle(B,Wb,l2[ids],m[ids])
-    out=(1-alpha)*Wa+alpha*Wb
-    return (out*255).astype(np.uint8)
+        # Get triangles from each set of landmarks
+        t1 = l1[ids].astype(np.float32)
+        t2 = l2[ids].astype(np.float32)
+        tm = morphed_landmarks[ids].astype(np.float32)
+        
+        # Warp both images to the intermediate mesh
+        warped1 = np.zeros_like(morphed_img)
+        warped2 = np.zeros_like(morphed_img)
+        warp_triangle(img1_f, warped1, t1, tm)
+        warp_triangle(img2_f, warped2, t2, tm)
+
+        # Blend the two warped images
+        morphed_triangle = (1 - alpha) * warped1 + alpha * warped2
+        
+        # Add the blended triangle to the final image
+        mask = np.zeros((dim, dim), dtype=np.uint8)
+        cv2.fillConvexPoly(mask, tm.astype(np.int32), 255)
+        morphed_img[mask > 0] = morphed_triangle[mask > 0]
+
+    return np.uint8(morphed_img)
 
+# --- Main Video Processing Function (Optimized) ---
 
 def process_video(video_path, ref_img, trans, res, step, feather, padding):
     cap = cv2.VideoCapture(video_path)
     fps = cap.get(cv2.CAP_PROP_FPS) or 24
-    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
 
-    # Prepare masked reference
+    # === 1. Process Reference Image ONCE before the loop ===
     ref_bgr = cv2.cvtColor(ref_img, cv2.COLOR_RGB2BGR)
-    mask_ref, ref_box = get_face_mask_box(ref_bgr, feather, padding)
-    if mask_ref is None:
-        return None, None, None, None
-    x_r, y_r, w_r, h_r = ref_box
-    ref_cut = ref_bgr[y_r:y_r+h_r, x_r:x_r+w_r]
-    mask_ref_norm = mask_ref.astype(np.float32)[..., None] / 255.0
-    ref_masked = (ref_cut.astype(np.float32) * mask_ref_norm).astype(np.uint8)
-    ref_morph = cv2.resize(ref_masked, (res, res))
-
-    # Output video setup
+    ref_results = face_mesh.process(ref_bgr)
+    ref_landmarks_full = get_landmarks_from_result(ref_results, ref_bgr.shape)
+    if ref_landmarks_full is None:
+        raise gr.Error("No face detected in the reference image. Please use a clear photo.")
+
+    _, ref_box = get_face_mask_box_from_landmarks(ref_landmarks_full, ref_bgr.shape, feather, padding)
+    xr, yr, wr, hr = ref_box
+    ref_cut = ref_bgr[yr:yr+hr, xr:xr+wr]
+    ref_morph_in = cv2.resize(ref_cut, (res, res))
+    
+    # Scale reference landmarks to the 'res x res' morphing space
+    ref_landmarks_scaled = (ref_landmarks_full - [xr, yr]) * [res/wr, res/hr]
+    ref_landmarks_scaled = ref_landmarks_scaled[::step] # Apply sub-sampling
+
+    # === 2. Setup Output Video ===
     w_o = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
     h_o = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
     tmp_vid = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4').name
     out_vid = cv2.VideoWriter(tmp_vid, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w_o, h_o))
-
+    
+    # Storage for first frame previews
     first_crop = None
-    first_mask = None
-    first_ref = None
     first_morphed = None
 
-    for i in range(total):
+    # === 3. Process Video Frame by Frame ===
+    for i in range(total_frames):
         ret, frame = cap.read()
         if not ret: break
-        mask_roi, box = get_face_mask_box(frame, feather, padding)
-        if mask_roi is None:
-            out_vid.write(frame)
+
+        # Timestamp is crucial for tracking in video mode
+        timestamp_ms = int(cap.get(cv2.CAP_PROP_POS_MSEC))
+        
+        # Process frame ONCE
+        results_src = face_mesh.process(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+        
+        if not results_src.multi_face_landmarks:
+            out_vid.write(frame) # Write original frame if no face is found
             continue
+
+        src_landmarks_full = get_landmarks_from_result(results_src, frame.shape)
+        mask_roi, box = get_face_mask_box_from_landmarks(src_landmarks_full, frame.shape, feather, padding)
         x, y, w, h = box
-        # Crop and resize original ROI
+
+        # Crop and resize source face
         crop = frame[y:y+h, x:x+w]
         crop_resized = cv2.resize(crop, (res, res))
-        # Morph
-        alpha = float(np.clip((trans+1)/2, 0, 1))
-        mor = morph_faces(crop_resized, ref_morph, alpha, res, step)
-        # Store first
+
+        # Scale source landmarks to the 'res x res' morphing space
+        src_landmarks_scaled = (src_landmarks_full - [x, y]) * [res/w, res/h]
+        src_landmarks_scaled = src_landmarks_scaled[::step] # Apply sub-sampling
+        
+        # Perform morph
+        alpha = float(np.clip((trans + 1) / 2, 0, 1))
+        morphed_face = morph_faces(crop_resized, ref_morph_in, src_landmarks_scaled, ref_landmarks_scaled, alpha, res)
+
+        # Store first frame for preview
         if i == 0:
             first_crop = crop_resized.copy()
-            first_ref = ref_morph.copy()
-            first_mask = cv2.resize(mask_roi, (res, res), interpolation=cv2.INTER_LINEAR)
-            first_morphed = mor.copy()
-        # Resize morphed back
-        mor_back = cv2.resize(mor, (w, h))
-        # Composite with shape mask
-        mask_n = (mask_roi.astype(np.float32)[..., None] / 255.0)
-        region = frame[y:y+h, x:x+w].astype(np.float32)
-        blended = region * (1-mask_n) + mor_back.astype(np.float32) * mask_n
+            first_morphed = morphed_face.copy()
+
+        # Resize morphed face back and blend it onto the original frame
+        morphed_face_resized = cv2.resize(morphed_face, (w, h))
+        
+        # Create a feathered mask for seamless blending
+        mask_norm = mask_roi.astype(np.float32) / 255.0
+        mask_expanded = mask_norm[..., None] # Add channel dimension for broadcasting
+
+        # Composite the morphed face
+        region = frame[y:y+h, x:x+w]
+        blended = region * (1 - mask_expanded) + morphed_face_resized * mask_expanded
         frame[y:y+h, x:x+w] = blended.astype(np.uint8)
+        
         out_vid.write(frame)
 
-    cap.release(); out_vid.release()
-
-    # Apply mask to first_morphed for preview
-    if first_morphed is not None and first_mask is not None:
-        mask_n0 = first_mask.astype(np.float32)[..., None] / 255.0
-        first_morphed = (first_morphed.astype(np.float32) * mask_n0).astype(np.uint8)
-    else:
-        first_morphed = np.zeros((res, res, 3), dtype=np.uint8)
-        first_crop = first_crop if first_crop is not None else np.zeros((res, res,3),np.uint8)
-        first_ref = first_ref if first_ref is not None else ref_morph.copy()
+    cap.release()
+    out_vid.release()
 
-    # Convert for Gradio
-    return tmp_vid, cv2.cvtColor(first_crop, cv2.COLOR_BGR2RGB), cv2.cvtColor(first_ref, cv2.COLOR_BGR2RGB), cv2.cvtColor(first_morphed, cv2.COLOR_BGR2RGB)
+    # Convert preview images for Gradio output
+    first_crop_rgb = cv2.cvtColor(first_crop, cv2.COLOR_BGR2RGB) if first_crop is not None else np.zeros((res,res,3), np.uint8)
+    ref_morph_in_rgb = cv2.cvtColor(ref_morph_in, cv2.COLOR_BGR2RGB)
+    first_morphed_rgb = cv2.cvtColor(first_morphed, cv2.COLOR_BGR2RGB) if first_morphed is not None else np.zeros((res,res,3), np.uint8)
+    
+    return tmp_vid, first_crop_rgb, ref_morph_in_rgb, first_morphed_rgb
 
 # --- Gradio App ---
 css = """video, img { object-fit: contain !important; }"""
-with gr.Blocks(css=css) as iface:
-    gr.Markdown("# Morph with Face-Shaped Composite and Padding")
+with gr.Blocks(css=css, theme=gr.themes.Soft()) as iface:
+    gr.Markdown("# ✨ Optimized Face Morphing ✨")
+    gr.Markdown("This version uses MediaPipe's video tracking for a **faster and smoother** result. Jitter is reduced by maintaining landmark context between frames.")
     with gr.Row():
         vid = gr.Video(label='Input Video')
         ref = gr.Image(type='numpy', label='Reference Image')
     with gr.Row():
-        res = gr.Dropdown([256,384,512,768], value=512, label='Resolution')
-        step = gr.Slider(1,4,value=4,step=1,label='Landmark Sub-sampling')
-        feather = gr.Slider(0,50,value=10,step=1,label='Feather Radius')
-        padding = gr.Slider(0,100,value=24,step=1,label='Crop Padding (px)')
-    trans = gr.Slider(-1.0,1.0,value=-0.35,step=0.05,label='Transition Level')
-    btn = gr.Button('Generate Morph 🚀')
-    out_vid = gr.Video(label='Morphed Video')
-    out_crop = gr.Image(label='First Frame Crop')
-    out_ref = gr.Image(label='Masked Reference')
-    out_morph = gr.Image(label='Masked Morphed First Frame')
+        res = gr.Dropdown([256, 384, 512], value=384, label='Morph Resolution')
+        step = gr.Slider(1, 4, value=4, step=1, label='Landmark Sub-sampling', info="Higher value is faster but less detailed.")
+        feather = gr.Slider(0, 50, value=15, step=1, label='Feather Radius', info="Softens the blend edge.")
+        padding = gr.Slider(0, 100, value=25, step=1, label='Crop Padding (px)', info="Expands the face area.")
+    trans = gr.Slider(-1.0, 1.0, value=-0.3, step=0.05, label='Morph Transition', info="-1.0 is original face, 1.0 is reference face.")
+    
+    btn = gr.Button('Generate Morph 🚀', variant='primary')
+    
+    with gr.Row():
+        out_vid = gr.Video(label='Morphed Video')
+    with gr.Row():
+        out_crop = gr.Image(label='First Frame Crop')
+        out_ref = gr.Image(label='Reference Face')
+        out_morph = gr.Image(label='Morphed First Frame')
 
     btn.click(
         fn=process_video,
-        inputs=[vid,ref,trans,res,step,feather,padding],
-        outputs=[out_vid,out_crop,out_ref,out_morph],
-        show_progress=True
+        inputs=[vid, ref, trans, res, step, feather, padding],
+        outputs=[out_vid, out_crop, out_ref, out_morph],
+        show_progress='full'
     )
-    gr.Markdown("---\n*Added padding to the face crop for better framing.*")
 
-if __name__=='__main__':
+if __name__ == '__main__':
     iface.launch(debug=True)
-