Update app.py

app.py

@@ -159,7 +159,7 @@ def get_hair_and_exclude_masks(pil_image: Image.Image):
 
     return hair, exclude, mustache, lip_mask
 
-# ====================== BEARD MASK ======================
+# ====================== BEARD MASK (FIXED: returns beard_present flag) ======================
 @timed("Beard Mask")
 def get_beard_mask_fast(pil_image: Image.Image, exclude_mask: np.ndarray, lip_mask: np.ndarray):
     model = load_beard_model()
@@ -180,16 +180,22 @@ def get_beard_mask_fast(pil_image: Image.Image, exclude_mask: np.ndarray, lip_ma
     )
 
     mask = np.zeros((orig_h, orig_w), dtype=np.float32)
+    beard_present = False  # <-- NEW FLAG
+
     if results[0].masks is not None:
         for i, cls in enumerate(results[0].boxes.cls):
             if int(cls) == 0:
-                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.25).astype(np.float32))
+                conf = results[0].boxes.conf[i].item()
+                if conf > 0.25:   # confidence threshold for considering a real beard
+                    beard_present = True
+                    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.25).astype(np.float32))
 
     mask = np.maximum(mask - exclude_mask * 0.6, 0)
 
-    if mask.sum() > 25:
+    # Only apply morphological refinements if beard is actually present
+    if beard_present and mask.sum() > 25:
         kernel_erode = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
         mask = cv2.erode(mask, kernel_erode, iterations=2)
         kernel_close = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (13, 13))
@@ -212,9 +218,9 @@ def get_beard_mask_fast(pil_image: Image.Image, exclude_mask: np.ndarray, lip_ma
         mask = cv2.erode(mask, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3)), iterations=1)
 
     mask[lip_mask > 0] = 0
-    return mask
+    return mask, beard_present   # <-- RETURN BOTH
 
-# ====================== COLOR TRANSFER - BEARD SAME AS HAIR (FIXED) ======================
+# ====================== COLOR TRANSFER - BEARD SAME AS HAIR ======================
 @timed("Color Transfer")
 def apply_strong_grey_hair(image: Image.Image, hair_mask: np.ndarray, beard_mask: np.ndarray):
     # Combine hair and beard masks
@@ -225,11 +231,8 @@ def apply_strong_grey_hair(image: Image.Image, hair_mask: np.ndarray, beard_mask
     img = np.array(image).astype(np.float32) / 255.0
     hsv = cv2.cvtColor((img * 255).astype(np.uint8), cv2.COLOR_RGB2HSV).astype(np.float32)
 
-    # Apply hair-style transformation using combined mask (hair + beard)
     hsv_transformed = hsv.copy()
-    # Desaturate based on combined mask
     hsv_transformed[..., 1] = hsv_transformed[..., 1] * (1 - 0.78 * combined_mask)
-    # Boost brightness
     original_v = hsv[..., 2]
     boost_amount = 89 * combined_mask
     hsv_transformed[..., 2] = np.clip(
@@ -238,11 +241,9 @@ def apply_strong_grey_hair(image: Image.Image, hair_mask: np.ndarray, beard_mask
     )
     transformed_rgb = cv2.cvtColor(hsv_transformed.astype(np.uint8), cv2.COLOR_HSV2RGB).astype(np.float32) / 255.0
 
-    # Final blend: only change combined mask area, keep rest original
     combined_mask_3ch = np.stack([combined_mask, combined_mask, combined_mask], axis=2)
     final = transformed_rgb * combined_mask_3ch + img * (1 - combined_mask_3ch)
 
-    # Light cool tint for natural look
     final = final + (np.array([9, 7, 5], dtype=np.float32) / 255.0 * combined_mask[..., None] * 0.18)
 
     final = np.clip(final * 255, 0, 255).astype(np.uint8)
@@ -251,7 +252,7 @@ def apply_strong_grey_hair(image: Image.Image, hair_mask: np.ndarray, beard_mask
 
     return result
 
-# ====================== MAIN PROCESSING ======================
+# ====================== MAIN PROCESSING (FIXED: mustache only if beard detected) ======================
 @timed("Total Processing")
 def process_face_whitening(input_image: Image.Image):
     orig = input_image.convert("RGB")
@@ -264,22 +265,15 @@ def process_face_whitening(input_image: Image.Image):
     img_resized = orig.resize((target, target), Image.BILINEAR)
 
     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)
-
-    # ========== NEW: Apply mustache only if beard is detected ==========
-    beard_detected = beard_mask.sum() > 100  # Threshold for beard presence
+    beard_mask, beard_present = get_beard_mask_fast(img_resized, exclude_mask, lip_mask)
 
-    if beard_detected:
-        # Blend mustache into beard mask
+    # ========== KEY FIX: Apply mustache ONLY if a beard is present ==========
+    if beard_present:
         beard_mask = np.maximum(beard_mask, mustache_mask * 0.98)
-
-        # Extra boost for thin mustache areas
         weak_mustache = (mustache_mask > 0.18) & (beard_mask < 0.48)
         beard_mask[weak_mustache] = np.maximum(beard_mask[weak_mustache], 0.75)
-
-        # Ensure lip region stays clear
         beard_mask[lip_mask > 0] = 0
-    # else: no beard → beard_mask remains as is (likely zero, mustache not added)
+    # else: no beard → mustache mask is ignored completely
 
     final_resized = apply_strong_grey_hair(img_resized, hair_mask, beard_mask)
     final_img = final_resized.resize((ow, oh), Image.LANCZOS)