Update app.py

app.py

@@ -73,13 +73,46 @@ def load_beard_model():
         beard_model = YOLO(BEARD_MODEL_PATH)
     return beard_model
 
-# ====================== MASKS ======================
+# ====================== IMPROVED MUSTACHE MASK ======================
+@timed("Mustache Mask")
+def get_mustache_mask(probs, orig_w, orig_h, exclude_mask):
+    # More sensitive for thin, light or sparse mustache
+    u_lip = (probs[11].numpy() > 0.16).astype(np.float32)
+    l_lip = (probs[12].numpy() > 0.16).astype(np.float32)
+    mouth = (probs[10].numpy() > 0.20).astype(np.float32)
+
+    mustache = np.maximum(u_lip, l_lip)
+    mustache = np.maximum(mustache, mouth * 0.55)
+
+    # Better morphology for mustache shape (horizontal bias)
+    kernel_small = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
+    kernel_horizontal = cv2.getStructuringElement(cv2.MORPH_RECT, (7, 3))
+
+    mustache = cv2.dilate(mustache, kernel_small, iterations=1)
+    mustache = cv2.morphologyEx(mustache, cv2.MORPH_CLOSE, kernel_horizontal, iterations=2)
+    mustache = cv2.GaussianBlur(mustache, (5, 5), 1.0)
+
+    # 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.65, 0)
+    mustache = (mustache > 0.20).astype(np.float32)
+
+    return mustache
+
+# ====================== HAIR + EXCLUDE + LIP MASK ======================
 @timed("Hair + Exclude + Lip Mask")
 def get_hair_and_exclude_masks(pil_image: Image.Image):
     load_face_parser()
     orig_w, orig_h = pil_image.size
-    img_small = pil_image.resize((128, 128), Image.BILINEAR)
 
+    img_small = pil_image.resize((128, 128), Image.BILINEAR)
     inputs = face_processor(images=img_small, return_tensors="pt").to(DEVICE)
 
     with torch.inference_mode():
@@ -92,11 +125,10 @@ def get_hair_and_exclude_masks(pil_image: Image.Image):
     hair = (probs[13].numpy() > 0.035).astype(np.float32)
     hair = cv2.GaussianBlur(hair, (3, 3), 1.0)
 
-    # Face mask
+    # Face mask (exclude forehead & other areas)
     parsing = up.argmax(dim=1).squeeze(0).cpu().numpy()
     face_cls = list(range(1,6)) + list(range(8,13)) + [17,18]
     face_m = np.isin(parsing, face_cls).astype(np.float32)
-
     kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3,3))
     face_m = cv2.dilate(face_m, kernel, iterations=1)
 
@@ -105,6 +137,7 @@ def get_hair_and_exclude_masks(pil_image: Image.Image):
     forehead[:int(h * 0.32)] = 1.0
     face_m = face_m * (1 - forehead * 0.45)
     hair = hair * (1 - face_m)
+
     hair = cv2.resize(hair, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR)
 
     # ====================== EXCLUDE MASK ======================
@@ -114,56 +147,37 @@ def get_hair_and_exclude_masks(pil_image: Image.Image):
     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 (IMPROVED + SHIFTED DOWN) ======================
-    mustache_upper = (probs[11].numpy() > 0.18).astype(np.float32)
-    mustache_lower = (probs[12].numpy() > 0.18).astype(np.float32)
-    mouth = (probs[10].numpy() > 0.18).astype(np.float32)
-
-    mustache = np.maximum(mustache_upper, mustache_lower)
-    mustache = np.maximum(mustache, mouth)
-
-    kernel_must = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
-    mustache = cv2.morphologyEx(mustache, cv2.MORPH_CLOSE, kernel_must, iterations=3)
-    mustache = cv2.dilate(mustache, kernel_must, iterations=2)
-    mustache = cv2.morphologyEx(mustache, cv2.MORPH_OPEN, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3,3)), iterations=1)
-    mustache = cv2.GaussianBlur(mustache, (5, 5), 1.2)
-
-    # --- Shift mustache mask downward by 3 pixels (at 128x128 scale) ---
-    shift_y = 3  # small downward shift
-    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)
-    mustache = np.maximum(mustache - exclude, 0)
-
-    # ====================== LIP MASK (HARD PROTECTION) ======================
+    # ====================== LIP MASK (Hard Protection) ======================
     lip_mask = np.zeros((128, 128), dtype=np.float32)
