Update app.py

app.py

@@ -106,10 +106,19 @@ 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 = np.zeros((128, 128), dtype=np.float32)
-    exclude = np.maximum(exclude, (probs[2] > 0.5).cpu().numpy().astype(np.float32))
-    exclude = np.maximum(exclude, (probs[11] > 0.5).cpu().numpy().astype(np.float32))
-    exclude = np.maximum(exclude, (probs[12] > 0.5).cpu().numpy().astype(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)
+    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)
 
@@ -122,11 +131,12 @@ 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.28,
-        iou=0.50,
+        conf=0.20,          # was 0.28 – catch more beard
+        iou=0.45,           # was 0.50
         imgsz=128,
         half=(DEVICE.type == "cuda"),
         verbose=False,
@@ -136,21 +146,25 @@ def get_beard_mask_fast(pil_image: Image.Image, exclude_mask: np.ndarray):
     mask = np.zeros((orig_h, orig_w), dtype=np.float32)
     if results[0].masks is not None:
         for i, cls in enumerate(results[0].boxes.cls):
-            if int(cls) == 0:
+            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))
+                # ====================== LOWER MASK THRESHOLD ======================
+                mask = np.maximum(mask, (m > 0.35).astype(np.float32))  # was 0.42
 
+    # Remove areas that should not change (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.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=1)
         mask = cv2.GaussianBlur(mask, (3,3), 0.8)
 
     return mask
 
-# ====================== UPDATED COLOR TRANSFER (Balanced) ======================
+# ====================== COLOR TRANSFER (Unchanged, Balanced) ======================
 @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)
@@ -167,12 +181,12 @@ def apply_strong_grey_hair(image: Image.Image, hair_mask: np.ndarray, beard_mask
     # Desaturation (grey effect)
     hsv_hair[..., 1] = hsv_hair[..., 1] * (1 - 0.78 * hair_mask)
     
-    # **Balanced Brightness** - Na zyada bright, na dark
+    # Balanced Brightness
     original_v = hsv[..., 2]
-    boost_amount = 89 * hair_mask                    # 145 → 105 (kam kiya)
+    boost_amount = 89 * hair_mask
     hsv_hair[..., 2] = np.clip(
         original_v + boost_amount - (original_v * 0.35 * hair_mask), 
-        110, 210                                  # upper limit kam kiya
+        110, 210
     )
 
     hair_grey = cv2.cvtColor(hsv_hair.astype(np.uint8), cv2.COLOR_HSV2RGB).astype(np.float32) / 255.0
@@ -205,7 +219,7 @@ 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 (natural grey look)
+    # 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)