Update app.py

app.py

@@ -19,6 +19,7 @@ import io
 import asyncio
 from concurrent.futures import ThreadPoolExecutor
 import logging
+
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
@@ -26,12 +27,15 @@ logger = logging.getLogger(__name__)
 # 2. CONFIG
 # ================================================
 BEARD_MODEL_PATH = "models/best_hair_117_epoch_v4.pt"
-SAFE_IMG_SIZE = 768
+SAFE_IMG_SIZE = 640                    # Reduced from 768 → big speed gain
 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"
+
 logger.info(f"🚀 Device: {DEVICE}, FP16: {USE_FP16}")
+
 os.environ["HF_HOME"] = "/tmp/hf_cache"
 os.makedirs("/tmp/hf_cache", exist_ok=True)
+
 executor = ThreadPoolExecutor(max_workers=2)
 
 face_processor = None
@@ -45,17 +49,20 @@ 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...")
+
+    logger.info("Loading Segformer face-parsing (FP16)...")
     face_processor = SegformerImageProcessor.from_pretrained("jonathandinu/face-parsing")
     face_parser = SegformerForSemanticSegmentation.from_pretrained("jonathandinu/face-parsing")
+    
     face_parser.to(DEVICE)
     face_parser.eval()
     if USE_FP16:
         face_parser = face_parser.half()
+    
     logger.info("✅ Face parser loaded!")
     return face_processor, face_parser
 
+
 def load_beard_model():
     global beard_model
     if beard_model is None:
@@ -66,235 +73,243 @@ def load_beard_model():
     return beard_model
 
 # ================================================
-# 4. EXCLUDE MASK (NOSE + LIPS) - BEARD SE HATANA
+# 4. LIGHTWEIGHT NOSE + LIPS MASK (bahut important optimization)
 # ================================================
-def get_exclude_mask(pil_image: Image.Image, max_size=512) -> np.ndarray:
-    """Returns binary mask for nose (2) and lips (11,12) at original size."""
+def get_nose_lips_mask(pil_image: Image.Image) -> np.ndarray:
+    """Fast nose + lips exclusion mask"""
     processor, parser = load_face_parser()
     
     w, h = pil_image.size
-    if max(w, h) > max_size:
-        ratio = max_size / max(w, h)
-        new_size = (int(w * ratio), int(h * ratio))
-        img_resized = pil_image.resize(new_size, Image.LANCZOS)
-    else:
-        img_resized = pil_image
+    # Small size for speed
+    small = pil_image.resize((192, 192), Image.LANCZOS)
     
-    inputs = processor(images=img_resized, return_tensors="pt").to(DEVICE)
+    inputs = processor(images=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=small.size[::-1], 
+                                           mode="bilinear", align_corners=False)
+        probs = torch.softmax(up.float() if USE_FP16 else up, dim=1)[0]
+
+    # Nose (2), Upper lip (11), Lower lip (12)
+    nose = (probs[2] > 0.5).cpu().numpy().astype(np.float32)
+    lips = ((probs[11] > 0.45) + (probs[12] > 0.45)).cpu().numpy().astype(np.float32)
     
-    logits = out.logits
-    up = torch.nn.functional.interpolate(logits, size=img_resized.size[::-1], mode="bilinear", align_corners=False)
-    probs = torch.softmax(up.float() if USE_FP16 else up, dim=1)[0]
-    
-    # Nose (class 2), Upper lip (11), Lower lip (12)
-    nose = (probs[2].cpu().numpy() > 0.5).astype(np.float32)
-    lip_upper = (probs[11].cpu().numpy() > 0.5).astype(np.float32)
-    lip_lower = (probs[12].cpu().numpy() > 0.5).astype(np.float32)
-    exclude = np.clip(nose + lip_upper + lip_lower, 0, 1)
+    exclude = np.clip(nose + lips, 0, 1)
     
-    # Resize back to original dimensions
+    # Resize to original + light dilation
     exclude = cv2.resize(exclude, (w, h), interpolation=cv2.INTER_NEAREST)
-    # Strong dilation to remove half lips and nose edges completely
-    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
-    exclude = cv2.dilate(exclude, kernel, iterations=2)
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
+    exclude = cv2.dilate(exclude, kernel, iterations=1)
+    
     return exclude
 
 # ================================================
-# 5. MASK FUNCTIONS (BEARD EXCLUDES NOSE + LIPS)
+# 5. MASK FUNCTIONS
 # ================================================
 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.25, iou=0.45, 
+                   verbose=False, half=USE_FP16, imgsz=256)
+
         h, w = np.array(pil_image).shape[:2]
         mask = np.zeros((h, w), dtype=np.float32)
-       
+
         if res[0].masks is not None:
             for i, cls in enumerate(res[0].boxes.cls):
                 if int(cls) == 0:  # beard class
