Update app.py

app.py

@@ -18,19 +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}")
 
@@ -41,7 +40,7 @@ executor = ThreadPoolExecutor(max_workers=2)
 
 face_processor = None
 face_parser = None
-beard_model = None
+beard_model = None 
 
 # ================================================
 # 3. LOAD MODELS
@@ -50,6 +49,7 @@ def load_face_parser():
     global face_processor, face_parser
     if face_parser is not None:
         return face_processor, face_parser
+    
     logger.info("Loading Segformer face-parsing...")
     face_processor = SegformerImageProcessor.from_pretrained("jonathandinu/face-parsing")
     face_parser = SegformerForSemanticSegmentation.from_pretrained("jonathandinu/face-parsing")
@@ -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,20 +67,18 @@ 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
+# 4. MASK FUNCTIONS (Improved Hair Mask)
 # ================================================
-
 def get_beard_mask_fast(pil_image: Image.Image) -> np.ndarray:
     temp = f"temp_{uuid.uuid4().hex[:8]}.jpg"
     try:
         img_small = pil_image.resize((256, 256), Image.LANCZOS)
         img_small.save(temp)
         model = load_beard_model()
-        res = model(temp, device=DEVICE.type, conf=0.3, iou=0.5, 
-                    verbose=False, half=USE_FP16, imgsz=256)
+        res = model(temp, device=DEVICE.type, conf=0.3, iou=0.5, verbose=False, half=USE_FP16, imgsz=256)
         
         h, w = np.array(pil_image).shape[:2]
         mask = np.zeros((h, w), dtype=np.float32)
@@ -104,6 +102,7 @@ def get_beard_mask_fast(pil_image: Image.Image) -> np.ndarray:
             os.remove(temp)
 
 def get_hair_mask_fast(pil_image: Image.Image) -> np.ndarray:
+    """Improved version - better captures forehead & 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)
@@ -113,22 +112,36 @@ def get_hair_mask_fast(pil_image: Image.Image) -> np.ndarray:
     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)
+    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]
-    hair = (probs[13].cpu().numpy() > 0.12).astype(np.float32)
+
+    # ================== 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)
+
+    # Face mask - exclude inner face but keep more forehead area
     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)
-    face_m = cv2.dilate(face_m, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5)), iterations=1)
-    hair = hair * (1 - face_m)
     
+    # Reduced dilation so forehead hairs are not suppressed
+    face_m = cv2.dilate(face_m, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3)), iterations=1)
+
+    hair = hair * (1 - face_m)
+
+    # Morphology - better connection for scattered 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=1)
-    hair = cv2.GaussianBlur(hair, (3, 3), 1)
-    
+    hair = cv2.morphologyEx(hair, cv2.MORPH_CLOSE, kernel, iterations=2)   # Increased for better thin hair connection
+
+    hair = cv2.GaussianBlur(hair, (5, 5), 1.5)   # Smoother blend
+
     oh, ow = np.array(pil_image).shape[:2]
     return np.clip(cv2.resize(hair, (ow, oh)), 0, 1)
 
@@ -138,59 +151,49 @@ def get_masks_sequential(image):
 # ================================================
 # 5. 🔥 STRONG GREY HAIR - INCREASED INTENSITY
 # ================================================
-
 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
-    """
+    """ Strong grey effect - very visible but natural """
     comb = np.maximum(hair_mask, beard_mask)
-    
     if np.sum(comb) < 100:
         logger.warning("⚠️ Small mask area")
-    
+
     # Soften edges
     comb = cv2.GaussianBlur(comb, (7, 7), 2)
-    
+
     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%)
-    hsv[:,:,1] = hsv[:,:,1] * (1 - 1.0 * comb)  # Full desaturate where mask is
-    
+    hsv[:,:,1] = hsv[:,:,1] * (1 - 1.0 * comb)
+
     # 2. Increase brightness to grey level (STRONG boost)
-    # Old: +80, New: +110 for much lighter grey
     hsv[:,:,2] = hsv[:,:,2] + (110 * comb)
-    hsv[:,:,2] = np.clip(hsv[:,:,2], 120, 230)  # Higher minimum brightness
-    
+    hsv[:,:,2] = np.clip(hsv[:,:,2], 120, 230)
+
     # 3. NO hue change - prevents blue/green
-    
-    # Convert back
     grey_result = cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2RGB).astype(np.float32)
-    
+
     # ========================================
     # Blend - 85% new, 15% original (stronger)
     # ========================================
-    blend = 0.85 * comb  # Increased from 0.75 to 0.85
+    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)
     # ========================================
     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
     result = Image.fromarray(final)
     result = result.filter(ImageFilter.UnsharpMask(radius=0.5, percent=50, threshold=0))
-    
     return result
 
 def process_face_whitening(input_image: Image.Image) -> Image.Image:
@@ -199,24 +202,24 @@ 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)
-        
     except Exception as e:
         logger.error(f"❌ Processing error: {e}")
         logger.error(traceback.format_exc())
@@ -228,10 +231,11 @@ def process_face_whitening(input_image: Image.Image) -> Image.Image:
 app = FastAPI(title="Strong Grey Hair API")
 
 app.add_middleware(CORSMiddleware,
-                   allow_origins=["*"],
-                   allow_credentials=True,
-                   allow_methods=["*"],
-                   allow_headers=["*"])
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"]
+)
 
 @app.on_event("startup")
 async def startup_event():
@@ -245,11 +249,13 @@ async def startup_event():
 async def age_face(file: UploadFile = File(...)):
     if not file.content_type.startswith("image/"):
         raise HTTPException(400, "Only image files allowed")
+    
     contents = await file.read()
     try:
         input_image = Image.open(io.BytesIO(contents)).convert("RGB")
         loop = asyncio.get_event_loop()
         result = await loop.run_in_executor(executor, process_face_whitening, input_image)
+        
         buf = io.BytesIO()
         result.save(buf, format="JPEG", quality=92, optimize=True)
         buf.seek(0)