Update app.py

app.py

@@ -146,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 (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))
@@ -214,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
 
-# ====================== FIXED COLOR TRANSFER (NO RED TINT + BEARD SAME AS HAIR) ======================
+# ====================== COLOR TRANSFER - BEARD SAME AS HAIR (FIXED) ======================
 @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)
@@ -222,47 +222,41 @@ def apply_strong_grey_hair(image: Image.Image, hair_mask: np.ndarray, beard_mask
         comb = cv2.GaussianBlur(comb, (5,5), 1.5)
 
     img = np.array(image).astype(np.float32) / 255.0
-    
-    # LAB for better color control
-    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 Grey Base
+    # === Hair Grey (Base) ===
     hsv_hair = hsv.copy()
-    hsv_hair[..., 1] = hsv_hair[..., 1] * (1 - 0.85 * hair_mask)   # strong desaturation
+    hsv_hair[..., 1] = hsv_hair[..., 1] * (1 - 0.78 * hair_mask)
     original_v = hsv[..., 2]
+    boost_amount = 89 * hair_mask
     hsv_hair[..., 2] = np.clip(
-        original_v + (85 * hair_mask) - (original_v * 0.30 * hair_mask),
-        105, 205
+        original_v + boost_amount - (original_v * 0.35 * hair_mask),
+        110, 210
     )
-    hair_grey_rgb = cv2.cvtColor(hsv_hair.astype(np.uint8), cv2.COLOR_HSV2RGB).astype(np.float32) / 255.0
-    hair_grey_lab = cv2.cvtColor((hair_grey_rgb * 255).astype(np.uint8), cv2.COLOR_RGB2LAB).astype(np.float32)
+    hair_grey = cv2.cvtColor(hsv_hair.astype(np.uint8), cv2.COLOR_HSV2RGB).astype(np.float32) / 255.0
 
-    # Apply to Beard + Mustache (Exactly same as hair)
-    beard_mask_3ch = np.stack([beard_mask, beard_mask, beard_mask], axis=2)
-    final_lab = lab.copy()
-    final_lab = final_lab * (1 - beard_mask_3ch) + hair_grey_lab * beard_mask_3ch
+    # === Create a fully grey image for beard (FIX) ===
+    gray = cv2.cvtColor((img * 255).astype(np.uint8), cv2.COLOR_RGB2GRAY)
+    gray = np.stack([gray, gray, gray], axis=2).astype(np.float32) / 255.0
 
-    # Remove Red Tint forcefully (a-channel = red-green)
-    beard_area = beard_mask > 0.12
-    if np.any(beard_area):
-        final_lab[beard_area, 1] = final_lab[beard_area, 1] * 0.20   # strong red neutralization
+    # === Beard + Mustache get the grey image (no more red!) ===
+    beard_mask_3ch = np.stack([beard_mask, beard_mask, beard_mask], axis=2)
+    final = gray * beard_mask_3ch + img * (1 - beard_mask_3ch)
 
-    # Apply Hair on top
+    # === Hair area gets the hair_grey (preserves texture) ===
     hair_mask_3ch = np.stack([hair_mask, hair_mask, hair_mask], axis=2)
-    final_lab = final_lab * (1 - hair_mask_3ch) + hair_grey_lab * hair_mask_3ch
+    final = hair_grey * hair_mask_3ch + final * (1 - hair_mask_3ch)
 
-    # Convert back to RGB
-    final = cv2.cvtColor(final_lab.astype(np.uint8), cv2.COLOR_LAB2RGB).astype(np.float32) / 255.0
-
-    # Final protection + cool tint
+    # === Final protection (avoid artifacts outside hair+beard) ===
     comb_3ch = np.stack([comb, comb, comb], axis=2)
     final = final * comb_3ch + img * (1 - comb_3ch)
-    final = final + (np.array([8, 6, 4], dtype=np.float32)/255.0 * comb[..., None] * 0.15)
+
+    # === Light cool tint for natural look ===
+    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.8, percent=80, threshold=1))
+    result = result.filter(ImageFilter.UnsharpMask(radius=0.8, percent=75, threshold=1))
 
     return result
 
@@ -281,7 +275,7 @@ 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 blending
+    # Strong blending for beard + mustache
     beard_mask = np.maximum(beard_mask, mustache_mask * 0.98)
 
     # Extra boost for thin mustache