Update app.py

app.py

@@ -18,18 +18,18 @@ from fastapi.middleware.cors import CORSMiddleware
 import io
 import asyncio
 from concurrent.futures import ThreadPoolExecutor
-import logging 
+import logging
 
 logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__) 
+logger = logging.getLogger(__name__)
 
 # ================================================
 # 2. CONFIG
 # ================================================
 BEARD_MODEL_PATH = "models/best_hair_117_epoch_v4.pt"
-SAFE_IMG_SIZE = 768 
+SAFE_IMG_SIZE = 768
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-USE_FP16 = DEVICE.type == "cuda" and torch.cuda.is_available() 
+USE_FP16 = DEVICE.type == "cuda" and torch.cuda.is_available()
 
 logger.info(f"🚀 Device: {DEVICE}, FP16: {USE_FP16}")
 
@@ -40,7 +40,7 @@ executor = ThreadPoolExecutor(max_workers=2)
 
 face_processor = None
 face_parser = None
-beard_model = None 
+beard_model = None
 
 # ================================================
 # 3. LOAD MODELS
@@ -58,7 +58,7 @@ def load_face_parser():
     if USE_FP16:
         face_parser = face_parser.half()
     logger.info("✅ Face parser loaded!")
-    return face_processor, face_parser 
+    return face_processor, face_parser
 
 def load_beard_model():
     global beard_model
@@ -67,10 +67,10 @@ def load_beard_model():
             raise FileNotFoundError(f"❌ Beard model not found: {BEARD_MODEL_PATH}")
         logger.info("Loading Beard YOLO model...")
         beard_model = YOLO(BEARD_MODEL_PATH)
-    return beard_model 
+    return beard_model
 
 # ================================================
-# 4. MASK FUNCTIONS (Improved Hair Mask)
+# 4. MASK FUNCTIONS (Fixed Hair Mask)
 # ================================================
 def get_beard_mask_fast(pil_image: Image.Image) -> np.ndarray:
     temp = f"temp_{uuid.uuid4().hex[:8]}.jpg"
@@ -96,60 +96,66 @@ def get_beard_mask_fast(pil_image: Image.Image) -> np.ndarray:
         mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=1)
         mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=1)
         mask = cv2.GaussianBlur(mask, (5, 5), 1)
+        
         return np.clip(mask, 0, 1)
     finally:
         if os.path.exists(temp):
             os.remove(temp)
 
 def get_hair_mask_fast(pil_image: Image.Image) -> np.ndarray:
-    """Improved version - better captures forehead & thin hairs"""
+    """Fixed & Improved version - Now properly captures forehead, hairline & thin hairs"""
     processor, parser = load_face_parser()
     img_small = pil_image.resize((256, 256), Image.LANCZOS)
     inputs = processor(images=img_small, return_tensors="pt").to(DEVICE)
     if USE_FP16:
         inputs['pixel_values'] = inputs['pixel_values'].half()
-    
+
     with torch.no_grad():
         out = parser(**inputs)
     logits = out.logits
     up = torch.nn.functional.interpolate(logits, size=img_small.size[::-1], mode="bilinear", align_corners=False)
     probs = torch.softmax(up.float() if USE_FP16 else up, dim=1)[0]
 
-    # ================== IMPROVED HAIR MASK ==================
-    # Lower threshold to catch thin & forehead hairs
-    hair = (probs[13].cpu().numpy() > 0.07).astype(np.float32)
-    
-    # Extra soft hair capture for very fine hairs
-    soft_hair = (probs[13].cpu().numpy() > 0.03).astype(np.float32)
-    
-    # Combine strong + soft hairs
-    hair = np.maximum(hair, soft_hair * 0.55)
+    # ================== STRONGER HAIR CAPTURE FOR FOREHEAD & THIN HAIRS ==================
+    strong_hair = (probs[13].cpu().numpy() > 0.055).astype(np.float32)
+    soft_hair   = (probs[13].cpu().numpy() > 0.022).astype(np.float32)   # Very fine hairs
+
+    hair = np.maximum(strong_hair, soft_hair * 0.68)
 
-    # Face mask - exclude inner face but keep more forehead area
+    # Face mask - exclude inner face but much softer on forehead & hairline
     face_cls = list(range(1, 6)) + list(range(8, 13)) + [17, 18]
     parsing = up.argmax(dim=1).squeeze(0).cpu().numpy()
     face_m = np.isin(parsing, face_cls).astype(np.float32)
-    
-    # Reduced dilation so forehead hairs are not suppressed
-    face_m = cv2.dilate(face_m, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3)), iterations=1)
+
+    # Very light dilation so forehead hairs are not erased
+    kernel_face = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
+    face_m = cv2.dilate(face_m, kernel_face, iterations=1)
+
+    # Special relaxation for forehead area (top part of image)
+    h, w = face_m.shape
+    forehead_region = np.zeros_like(face_m)
+    forehead_region[0:int(h * 0.30), :] = 1.0          # Top 30% as forehead
+    face_m = face_m * (1 - forehead_region * 0.45)     # Reduce face suppression in forehead
 
     hair = hair * (1 - face_m)
 
