Update app.py

app.py

@@ -106,23 +106,27 @@ def get_hair_and_exclude_masks(pil_image: Image.Image):
 
     hair = cv2.resize(hair, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR)
 
-    # ====================== IMPROVED EXCLUDE MASK (LIPS + EYES) ======================
+    # ====================== EXCLUDE MASK (LIPS + EYES) ======================
     exclude = np.zeros((128, 128), dtype=np.float32)
-    # Mouth / lips classes (common indices from CelebAMask-HQ via Segformer)
-    # 10 = mouth, 11 = upper lip, 12 = lower lip
-    exclude = np.maximum(exclude, (probs[10] > 0.4).cpu().numpy().astype(np.float32))
-    exclude = np.maximum(exclude, (probs[11] > 0.4).cpu().numpy().astype(np.float32))
-    exclude = np.maximum(exclude, (probs[12] > 0.4).cpu().numpy().astype(np.float32))
-    # Eyes protection (optional but safe)
+    # Mouth / lips classes
+    exclude = np.maximum(exclude, (probs[10] > 0.4).cpu().numpy().astype(np.float32))  # mouth
+    exclude = np.maximum(exclude, (probs[11] > 0.4).cpu().numpy().astype(np.float32))  # upper lip
+    exclude = np.maximum(exclude, (probs[12] > 0.4).cpu().numpy().astype(np.float32))  # lower lip
+    # Eyes protection
     exclude = np.maximum(exclude, (probs[4] > 0.4).cpu().numpy().astype(np.float32))   # left eye
     exclude = np.maximum(exclude, (probs[5] > 0.4).cpu().numpy().astype(np.float32))   # right eye
-    # Original nose (class 2) – keep as is (beard can overlap)
-    # But we already have exclude from lips/eyes only.
     
     exclude = cv2.resize(exclude, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)
     exclude = cv2.dilate(exclude, kernel, iterations=1)
 
-    return hair, exclude
+    # ====================== MUSTACHE MASK (from upper lip region) ======================
+    mustache = (probs[11].cpu().numpy() > 0.35).astype(np.float32)   # upper lip as mustache
+    mustache = cv2.resize(mustache, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR)
+    mustache = cv2.dilate(mustache, kernel, iterations=1)
+    # Remove overlap with exclude (lips) to avoid coloring lips
+    mustache = np.maximum(mustache - exclude, 0)
+
+    return hair, exclude, mustache
 
 @timed("Beard Mask")
 def get_beard_mask_fast(pil_image: Image.Image, exclude_mask: np.ndarray):
@@ -131,12 +135,11 @@ def get_beard_mask_fast(pil_image: Image.Image, exclude_mask: np.ndarray):
     img_small = pil_image.resize((128, 128), Image.BILINEAR)
     img_array = np.array(img_small)
 
-    # ====================== LOWER CONFIDENCE FOR MORE DETECTION ======================
     results = model.predict(
         img_array,
         device=DEVICE.type,
-        conf=0.20,          # was 0.28 – catch more beard
-        iou=0.45,           # was 0.50
+        conf=0.20,          # lower confidence for more beard
+        iou=0.45,
         imgsz=128,
         half=(DEVICE.type == "cuda"),
         verbose=False,
@@ -149,22 +152,20 @@ 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)
-                # ====================== LOWER MASK THRESHOLD ======================
-                mask = np.maximum(mask, (m > 0.35).astype(np.float32))  # was 0.42
+                mask = np.maximum(mask, (m > 0.35).astype(np.float32))
 
-    # Remove areas that should not change (lips, eyes)
+    # Remove lips/eyes
     mask = np.maximum(mask - exclude_mask, 0)
 
     if mask.sum() > 30:
         kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3,3))
-        # ====================== ADD DILATION TO COVER HAIRS ======================
-        mask = cv2.dilate(mask, kernel, iterations=1)          # NEW: expand beard coverage
+        mask = cv2.dilate(mask, kernel, iterations=1)
         mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=1)
         mask = cv2.GaussianBlur(mask, (3,3), 0.8)
 
     return mask
 
-# ====================== COLOR TRANSFER (Unchanged, Balanced) ======================
+# ====================== COLOR TRANSFER ======================
 @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)
@@ -175,13 +176,9 @@ def apply_strong_grey_hair(image: Image.Image, hair_mask: np.ndarray, beard_mask
     orig_lab = cv2.cvtColor((img * 255).astype(np.uint8), cv2.COLOR_RGB2LAB).astype(np.float32)
     hsv = cv2.cvtColor((img * 255).astype(np.uint8), cv2.COLOR_RGB2HSV).astype(np.float32)
 
-    # ====================== HAIR - Balanced Light Grey ======================
+    # Hair greying
     hsv_hair = hsv.copy()
-    
-    # Desaturation (grey effect)
     hsv_hair[..., 1] = hsv_hair[..., 1] * (1 - 0.78 * hair_mask)
-    
-    # Balanced Brightness
     original_v = hsv[..., 2]
     boost_amount = 89 * hair_mask
     hsv_hair[..., 2] = np.clip(
@@ -200,7 +197,7 @@ def apply_strong_grey_hair(image: Image.Image, hair_mask: np.ndarray, beard_mask
         mean_hair = np.array([130., 0., 0.])
         std_hair = np.array([30., 10., 10.])
 
-    # Beard Transfer
+    # Beard transfer
     beard_bin = beard_mask > 0.5
     if np.sum(beard_bin) > 30:
         beard_pix = orig_lab[beard_bin]
@@ -219,12 +216,9 @@ def apply_strong_grey_hair(image: Image.Image, hair_mask: np.ndarray, beard_mask
     final = hair_grey * hair_mask_3ch + final * (1 - hair_mask_3ch)
     final = final * comb_3ch + img * (1 - comb_3ch)
 
-    # Soft cool tint
     final = final + (np.array([9, 7, 5], dtype=np.float32)/255.0 * comb[..., None] * 0.18)
-
     final = np.clip(final * 255, 0, 255).astype(np.uint8)
     result = Image.fromarray(final)
-    
     result = result.filter(ImageFilter.UnsharpMask(radius=0.7, percent=70, threshold=2))
     
     return result
@@ -241,9 +235,12 @@ def process_face_whitening(input_image: Image.Image):
 
     img_resized = orig.resize((target, target), Image.BILINEAR)
 
-    hair_mask, exclude_mask = get_hair_and_exclude_masks(img_resized)
+    hair_mask, exclude_mask, mustache_mask = get_hair_and_exclude_masks(img_resized)
     beard_mask = get_beard_mask_fast(img_resized, exclude_mask)
 
+    # ====================== MERGE MUSTACHE INTO BEARD MASK ======================
+    beard_mask = np.maximum(beard_mask, mustache_mask)
+
     final_resized = apply_strong_grey_hair(img_resized, hair_mask, beard_mask)
 
     final_img = final_resized.resize((ow, oh), Image.LANCZOS)