Update app.py

app.py

@@ -107,30 +107,25 @@ def get_hair_and_exclude_masks(pil_image: Image.Image):
     hair = hair * (1 - face_m)
     hair = cv2.resize(hair, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR)
 
-    # ====================== STRONGER LIP PROTECTION (Exclude Mask) ======================
+    # ====================== EXCLUDE MASK ======================
     exclude = np.zeros((128, 128), dtype=np.float32)
-    
-    # Mouth, upper lip, lower lip, teeth ko strong exclude
-    exclude = np.maximum(exclude, (probs[10].numpy() > 0.22).astype(np.float32))  # mouth
-    exclude = np.maximum(exclude, (probs[11].numpy() > 0.25).astype(np.float32))  # upper lip
-    exclude = np.maximum(exclude, (probs[12].numpy() > 0.25).astype(np.float32))  # lower lip
-    exclude = np.maximum(exclude, (probs[4].numpy() > 0.28).astype(np.float32))   # teeth/nose area
-    exclude = np.maximum(exclude, (probs[5].numpy() > 0.28).astype(np.float32))
-
-    # Strong dilation for full coverage
-    kernel_big = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
-    exclude = cv2.dilate(exclude, kernel_big, iterations=3)
-    exclude = cv2.GaussianBlur(exclude, (9, 9), 1.8)
+    exclude = np.maximum(exclude, (probs[10].numpy() > 0.35).astype(np.float32))
+    exclude = np.maximum(exclude, (probs[11].numpy() > 0.35).astype(np.float32))
+    exclude = np.maximum(exclude, (probs[12].numpy() > 0.35).astype(np.float32))
+    exclude = np.maximum(exclude, (probs[4].numpy() > 0.35).astype(np.float32))
+    exclude = np.maximum(exclude, (probs[5].numpy() > 0.35).astype(np.float32))
+
+    exclude = cv2.dilate(exclude, kernel, iterations=2)
+    exclude = cv2.GaussianBlur(exclude, (5, 5), 1.2)
     exclude = cv2.resize(exclude, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR)
 
     # ====================== MUSTACHE MASK ======================
-    mustache_upper = (probs[11].numpy() > 0.28).astype(np.float32)
-    mustache_lower = (probs[12].numpy() > 0.28).astype(np.float32)
+    mustache_upper = (probs[11].numpy() > 0.30).astype(np.float32)
+    mustache_lower = (probs[12].numpy() > 0.30).astype(np.float32)
     mustache = np.maximum(mustache_upper, mustache_lower)
    
-    mouth = (probs[10].numpy() > 0.28).astype(np.float32)
-    mustache = np.maximum(mustache, mouth * 0.50)
-    
+    mouth = (probs[10].numpy() > 0.30).astype(np.float32)
+    mustache = np.maximum(mustache, mouth * 0.6)
     kernel_must = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5,5))
     mustache = cv2.morphologyEx(mustache, cv2.MORPH_CLOSE, kernel_must, iterations=2)
     mustache = cv2.GaussianBlur(mustache, (5, 5), 1.2)
@@ -150,12 +145,12 @@ def get_beard_mask_fast(pil_image: Image.Image, exclude_mask: np.ndarray):
     results = model.predict(
         img_array,
         device=DEVICE.type,
-        conf=0.28,      # Higher confidence = less false beard
-        iou=0.50,
+        conf=0.20,
+        iou=0.45,
         imgsz=128,
         half=False,
         verbose=False,
-        max_det=6
+        max_det=8
     )
 
     mask = np.zeros((orig_h, orig_w), dtype=np.float32)
@@ -165,46 +160,44 @@ 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.42).astype(np.float32))   # stricter threshold
+                mask = np.maximum(mask, (m > 0.32).astype(np.float32))   # slightly lower threshold
 
-    # Strong lip protection
+    # Remove protected areas
     mask = np.maximum(mask - exclude_mask, 0)
 
-    if mask.sum() > 25:
-        # ====================== STABLE ROUND BEARD MASK ======================
+    if mask.sum() > 25:   # lowered a bit
+        # ====================== HIGHLY IMPROVED ROUNDING ======================
         
-        # Aggressive erosion to remove extra space
+        # 1. Aggressive erosion to kill extra space
         kernel_erode = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
         mask = cv2.erode(mask, kernel_erode, iterations=2)
 
-        # Big close for perfect round corners
-        kernel_close = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 15))
+        # 2. Strong close with very big kernel for perfect round corners
+        kernel_close = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (13, 13))
         mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel_close, iterations=3)
 
-        # Clean small noise
+        # 3. Clean noise
         kernel_open = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
         mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel_open, iterations=1)
 
-        # Contour-based smoothing for natural round shape
-        contours, _ = cv2.findContours((mask > 0.12).astype(np.uint8), 
-                                       cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
-        
+        # 4. Contour based smoothing (Best for round & tight beard shape)
+        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) > 100:   # ignore tiny noisy parts
-                    epsilon = 0.004 * cv2.arcLength(cnt, True)  # smoother rounding
+                if cv2.contourArea(cnt) > 50:  # ignore very small noise
+                    epsilon = 0.008 * cv2.arcLength(cnt, True)   # lower = more round
                     approx = cv2.approxPolyDP(cnt, epsilon, True)
                     cv2.drawContours(smooth_mask, [approx], -1, 1.0, thickness=cv2.FILLED)
             mask = smooth_mask
 
-        # Soft final blur
+        # 5. Final soft blur
         mask = cv2.GaussianBlur(mask, (9, 9), 2.0)
 
-        # Tight final threshold
-        mask = (mask > 0.32).astype(np.float32)
+        # 6. Very tight final threshold (no extra space)
+        mask = (mask > 0.28).astype(np.float32)
 
-        # Final safety erosion
+        # 7. Final light erosion for extra safety
         mask = cv2.erode(mask, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3)), iterations=1)
 
     return mask
@@ -280,8 +273,8 @@ def process_face_whitening(input_image: Image.Image):
     hair_mask, exclude_mask, mustache_mask = get_hair_and_exclude_masks(img_resized)
     beard_mask = get_beard_mask_fast(img_resized, exclude_mask)
     
-    # Combine beard + mustache
-    beard_mask = np.maximum(beard_mask, mustache_mask * 0.85)
+    # Combine with mustache
+    beard_mask = np.maximum(beard_mask, mustache_mask * 0.88)
 
     final_resized = apply_strong_grey_hair(img_resized, hair_mask, beard_mask)
     final_img = final_resized.resize((ow, oh), Image.LANCZOS)