-    lip_mask = np.maximum(lip_mask, (probs[10].numpy() > 0.35).astype(np.float32))
-    lip_mask = np.maximum(lip_mask, (probs[11].numpy() > 0.35).astype(np.float32))
-    lip_mask = np.maximum(lip_mask, (probs[12].numpy() > 0.35).astype(np.float32))
-    lip_mask = cv2.dilate(lip_mask, kernel, iterations=2)
+    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))
+    lip_mask = np.maximum(lip_mask, (probs[12].numpy() > 0.42).astype(np.float32))
+    lip_mask = cv2.dilate(lip_mask, kernel, iterations=1)
     lip_mask = cv2.resize(lip_mask, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)
     lip_mask = (lip_mask > 0.5).astype(np.float32)
 
-    return hair, exclude, mustache, lip_mask
+    # Get improved mustache
+    mustache = get_mustache_mask(probs, orig_w, orig_h, exclude)
 
+    return hair, exclude, mustache, lip_mask
 
+# ====================== BEARD MASK ======================
 @timed("Beard Mask")
 def get_beard_mask_fast(pil_image: Image.Image, exclude_mask: np.ndarray, lip_mask: np.ndarray):
     model = load_beard_model()
     orig_w, orig_h = pil_image.size
+
     img_small = pil_image.resize((128, 128), Image.BILINEAR)
     img_array = np.array(img_small)
 
     results = model.predict(
         img_array,
         device=DEVICE.type,
-        conf=0.20,
+        conf=0.18,      # thoda lower for better mustache detection
         iou=0.45,
         imgsz=128,
         half=False,
@@ -172,15 +186,14 @@ def get_beard_mask_fast(pil_image: Image.Image, exclude_mask: np.ndarray, lip_ma
     )
 
     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:  # 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.32).astype(np.float32))
+                mask = np.maximum(mask, (m > 0.26).astype(np.float32))   # lower threshold
 
-    mask = np.maximum(mask - exclude_mask, 0)
+    mask = np.maximum(mask - exclude_mask * 0.6, 0)
 
     if mask.sum() > 25:
         kernel_erode = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
@@ -209,7 +222,6 @@ def get_beard_mask_fast(pil_image: Image.Image, exclude_mask: np.ndarray, lip_ma
     mask[lip_mask > 0] = 0
     return mask
 
-
 # ====================== COLOR TRANSFER ======================
 @timed("Color Transfer")
 def apply_strong_grey_hair(image: Image.Image, hair_mask: np.ndarray, beard_mask: np.ndarray):
@@ -259,27 +271,34 @@ 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)
     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)
 
+    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
 
-
 # ====================== MAIN PROCESSING ======================
 @timed("Total Processing")
 def process_face_whitening(input_image: Image.Image):
     orig = input_image.convert("RGB")
     ow, oh = orig.size
+
     target = min(SAFE_IMG_SIZE, max(ow, oh))
     if target % 2 != 0:
         target -= 1
+
     img_resized = orig.resize((target, target), Image.BILINEAR)
 
     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)
-    
-    beard_mask = np.maximum(beard_mask, mustache_mask * 0.88)
+
+    # Improved mustache blending
+    beard_mask = np.maximum(beard_mask, mustache_mask * 0.96)
+
+    # Extra boost for thin mustache areas
+    thin_mustache = (mustache_mask > 0.25) & (beard_mask < 0.45)
+    beard_mask[thin_mustache] = np.maximum(beard_mask[thin_mustache], 0.68)
+
     beard_mask[lip_mask > 0] = 0
 
     final_resized = apply_strong_grey_hair(img_resized, hair_mask, beard_mask)
@@ -291,7 +310,6 @@ def process_face_whitening(input_image: Image.Image):
 
     return final_img
 
-
 # ====================== FASTAPI ======================
 app = FastAPI()
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
@@ -307,7 +325,6 @@ async def startup():
     _ = process_face_whitening(dummy)
     logger.info("✅ Server Ready!")
 
-
 @app.post("/age-face")
 async def age_face(file: UploadFile = File(...)):
     start_total = time.perf_counter()
@@ -326,7 +343,6 @@ async def age_face(file: UploadFile = File(...)):
 
     return StreamingResponse(buf, media_type="image/jpeg")
 
-
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860, workers=1)