-                    m = cv2.resize(res[0].masks.data[i].cpu().numpy(), (w, h))
-                    mask = np.maximum(mask, (m > 0.45).astype(np.float32))
-       
-        # ========== REMOVE NOSE AND LIPS FROM BEARD MASK ==========
-        exclude_mask = get_exclude_mask(pil_image)   # nose + lips
-        mask = np.maximum(mask - exclude_mask, 0.0)
-       
+                    m = cv2.resize(res[0].masks.data[i].cpu().numpy(), (w, h), 
+                                 interpolation=cv2.INTER_LINEAR)
+                    mask = np.maximum(mask, (m > 0.40).astype(np.float32))
+
+        # Remove nose + lips
+        exclude = get_nose_lips_mask(pil_image)
+        mask = np.maximum(mask - exclude, 0.0)
+
         if np.sum(mask) == 0:
             return mask
-       
+
+        # Light morphology + blur
         kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
         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)
-       
+        mask = cv2.GaussianBlur(mask, (5, 5), 0.8)
+
         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:
-    """Fixed & Improved version - Now properly captures forehead, hairline & thin hairs"""
+    """Improved + faster hair mask"""
     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]
-    # ================== 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
+        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]
+
+    # Strong + soft hair capture (forehead & thin hairs)
+    strong_hair = (probs[13] > 0.055).astype(np.float32)
+    soft_hair = (probs[13] > 0.022).astype(np.float32)
     hair = np.maximum(strong_hair, soft_hair * 0.68)
-    # Face mask - exclude inner face but much softer on forehead & hairline
+
+    # Face exclusion (softer on forehead)
     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)
-    # 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)
+
+    # Forehead relaxation
     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
+    forehead_region[0:int(h * 0.32), :] = 1.0
+    face_m = face_m * (1 - forehead_region * 0.45)
+
     hair = hair * (1 - face_m)
-    # Morphology - better connection for scattered forehead hairs
+
+    # Morphology for better connection
     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=3) # Increased for better thin hair connection
-    # Smoother blend
-    hair = cv2.GaussianBlur(hair, (5, 5), 1.8)
-    oh, ow = np.array(pil_image).shape[:2]
+    hair = cv2.morphologyEx(hair, cv2.MORPH_CLOSE, kernel, iterations=2)
+
+    hair = cv2.GaussianBlur(hair, (5, 5), 1.5)
+
+    ow, oh = pil_image.size
     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)
+    hair = get_hair_mask_fast(image)
+    beard = get_beard_mask_fast(image)
+    return hair, beard
 
 # ================================================
-# 6. GREY HAIR & BEARD - EXACT COLOR TRANSFER
+# 6. GREY HAIR & BEARD COLOR TRANSFER (Optimized)
 # ================================================
 def apply_strong_grey_hair(image: Image.Image, hair_mask: np.ndarray, beard_mask: np.ndarray) -> Image.Image:
-    """Beard color exactly matches hair using LAB color transfer"""
     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)
+    if np.sum(comb) < 80:
+        logger.warning("⚠️ Very small mask area")
+
+    comb = cv2.GaussianBlur(comb, (7, 7), 2.0)
     img = np.array(image).astype(np.float32) / 255.0
-    
-    # Step 1: Apply grey transformation to hair area only
+
+    # HSV based grey for hair area (faster)
     hsv = cv2.cvtColor((img * 255).astype(np.uint8), cv2.COLOR_RGB2HSV).astype(np.float32)
     hsv_hair = hsv.copy()
-    # Desaturate and boost brightness only on hair
-  # Gentler grey transformation
-    saturation_factor = 0.8  # Less aggressive (was 1.0)
-    brightness_boost = 90     # Less boost (was 110)
-
+    
+    saturation_factor = 0.78
+    brightness_boost = 82
     hsv_hair[:,:,1] = hsv_hair[:,:,1] * (1 - saturation_factor * hair_mask)
-    hsv_hair[:,:,2] = hsv_hair[:,:,2] + (brightness_boost * hair_mask)
-    hsv_hair[:,:,2] = np.clip(hsv_hair[:,:,2], 100, 200) 
-    hair_grey = cv2.cvtColor(hsv_hair.astype(np.uint8), cv2.COLOR_HSV2RGB).astype(np.float32) / 255.0
+    hsv_hair[:,:,2] = np.clip(hsv_hair[:,:,2] + (brightness_boost * hair_mask), 95, 205)
     
