Update app.py

app.py

@@ -4,90 +4,92 @@ import numpy as np
 from PIL import Image
 import mediapipe as mp
 
-# MediaPipe face mesh setup
+# Setup MediaPipe
 mp_face_mesh = mp.solutions.face_mesh
-face_mesh = mp_face_mesh.FaceMesh(static_image_mode=True)
 
-# Get facial landmarks
 def get_landmarks(image):
-    rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-    results = face_mesh.process(rgb)
-    h, w, _ = image.shape
-
-    if results.multi_face_landmarks:
-        landmarks = []
-        for pt in results.multi_face_landmarks[0].landmark:
-            x, y = int(pt.x * w), int(pt.y * h)
-            landmarks.append((x, y))
-        return np.array(landmarks, dtype=np.int32)
+    with mp_face_mesh.FaceMesh(static_image_mode=True) as face_mesh:
+        rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        results = face_mesh.process(rgb)
+        h, w, _ = image.shape
+
+        if results.multi_face_landmarks:
+            landmarks = []
+            for pt in results.multi_face_landmarks[0].landmark:
+                x, y = int(pt.x * w), int(pt.y * h)
+                landmarks.append((x, y))
+            return np.array(landmarks, dtype=np.int32)
     return None
 
-# Morph images based on landmarks and alpha
 def morph_images(img1, img2, alpha):
     lm1 = get_landmarks(img1)
     lm2 = get_landmarks(img2)
 
     if lm1 is None or lm2 is None:
-        return Image.fromarray(img1)  # fallback
+        return Image.fromarray(cv2.cvtColor(img1, cv2.COLOR_BGR2RGB))  # fallback
 
     lm_avg = ((1 - alpha) * lm1 + alpha * lm2).astype(np.int32)
 
-    # Triangulation
     rect = (0, 0, img1.shape[1], img1.shape[0])
     subdiv = cv2.Subdiv2D(rect)
     for p in lm_avg:
-        subdiv.insert((p[0], p[1]))
+        subdiv.insert((float(p[0]), float(p[1])))
     triangles = subdiv.getTriangleList().astype(np.int32)
 
-    def get_indices(tri, points):
+    def get_indices(tri_pts, ref_pts):
         idxs = []
-        for p in tri:
-            for i, pt in enumerate(points):
-                if np.linalg.norm(np.array(p) - np.array(pt)) < 1.0:
+        for tp in tri_pts:
+            for i, pt in enumerate(ref_pts):
+                if abs(tp[0] - pt[0]) < 2 and abs(tp[1] - pt[1]) < 2:
                     idxs.append(i)
+                    break
         return idxs if len(idxs) == 3 else None
 
-    morphed = np.zeros_like(img1)
+    morphed = np.zeros_like(img1, dtype=np.float32)
     for tri in triangles:
         pts = [(tri[0], tri[1]), (tri[2], tri[3]), (tri[4], tri[5])]
         idxs = get_indices(pts, lm_avg.tolist())
-        if idxs is None: continue
+        if idxs is None:
+            continue
 
         t1 = np.float32([lm1[i] for i in idxs])
         t2 = np.float32([lm2[i] for i in idxs])
         t = np.float32([lm_avg[i] for i in idxs])
 
         def warp_triangle(src, t_src, t_dst):
-            rect_src = cv2.boundingRect(t_src)
-            rect_dst = cv2.boundingRect(t_dst)
+            r_src = cv2.boundingRect(t_src)
+            r_dst = cv2.boundingRect(t_dst)
 
-            t_src_offset = np.array([[pt[0]-rect_src[0], pt[1]-rect_src[1]] for pt in t_src], np.float32)
-            t_dst_offset = np.array([[pt[0]-rect_dst[0], pt[1]-rect_dst[1]] for pt in t_dst], np.float32)
+            t_src_offset = np.array([[pt[0] - r_src[0], pt[1] - r_src[1]] for pt in t_src], np.float32)
+            t_dst_offset = np.array([[pt[0] - r_dst[0], pt[1] - r_dst[1]] for pt in t_dst], np.float32)
 
-            mask = np.zeros((rect_dst[3], rect_dst[2], 3), dtype=np.float32)
+            mask = np.zeros((r_dst[3], r_dst[2], 3), dtype=np.float32)
             cv2.fillConvexPoly(mask, np.int32(t_dst_offset), (1.0, 1.0, 1.0), 16, 0)
 
-            src_crop = src[rect_src[1]:rect_src[1]+rect_src[3], rect_src[0]:rect_src[0]+rect_src[2]]
+            src_crop = src[r_src[1]:r_src[1]+r_src[3], r_src[0]:r_src[0]+r_src[2]]
             warp_mat = cv2.getAffineTransform(t_src_offset, t_dst_offset)
-            dst_crop = cv2.warpAffine(src_crop, warp_mat, (rect_dst[2], rect_dst[3]), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT_101)
-            morphed[rect_dst[1]:rect_dst[1]+rect_dst[3], rect_dst[0]:rect_dst[0]+rect_dst[2]] *= (1 - mask)
-            morphed[rect_dst[1]:rect_dst[1]+rect_dst[3], rect_dst[0]:rect_dst[0]+rect_dst[2]] += dst_crop * mask
+            dst_crop = cv2.warpAffine(src_crop, warp_mat, (r_dst[2], r_dst[3]), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT_101)
+
+            morphed[r_dst[1]:r_dst[1]+r_dst[3], r_dst[0]:r_dst[0]+r_dst[2]] *= (1 - mask)
+            morphed[r_dst[1]:r_dst[1]+r_dst[3], r_dst[0]:r_dst[0]+r_dst[2]] += dst_crop * mask
 
         warp_triangle(img1, t1, t)
         warp_triangle(img2, t2, t)
 
-    return Image.fromarray(np.uint8(morphed))
+    return Image.fromarray(cv2.cvtColor(np.uint8(morphed), cv2.COLOR_BGR2RGB))
 
 def process(mm_image, aa_image, ee_image, oo_image, ww_image, na_image, slider_aa, slider_oo, slider_ee, slider_ww, slider_na):
-    def load(img): return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+    def to_bgr(img): return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+
+    base = to_bgr(mm_image)
+    morph_order = [(slider_aa, aa_image), (slider_oo, oo_image), (slider_ee, ee_image),
+                   (slider_ww, ww_image), (slider_na, na_image)]
 
-    img = load(mm_image)
-    for s, phoneme in [(slider_aa, aa_image), (slider_oo, oo_image), (slider_ee, ee_image), (slider_ww, ww_image), (slider_na, na_image)]:
-        if s > 0.0:
-            target = load(phoneme)
-            img = np.array(morph_images(img, target, s))
+    for strength, image in morph_order:
+        if strength > 0:
+            base = cv2.cvtColor(np.array(morph_images(base, to_bgr(image), strength)), cv2.COLOR_RGB2BGR)
 
-    return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
+    return Image.fromarray(cv2.cvtColor(base, cv2.COLOR_BGR2RGB))
 
 iface = gr.Interface(
     fn=process,