Update app.py

app.py

@@ -6,288 +6,262 @@ import time
 import tempfile
 import os
 
-# --- MediaPipe Initialization for Video ---
-# Optimized for video by setting static_image_mode to False
-# This enables tracking and significantly improves speed and stability.
+# --- MediaPipe Initialization ---
 try:
     mp_face_mesh = mp.solutions.face_mesh
     face_mesh = mp_face_mesh.FaceMesh(
-        static_image_mode=False, # <-- The most important change for video
+        static_image_mode=True,
         max_num_faces=1,
         refine_landmarks=True,
-        min_detection_confidence=0.5,
-        min_tracking_confidence=0.5 # Confidence for tracking across frames
+        min_detection_confidence=0.5
     )
-    print("MediaPipe Face Mesh initialized for VIDEO successfully.")
+    print("MediaPipe Face Mesh initialized successfully.")
 except (ImportError, AttributeError):
     print("Error: Could not initialize MediaPipe Face Mesh. Is mediapipe installed correctly?")
     face_mesh = None
 
-# --- Helper Functions (Refactored for Efficiency) ---
+# --- Helper Functions ---
 
-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_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))
+def get_face_mask_box(img, feather, padding=0):
+    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 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_pad = min(x + bw + padding, w)
-    y2_pad = min(y + bh + padding, h)
-
-    # Feather the mask for smoother blending
-    if feather > 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)
-        if k % 2 == 0: k += 1 # Kernel size must be odd
-        mask = cv2.GaussianBlur(mask, (k, k), 0)
+        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)
+
+
+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):
-    """Calculates Delaunay triangles for a set of points."""
-    if points is None or len(points) < 3:
+    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)
-    # 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 = []
+    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]))
     
-    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)
-
-            if idx1 is not None and idx2 is not None and idx3 is not None:
-                delaunay_triangles.append([idx1, idx2, idx3])
-                
-    return delaunay_triangles
+    if r1[2] <= 0 or r1[3] <= 0 or r2[2] <= 0 or r2[3] <= 0:
+        return
 
+    img1_rect = img1[r1[1]:r1[1]+r1[3], r1[0]:r1[0]+r1[2]]
+    if img1_rect.size == 0: 
+        return
 
-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]
+    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)]
     
-    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)
-
+    mask=np.zeros((r2[3],r2[2],3),dtype=np.float32)
+    
+    cv2.fillConvexPoly(mask,np.int32(t2r),(1,1,1),16)
+    
+    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)
     for ids in tris:
-        # 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)
+        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)
 
-# --- 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_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
 
-    # === 1. Process Reference Image ONCE before the loop ===
+    # Prepare masked reference
     ref_bgr = cv2.cvtColor(ref_img, cv2.COLOR_RGB2BGR)
-    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 ===
+    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
     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
 
-    # === 3. Process Video Frame by Frame ===
-    for i in range(total_frames):
+    for i in range(total):
         ret, frame = cap.read()
         if not ret: break
-
-        # 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
+        mask_roi, box = get_face_mask_box(frame, feather, padding)
+        if mask_roi is None:
+            out_vid.write(frame)
             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 source face
+        # Crop and resize original ROI
         crop = frame[y:y+h, x:x+w]
         crop_resized = cv2.resize(crop, (res, res))
-
-        # 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
+        # Morph
+        alpha = float(np.clip((trans+1)/2, 0, 1))
+        mor = morph_faces(crop_resized, ref_morph, alpha, res, step)
+        # Store first
         if i == 0:
             first_crop = crop_resized.copy()
-            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
+            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
         frame[y:y+h, x:x+w] = blended.astype(np.uint8)
-        
         out_vid.write(frame)
 
-    cap.release()
-    out_vid.release()
+    cap.release(); out_vid.release()
 
-    # 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
+    # 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()
+
+    # 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)
 
 # --- Gradio App ---
 css = """video, img { object-fit: contain !important; }"""
-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.Blocks(css=css) as iface:
+    gr.Markdown("# Morph with Face-Shaped Composite and Padding")
     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], 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')
+        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')
 
     btn.click(
         fn=process_video,
-        inputs=[vid, ref, trans, res, step, feather, padding],
-        outputs=[out_vid, out_crop, out_ref, out_morph],
-        show_progress='full'
+        inputs=[vid,ref,trans,res,step,feather,padding],
+        outputs=[out_vid,out_crop,out_ref,out_morph],
+        show_progress=True
     )
+    gr.Markdown("---\n*Added padding to the face crop for better framing.*")
 
-if __name__ == '__main__':
-    iface.launch(debug=True)
+if __name__=='__main__':
+    iface.launch(debug=True)