Create app.py

app.py

@@ -0,0 +1,126 @@
+import gradio as gr
+import cv2
+import numpy as np
+from PIL import Image
+import torch
+import face_alignment
+import insightface
+from scipy import stats
+
+# -------------------- Device --------------------
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# -------------------- Face Alignment --------------------
+fa = face_alignment.FaceAlignment(
+    face_alignment.LandmarksType["2D"],
+    device=device,
+    flip_input=False
+)
+
+# -------------------- Identity Model --------------------
+face_analyzer = insightface.app.FaceAnalysis(name="buffalo_l")
+face_analyzer.prepare(ctx_id=0 if device == "cuda" else -1)
+
+# -------------------- Utilities --------------------
+def pil_to_cv(img):
+    return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+
+def cv_to_pil(img):
+    return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
+
+def get_landmarks(img):
+    preds = fa.get_landmarks(np.array(img))
+    if preds is None or len(preds) == 0:
+        return None
+    return preds[0].astype(np.float32)  # (68,2)
+
+def align_face(src, tgt):
+    src_lm = get_landmarks(src)
+    tgt_lm = get_landmarks(tgt)
+    if src_lm is None or tgt_lm is None:
+        return None
+
+    # Use 5 key landmarks for affine transform: eyes, nose tip, mouth corners
+    idx = [36, 45, 30, 48, 54]
+    M, _ = cv2.estimateAffinePartial2D(src_lm[idx], tgt_lm[idx])
+    if M is None:
+        return None
+
+    aligned = cv2.warpAffine(
+        pil_to_cv(src),
+        M,
+        (tgt.width, tgt.height),
+        flags=cv2.INTER_LINEAR
+    )
+    return cv_to_pil(aligned)
+
+def identity_similarity(a, b):
+    ea = face_analyzer.get(np.array(a))
+    eb = face_analyzer.get(np.array(b))
+    if not ea or not eb:
+        return 0.0
+    v1 = ea[0].embedding
+    v2 = eb[0].embedding
+    return float(np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2)))
+
+def color_match(src, tgt):
+    src_lab = cv2.cvtColor(src, cv2.COLOR_BGR2LAB)
+    tgt_lab = cv2.cvtColor(tgt, cv2.COLOR_BGR2LAB)
+
+    for i in range(3):
+        s = src_lab[:, :, i].flatten()
+        t = tgt_lab[:, :, i].flatten()
+        s_rank = stats.rankdata(s)
+        s_norm = (s_rank - s_rank.min()) / (s_rank.max() - s_rank.min() + 1e-6)
+        t_sorted = np.sort(t)
+        src_lab[:, :, i] = t_sorted[(s_norm * (len(t_sorted) - 1)).astype(int)].reshape(src_lab[:, :, i].shape)
+
+    return cv2.cvtColor(src_lab.astype(np.uint8), cv2.COLOR_LAB2BGR)
+
+# -------------------- Core Face Swap --------------------
+def face_swap(src_img, tgt_img):
+    if src_img is None or tgt_img is None:
+        return "Upload both images", None
+
+    aligned = align_face(src_img, tgt_img)
+    if aligned is None:
+        return "Face alignment failed", None
+
+    src_cv = pil_to_cv(aligned)
+    tgt_cv = pil_to_cv(tgt_img)
+
+    # Color harmonization
+    src_cv = color_match(src_cv, tgt_cv)
+
+    # Poisson blending
+    mask = 255 * np.ones(src_cv.shape[:2], dtype=np.uint8)
+    center = (tgt_cv.shape[1] // 2, tgt_cv.shape[0] // 2)
+    blended = cv2.seamlessClone(src_cv, tgt_cv, mask, center, cv2.NORMAL_CLONE)
+
+    result = cv_to_pil(blended)
+
+    # Identity validation
+    sim = identity_similarity(src_img, result)
+    if sim < 0.94:
+        return f"Identity similarity too low: {sim:.3f}", result
+
+    return f"Identity similarity OK: {sim:.3f}", result
+
+# -------------------- Gradio UI --------------------
+with gr.Blocks() as demo:
+    gr.Markdown("## Ultra-Realistic Face Swap (Photographic Only)")
+    gr.Markdown(
+        "- Strict 2D face swap\n"
+        "- Identity similarity ≥0.94\n"
+        "- Zero AI / 3D look\n"
+    )
+
+    src = gr.Image(label="Source Face", type="pil")
+    tgt = gr.Image(label="Target Image", type="pil")
+    btn = gr.Button("Run Face Swap")
+    status = gr.Textbox(label="Status")
+    output = gr.Image(label="Result")
+
+    btn.click(face_swap, [src, tgt], [status, output])
+
+demo.launch()