Update app.py

app.py

@@ -73,10 +73,9 @@ def load_beard_model():
         beard_model = YOLO(BEARD_MODEL_PATH)
     return beard_model
 
-# ====================== IMPROVED MUSTACHE MASK ======================
+# ====================== MUSTACHE MASK ======================
 @timed("Mustache Mask")
 def get_mustache_mask(probs, orig_w, orig_h, exclude_mask):
-    # Sensitive detection for thin/sparse mustache
     u_lip = (probs[11].numpy() > 0.13).astype(np.float32)
     l_lip = (probs[12].numpy() > 0.13).astype(np.float32)
     mouth = (probs[10].numpy() > 0.18).astype(np.float32)
@@ -84,7 +83,6 @@ def get_mustache_mask(probs, orig_w, orig_h, exclude_mask):
     mustache = np.maximum(u_lip * 1.15, l_lip)
     mustache = np.maximum(mustache, mouth * 0.45)
 
-    # Morphology for nice mustache shape
     kernel_h = cv2.getStructuringElement(cv2.MORPH_RECT, (9, 3))
     kernel_e = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
 
@@ -92,15 +90,12 @@ def get_mustache_mask(probs, orig_w, orig_h, exclude_mask):
     mustache = cv2.morphologyEx(mustache, cv2.MORPH_CLOSE, kernel_h, iterations=2)
     mustache = cv2.GaussianBlur(mustache, (7, 5), 1.2)
 
-    # Small downward shift
     shift_y = 1
     M = np.float32([[1, 0, 0], [0, 1, shift_y]])
     mustache = cv2.warpAffine(mustache, M, (mustache.shape[1], mustache.shape[0]),
                               flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=0)
 
     mustache = cv2.resize(mustache, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR)
-
-    # Less aggressive exclude
     mustache = np.maximum(mustache - exclude_mask * 0.5, 0)
     mustache = cv2.GaussianBlur(mustache, (5, 5), 1.0)
     mustache = (mustache > 0.15).astype(np.float32)
@@ -151,7 +146,7 @@ def get_hair_and_exclude_masks(pil_image: Image.Image):
     exclude = cv2.GaussianBlur(exclude, (5, 5), 1.2)
     exclude = cv2.resize(exclude, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR)
 
-    # Lip mask (hard protection)
+    # Lip mask
     lip_mask = np.zeros((128, 128), dtype=np.float32)
     lip_mask = np.maximum(lip_mask, (probs[10].numpy() > 0.42).astype(np.float32))
     lip_mask = np.maximum(lip_mask, (probs[11].numpy() > 0.42).astype(np.float32))
@@ -160,7 +155,6 @@ def get_hair_and_exclude_masks(pil_image: Image.Image):
     lip_mask = cv2.resize(lip_mask, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)
     lip_mask = (lip_mask > 0.5).astype(np.float32)
 
-    # Mustache mask
     mustache = get_mustache_mask(probs, orig_w, orig_h, exclude)
 
     return hair, exclude, mustache, lip_mask
@@ -220,7 +214,7 @@ def get_beard_mask_fast(pil_image: Image.Image, exclude_mask: np.ndarray, lip_ma
     mask[lip_mask > 0] = 0
     return mask
 
-# ====================== STRONG COLOR TRANSFER (Mustache Fix) ======================
+# ====================== 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):
     comb = np.maximum(hair_mask, beard_mask)
@@ -230,7 +224,7 @@ 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)
 
-    # Hair Grey
+    # === Hair Grey (Base) ===
     hsv_hair = hsv.copy()
     hsv_hair[..., 1] = hsv_hair[..., 1] * (1 - 0.78 * hair_mask)
     original_v = hsv[..., 2]
@@ -241,18 +235,15 @@ def apply_strong_grey_hair(image: Image.Image, hair_mask: np.ndarray, beard_mask
     )
     hair_grey = cv2.cvtColor(hsv_hair.astype(np.uint8), cv2.COLOR_HSV2RGB).astype(np.float32) / 255.0
 
-    # Beard + Mustache Grey (Strong Application)
-    beard_grey = hair_grey.copy()
-
-    # Apply grey to beard + mustache
+    # === Beard + Mustache ko BILKUL SAME Hair Grey do ===
     beard_mask_3ch = np.stack([beard_mask, beard_mask, beard_mask], axis=2)
-    final = beard_grey * beard_mask_3ch + img * (1 - beard_mask_3ch)
+    final = hair_grey * beard_mask_3ch + img * (1 - beard_mask_3ch)
 
-    # Apply hair grey on top
+    # Hair area (on top)
     hair_mask_3ch = np.stack([hair_mask, hair_mask, hair_mask], axis=2)
     final = hair_grey * hair_mask_3ch + final * (1 - hair_mask_3ch)
 
-    # Final comb protection
+    # Final protection
     comb_3ch = np.stack([comb, comb, comb], axis=2)
     final = final * comb_3ch + img * (1 - comb_3ch)
 
@@ -280,10 +271,10 @@ def process_face_whitening(input_image: Image.Image):
     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)
 
-    # Strong mustache blending
+    # Strong blending for beard + mustache
     beard_mask = np.maximum(beard_mask, mustache_mask * 0.98)
 
-    # Extra boost for thin/weak mustache areas
+    # Extra boost for thin mustache
     weak_mustache = (mustache_mask > 0.18) & (beard_mask < 0.48)
     beard_mask[weak_mustache] = np.maximum(beard_mask[weak_mustache], 0.75)