Update app.py

app.py

@@ -1,34 +1,120 @@
 import gradio as gr
-from faceflux import process_video
+import cv2
+import numpy as np
+import mediapipe as mp
+import time
+import tempfile
+import os
+import faceflux
 
-css = """video { object-fit: contain !important; }"""
 
-with gr.Blocks(css=css) as iface:
-    gr.Markdown("# Super Fast Face Swap – FACEFLUX")
-    gr.Markdown(
-        "**FACEFLUX**: Ultra-lightweight, CPU-only face swap. "
-        "Ideal for small or distant faces; offline & privacy-preserving. "
-        "Weakness: large up-close or extreme angles."
-    )
+def process_video(
+    video_path, ref_img, trans, res, step, feather_pct, padding_pct,
+    progress=gr.Progress()
+):
+    # --- Initialization ---
+    cap = cv2.VideoCapture(video_path)
+    fps = cap.get(cv2.CAP_PROP_FPS) or 24
+    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    progress(0.0, desc="Initializing")
 
-    with gr.Row():
-        vid = gr.Video(label="Input Video")
-        ref = gr.Image(type="numpy", label="Reference Image")
+    # --- Prepare masked reference ---
+    ref_bgr = cv2.cvtColor(ref_img, cv2.COLOR_RGB2BGR)
+    mask_ref, ref_box = get_face_mask_box(ref_bgr, feather_pct, padding_pct)
+    if mask_ref is None:
+        progress(None)  # hide on error
+        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))
+    progress(0.1, desc="Reference ready")
 
-    with gr.Row():
-        res = gr.Dropdown([256,384,512,768], value=512, label="Resolution")
-        quality = gr.Slider(1,4, value=1, step=1, label="Swap Quality")
-        feather = gr.Slider(0.12, 0.24, value=0.12, step=0.01, label="Feather (%)")
-        strength = gr.Slider(-0.35, -0.15, value=-0.25, step=0.05, label="Strength")
+    # --- Output 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))
+
+    first_crop = first_ref = first_mask = first_morphed = None
+
+    # --- Frame-by-frame processing ---
+    for i in range(total):
+        ret, frame = cap.read()
+        if not ret:
+            break
+        progress(0.1 + 0.8 * (i / total), desc=f"Processing frame {i+1}/{total}")
+
+        mask_roi, box = get_face_mask_box(frame, feather_pct, padding_pct)
+        if mask_roi is None:
+            out_vid.write(frame)
+            continue
+        x, y, w, h = box
+        crop = frame[y:y+h, x:x+w]
+        crop_resized = cv2.resize(crop, (res, res))
+        alpha = float(np.clip((trans+1)/2, 0, 1))
+        mor = morph_faces(crop_resized, ref_morph, alpha, res, step)
+
+        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()
 
-    btn = gr.Button("Generate Morph 🚀")
-    out_vid = gr.Video(label="Morphed Video")
+        mor_back = cv2.resize(mor, (w, h))
+        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()
+
+    # --- First-frame outputs ---
+    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:
+        zero = np.zeros((res, res, 3), dtype=np.uint8)
+        first_crop = first_crop or zero
+        first_ref = first_ref or ref_morph
+        first_morphed = zero
+
+    progress(1.0, desc="Done")
+    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) as iface:
+    gr.Markdown("# Morph with Face-Shaped Composite and Padding Percentage")
+    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.0,0.5,value=0.1,step=0.01,label='Feather (%)')
+        padding = gr.Slider(0.0,0.5,value=0.24,step=0.01,label='Padding (%)')
+    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, strength, res, quality, feather, 0.24],
-        outputs=[out_vid],
+        inputs=[vid, ref, trans, res, step, feather, padding],
+        outputs=[out_vid, out_crop, out_ref, out_morph],
         show_progress=True
     )
 
-    iface.queue().launch(debug=True)
+    gr.Markdown("---\n*Default values set and feather/padding are now relative percentages.*")
+
+# Enable queueing so progress updates render
+iface.queue().launch(debug=True)