Update app.py

app.py

@@ -74,7 +74,7 @@ def load_beard_model():
     return beard_model
 
 # ====================== MASKS ======================
-@timed("Hair + Exclude Mask")
+@timed("Hair + Exclude + Lip Mask")
 def get_hair_and_exclude_masks(pil_image: Image.Image):
     load_face_parser()
     orig_w, orig_h = pil_image.size
@@ -132,11 +132,20 @@ def get_hair_and_exclude_masks(pil_image: Image.Image):
     mustache = cv2.resize(mustache, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR)
     mustache = np.maximum(mustache - exclude, 0)
 
-    return hair, exclude, mustache
+    # ====================== LIP MASK (HARD PROTECTION) ======================
+    lip_mask = np.zeros((128, 128), dtype=np.float32)
+    lip_mask = np.maximum(lip_mask, (probs[10].numpy() > 0.35).astype(np.float32))  # mouth
+    lip_mask = np.maximum(lip_mask, (probs[11].numpy() > 0.35).astype(np.float32))  # upper lip
+    lip_mask = np.maximum(lip_mask, (probs[12].numpy() > 0.35).astype(np.float32))  # lower lip
+    lip_mask = cv2.dilate(lip_mask, kernel, iterations=2)  # expand slightly to cover edges
+    lip_mask = cv2.resize(lip_mask, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)  # keep binary
+    lip_mask = (lip_mask > 0.5).astype(np.float32)
+
+    return hair, exclude, mustache, lip_mask
 
 
 @timed("Beard Mask")
-def get_beard_mask_fast(pil_image: Image.Image, exclude_mask: np.ndarray):
+def get_beard_mask_fast(pil_image: Image.Image, exclude_mask: np.ndarray, lip_mask: np.ndarray):
     model = load_beard_model()
     orig_w, orig_h = pil_image.size
     img_small = pil_image.resize((128, 128), Image.BILINEAR)
@@ -160,46 +169,43 @@ def get_beard_mask_fast(pil_image: Image.Image, exclude_mask: np.ndarray):
             if int(cls) == 0:  # beard class
                 m = results[0].masks.data[i].cpu().numpy()
                 m = cv2.resize(m, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR)
-                mask = np.maximum(mask, (m > 0.32).astype(np.float32))   # slightly lower threshold
+                mask = np.maximum(mask, (m > 0.32).astype(np.float32))
 
     # Remove protected areas
     mask = np.maximum(mask - exclude_mask, 0)
 
-    if mask.sum() > 25:   # lowered a bit
-        # ====================== HIGHLY IMPROVED ROUNDING ======================
-        
-        # 1. Aggressive erosion to kill extra space
+    if mask.sum() > 25:
+        # Aggressive erosion
         kernel_erode = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
         mask = cv2.erode(mask, kernel_erode, iterations=2)
 
-        # 2. Strong close with very big kernel for perfect round corners
+        # Strong close
         kernel_close = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (13, 13))
         mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel_close, iterations=3)
 
-        # 3. Clean noise
+        # Clean noise
         kernel_open = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
         mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel_open, iterations=1)
 
-        # 4. Contour based smoothing (Best for round & tight beard shape)
+        # Contour smoothing
         contours, _ = cv2.findContours((mask > 0.1).astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
         if contours:
             smooth_mask = np.zeros_like(mask, dtype=np.float32)
             for cnt in contours:
-                if cv2.contourArea(cnt) > 50:  # ignore very small noise
-                    epsilon = 0.008 * cv2.arcLength(cnt, True)   # lower = more round
+                if cv2.contourArea(cnt) > 50:
+                    epsilon = 0.008 * cv2.arcLength(cnt, True)
                     approx = cv2.approxPolyDP(cnt, epsilon, True)
                     cv2.drawContours(smooth_mask, [approx], -1, 1.0, thickness=cv2.FILLED)
             mask = smooth_mask
 
-        # 5. Final soft blur
         mask = cv2.GaussianBlur(mask, (9, 9), 2.0)
-
-        # 6. Very tight final threshold (no extra space)
         mask = (mask > 0.28).astype(np.float32)
-
-        # 7. Final light erosion for extra safety
         mask = cv2.erode(mask, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3)), iterations=1)
 
+    # ====================== FINAL LIP PROTECTION ======================
+    # Force lips to zero after all processing
+    mask[lip_mask > 0] = 0
+
     return mask
 
 
@@ -270,11 +276,13 @@ def process_face_whitening(input_image: Image.Image):
         target -= 1
     img_resized = orig.resize((target, target), Image.BILINEAR)
 
-    hair_mask, exclude_mask, mustache_mask = get_hair_and_exclude_masks(img_resized)
-    beard_mask = get_beard_mask_fast(img_resized, exclude_mask)
+    hair_mask, exclude_mask, mustache_mask, lip_mask = get_hair_and_exclude_masks(img_resized)
+    beard_mask = get_beard_mask_fast(img_resized, exclude_mask, lip_mask)
     
-    # Combine with mustache
+    # Combine with mustache, but lips already zeroed in beard_mask
     beard_mask = np.maximum(beard_mask, mustache_mask * 0.88)
+    # Reapply lip protection in case mustache added lip pixels
+    beard_mask[lip_mask > 0] = 0
 
     final_resized = apply_strong_grey_hair(img_resized, hair_mask, beard_mask)
     final_img = final_resized.resize((ow, oh), Image.LANCZOS)