Create app.py

app.py

@@ -0,0 +1,100 @@
+import gradio as gr
+import cv2
+import numpy as np
+from PIL import Image
+import mediapipe as mp
+import insightface
+import torch
+from scipy import stats
+
+# -------------------- Setup --------------------
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+mp_mesh = mp.solutions.face_mesh.FaceMesh(
+    static_image_mode=True,
+    max_num_faces=1,
+    refine_landmarks=True
+)
+
+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):
+    res = mp_mesh.process(cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR))
+    if not res.multi_face_landmarks:
+        return None
+    lm = res.multi_face_landmarks[0].landmark
+    return np.array([[p.x, p.y] for p in lm], dtype=np.float32)
+
+def align_face(src, tgt):
+    s = get_landmarks(src)
+    t = get_landmarks(tgt)
+    if s is None or t is None:
+        return None
+    idx = [33, 263, 1, 61, 291]
+    M, _ = cv2.estimateAffinePartial2D(s[idx], t[idx])
+    return cv_to_pil(cv2.warpAffine(pil_to_cv(src), M, tgt.size))
+
+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, v2 = ea[0].embedding, 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):
+        src_lab[:,:,i] = stats.rankdata(src_lab[:,:,i]).reshape(src_lab[:,:,i].shape)
+        src_lab[:,:,i] = np.interp(src_lab[:,:,i],
+                                   (src_lab[:,:,i].min(), src_lab[:,:,i].max()),
+                                   (tgt_lab[:,:,i].min(), tgt_lab[:,:,i].max()))
+    return cv2.cvtColor(src_lab.astype(np.uint8), cv2.COLOR_LAB2BGR)
+
+# -------------------- Core --------------------
+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)
+
+    src_cv = color_match(src_cv, tgt_cv)
+
+    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)
+
+    sim = identity_similarity(src_img, result)
+    if sim < 0.94:
+        return f"Similarity too low: {sim:.3f}", result
+
+    return f"Similarity OK: {sim:.3f}", result
+
+# -------------------- UI --------------------
+with gr.Blocks() as demo:
+    gr.Markdown("## Ultra-Realistic Face Swap (Photographic Only)")
+    s = gr.Image(label="Source Face", type="pil")
+    t = gr.Image(label="Target Image", type="pil")
+    b = gr.Button("Swap")
+    out_txt = gr.Textbox(label="Status")
+    out_img = gr.Image(label="Result")
+
+    b.click(face_swap, [s, t], [out_txt, out_img])
+
+demo.launch()