-    # Morphology - better connection for scattered hairs
+    # Morphology - better connection for scattered forehead hairs
     kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
     hair = cv2.morphologyEx(hair, cv2.MORPH_OPEN, kernel, iterations=1)
-    hair = cv2.morphologyEx(hair, cv2.MORPH_CLOSE, kernel, iterations=2)   # Increased for better thin hair connection
+    hair = cv2.morphologyEx(hair, cv2.MORPH_CLOSE, kernel, iterations=3)   # Increased for better thin hair connection
 
-    hair = cv2.GaussianBlur(hair, (5, 5), 1.5)   # Smoother blend
+    # Smoother blend
+    hair = cv2.GaussianBlur(hair, (5, 5), 1.8)
 
     oh, ow = np.array(pil_image).shape[:2]
     return np.clip(cv2.resize(hair, (ow, oh)), 0, 1)
 
+
 def get_masks_sequential(image):
     return get_hair_mask_fast(image), get_beard_mask_fast(image)
 
 # ================================================
-# 5. 🔥 STRONG GREY HAIR - INCREASED INTENSITY
+# 5. STRONG GREY HAIR
 # ================================================
 def apply_strong_grey_hair(image: Image.Image, hair_mask: np.ndarray, beard_mask: np.ndarray) -> Image.Image:
     """ Strong grey effect - very visible but natural """
@@ -162,36 +168,28 @@ def apply_strong_grey_hair(image: Image.Image, hair_mask: np.ndarray, beard_mask
 
     img = np.array(image).astype(np.float32)
 
-    # ========================================
     # STRONG GREY CONVERSION
-    # ========================================
     hsv = cv2.cvtColor(img.astype(np.uint8), cv2.COLOR_RGB2HSV).astype(np.float32)
 
-    # 1. COMPLETE desaturation (100%)
+    # 1. COMPLETE desaturation
     hsv[:,:,1] = hsv[:,:,1] * (1 - 1.0 * comb)
-
-    # 2. Increase brightness to grey level (STRONG boost)
+    # 2. Increase brightness
     hsv[:,:,2] = hsv[:,:,2] + (110 * comb)
     hsv[:,:,2] = np.clip(hsv[:,:,2], 120, 230)
 
-    # 3. NO hue change - prevents blue/green
     grey_result = cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2RGB).astype(np.float32)
 
-    # ========================================
-    # Blend - 85% new, 15% original (stronger)
-    # ========================================
+    # Blend - 85% new, 15% original
     blend = 0.85 * comb
     mask_3ch = np.stack([blend, blend, blend], axis=2)
     final = grey_result * mask_3ch + img * (1 - mask_3ch)
 
-    # ========================================
-    # Very slight warm tint (barely noticeable)
-    # ========================================
+    # Very slight warm tint
     warm = np.array([5, 3, 0], dtype=np.float32)
     final = final + (warm * comb[..., None] * 0.2)
     final = np.clip(final, 0, 255).astype(np.uint8)
 
-    # Optional: Slight sharpening for definition
+    # Slight sharpening
     result = Image.fromarray(final)
     result = result.filter(ImageFilter.UnsharpMask(radius=0.5, percent=50, threshold=0))
     return result
@@ -202,21 +200,21 @@ def process_face_whitening(input_image: Image.Image) -> Image.Image:
         logger.info(f"→ Processing image: {input_image.size}")
         orig = input_image.convert("RGB")
         ow, oh = orig.size
-        
+
         target_size = min(SAFE_IMG_SIZE, max(ow, oh))
         if target_size % 2 == 1:
             target_size -= 1
-            
+
         img_resized = orig.resize((target_size, target_size), Image.LANCZOS)
-        
+
         logger.info(" Generating hair & beard masks...")
         hair_mask, beard_mask = get_masks_sequential(img_resized)
-        
+
         logger.info(f"Hair mask sum: {np.sum(hair_mask):.0f}, Beard mask sum: {np.sum(beard_mask):.0f}")
-        
+
         logger.info(" Applying STRONG GREY HAIR...")
         final_img = apply_strong_grey_hair(img_resized, hair_mask, beard_mask)
-        
+
         gc.collect()
         logger.info("✅ Processing completed!")
         return final_img.resize((ow, oh), Image.LANCZOS)
@@ -229,13 +227,11 @@ def process_face_whitening(input_image: Image.Image) -> Image.Image:
 # 6. FASTAPI
 # ================================================
 app = FastAPI(title="Strong Grey Hair API")
-
-app.add_middleware(CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"]
-)
+app.add_middleware(CORSMiddleware, 
+                   allow_origins=["*"], 
+                   allow_credentials=True, 
+                   allow_methods=["*"], 
+                   allow_headers=["*"])
 
 @app.on_event("startup")
 async def startup_event():