-    # Step 2: Transfer color from hair_grey to beard area (preserving beard texture)
-    # Convert to LAB color space
+    hair_grey = cv2.cvtColor(hsv_hair.astype(np.uint8), cv2.COLOR_HSV2RGB).astype(np.float32) / 255.0
+
+    # Simple LAB color transfer for beard (to match hair tone)
     hair_lab = cv2.cvtColor((hair_grey * 255).astype(np.uint8), cv2.COLOR_RGB2LAB).astype(np.float32)
     img_lab = cv2.cvtColor((img * 255).astype(np.uint8), cv2.COLOR_RGB2LAB).astype(np.float32)
-    
-    # Get hair region pixels in LAB
+
     hair_mask_binary = (hair_mask > 0.5)
-    if np.sum(hair_mask_binary) > 100:
-        hair_lab_pixels = hair_lab[hair_mask_binary]
-        mean_hair_lab = np.mean(hair_lab_pixels, axis=0)
-        std_hair_lab = np.std(hair_lab_pixels, axis=0)
+    if np.sum(hair_mask_binary) > 80:
+        mean_hair_lab = np.mean(hair_lab[hair_mask_binary], axis=0)
+        std_hair_lab = np.std(hair_lab[hair_mask_binary], axis=0)
     else:
-        # fallback grey
         mean_hair_lab = np.array([128, 0, 0])
-        std_hair_lab = np.array([30, 10, 10])
-    
-    # Apply to beard region
+        std_hair_lab = np.array([30, 12, 12])
+
     beard_mask_binary = (beard_mask > 0.5)
     if np.sum(beard_mask_binary) > 0:
         beard_pixels_lab = img_lab[beard_mask_binary]
-        mean_beard_lab = np.mean(beard_pixels_lab, axis=0)
-        std_beard_lab = np.std(beard_pixels_lab, axis=0)
+        mean_beard = np.mean(beard_pixels_lab, axis=0)
+        std_beard = np.maximum(np.std(beard_pixels_lab, axis=0), 1e-5)
         
-        # Avoid division by zero
-        std_beard_lab = np.maximum(std_beard_lab, 1e-5)
-        # Transfer: (beard - mean_beard)/std_beard * std_hair + mean_hair
-        beard_norm = (beard_pixels_lab - mean_beard_lab) / std_beard_lab
+        beard_norm = (beard_pixels_lab - mean_beard) / std_beard
         beard_transfer = beard_norm * std_hair_lab + mean_hair_lab
         beard_transfer = np.clip(beard_transfer, 0, 255)
-        
-        # Put back
+
         img_lab_transfer = img_lab.copy()
         img_lab_transfer[beard_mask_binary] = beard_transfer
     else:
         img_lab_transfer = img_lab
-    
-    # Convert back to RGB
+
     final = cv2.cvtColor(img_lab_transfer.astype(np.uint8), cv2.COLOR_LAB2RGB).astype(np.float32) / 255.0
-    
-    # Where hair is, use hair_grey
+
+    # Apply hair grey
     hair_mask_3ch = np.stack([hair_mask, hair_mask, hair_mask], axis=2)
     final = hair_grey * hair_mask_3ch + final * (1 - hair_mask_3ch)
-    
-    # Blend edges
+
+    # Edge blend
     comb_3ch = np.stack([comb, comb, comb], axis=2)
     final = final * comb_3ch + img * (1 - comb_3ch)
-    
-    # Slight warm tint (optional)
+
+    # Slight warm tint
     warm = np.array([5, 3, 0], dtype=np.float32) / 255.0
-    final = final + (warm * comb[..., None] * 0.2)
+    final = final + (warm * 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.5, percent=50, threshold=0))
+    result = result.filter(ImageFilter.UnsharpMask(radius=0.6, percent=55, threshold=0))
     return result
 
+
 def process_face_whitening(input_image: Image.Image) -> Image.Image:
-    """Main processing function"""
     try:
         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...")
+
+        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 with color transfer...")
+
+        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())
         raise HTTPException(status_code=500, detail=f"Processing failed: {str(e)}")
 
+
 # ================================================
 # 7. FASTAPI
 # ================================================
 app = FastAPI(title="Strong Grey Hair API")
-app.add_middleware(CORSMiddleware,
-                   allow_origins=["*"],
-                   allow_credentials=True,
-                   allow_methods=["*"],
+
+app.add_middleware(CORSMiddleware, 
+                   allow_origins=["*"], 
+                   allow_credentials=True, 
+                   allow_methods=["*"], 
                    allow_headers=["*"])
 
 @app.on_event("startup")
@@ -309,13 +324,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)
@@ -326,7 +341,8 @@ async def age_face(file: UploadFile = File(...)):
     finally:
         gc.collect()
 
+
 if __name__ == "__main__":
     import uvicorn
     logger.info("🚀 Starting STRONG GREY HAIR server on port 7860")
-    uvicorn.run(app, host="0.0.0.0", port=7860) 
+    uvicorn.run(app, host="0.0.0.0", port=7860)