Update app.py

app.py

@@ -2,13 +2,12 @@ import os
 import torch
 import numpy as np
 import cv2
-import traceback
 import gc
 import time
 from PIL import Image, ImageFilter
 from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
 from ultralytics import YOLO
-from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi import FastAPI, File, UploadFile
 from fastapi.responses import StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
 import io
@@ -27,7 +26,7 @@ DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 logger.info(f"Using Device: {DEVICE}")
 logger.info(f"CUDA Available: {torch.cuda.is_available()}")
 
-# CPU Optimization (stable settings)
+# CPU Optimization
 if DEVICE.type == "cpu":
     torch.set_num_threads(4)
     torch.set_num_interop_threads(1)
@@ -81,32 +80,33 @@ 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)
-    
+   
     inputs = face_processor(images=img_small, return_tensors="pt").to(DEVICE)
-    
+   
     with torch.inference_mode():
         out = face_parser(**inputs)
         logits = out.logits
         up = torch.nn.functional.interpolate(logits, size=(128, 128), mode="bilinear", align_corners=False)
         probs = torch.softmax(up, dim=1)[0]
 
-    # Optimized hair mask
-    hair = (probs[13].cpu().numpy() > 0.04).astype(np.float32)
-    hair = cv2.GaussianBlur(hair, (3, 3), 0.8)
+    # Hair Mask
+    hair = (probs[13].cpu().numpy() > 0.035).astype(np.float32)   # thoda lower threshold
+    hair = cv2.GaussianBlur(hair, (3, 3), 1.0)
 
-    # Face exclude
+    # Face exclude mask
     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)
-    
+   
     h, w = face_m.shape
     forehead = np.zeros_like(face_m, dtype=np.float32)
     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 (nose + lips)
@@ -119,7 +119,6 @@ def get_hair_and_exclude_masks(pil_image: Image.Image):
 
     return hair, exclude
 
-
 @timed("Beard Mask")
 def get_beard_mask_fast(pil_image: Image.Image, exclude_mask: np.ndarray):
     model = load_beard_model()
@@ -130,7 +129,7 @@ def get_beard_mask_fast(pil_image: Image.Image, exclude_mask: np.ndarray):
     results = model.predict(
         img_array,
         device=DEVICE.type,
-        conf=0.30,
+        conf=0.28,
         iou=0.50,
         imgsz=128,
         half=(DEVICE.type == "cuda"),
@@ -141,10 +140,10 @@ def get_beard_mask_fast(pil_image: Image.Image, exclude_mask: np.ndarray):
     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:
+            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.45).astype(np.float32))
+                mask = np.maximum(mask, (m > 0.42).astype(np.float32))
 
     mask = np.maximum(mask - exclude_mask, 0)
 
@@ -155,7 +154,7 @@ def get_beard_mask_fast(pil_image: Image.Image, exclude_mask: np.ndarray):
 
     return mask
 
-
+# ====================== MAIN COLOR FUNCTION (Updated) ======================
 @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)
@@ -164,11 +163,21 @@ def apply_strong_grey_hair(image: Image.Image, hair_mask: np.ndarray, beard_mask
 
     img = np.array(image).astype(np.float32) / 255.0
     orig_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 - Special Light Grey Treatment ======================
     hsv_hair = hsv.copy()
-    hsv_hair[..., 1] = hsv_hair[..., 1] * (1 - 0.82 * hair_mask)
-    hsv_hair[..., 2] = np.clip(hsv_hair[..., 2] + (88 * hair_mask), 95, 205)
+    
+    # Strong desaturation for grey look
+    hsv_hair[..., 1] = hsv_hair[..., 1] * (1 - 0.78 * hair_mask)
+    
+    # **Main Fix**: Much stronger + adaptive brightness for camera images
+    original_v = hsv[..., 2]
+    boost_amount = 145 * hair_mask                                   # 88 → 145
+    hsv_hair[..., 2] = np.clip(
+        original_v + boost_amount - (original_v * 0.42 * hair_mask), 
+        125, 245
+    )
 
     hair_grey = cv2.cvtColor(hsv_hair.astype(np.uint8), cv2.COLOR_HSV2RGB).astype(np.float32) / 255.0
     hair_grey_lab = cv2.cvtColor((hair_grey * 255).astype(np.uint8), cv2.COLOR_RGB2LAB).astype(np.float32)
@@ -178,9 +187,10 @@ def apply_strong_grey_hair(image: Image.Image, hair_mask: np.ndarray, beard_mask
         mean_hair = np.mean(hair_grey_lab[hair_bin], axis=0)
         std_hair = np.maximum(np.std(hair_grey_lab[hair_bin], axis=0), 5.0)
     else:
-        mean_hair = np.array([128., 0., 0.])
-        std_hair = np.array([30., 10., 10.])
+        mean_hair = np.array([135., 0., 0.])
+        std_hair = np.array([32., 12., 12.])
 
+    # Beard Color Transfer (same as before)
     beard_bin = beard_mask > 0.5
     if np.sum(beard_bin) > 30:
         beard_pix = orig_lab[beard_bin]
@@ -198,14 +208,18 @@ 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([5,3,0], dtype=np.float32)/255.0 * comb[..., None] * 0.18)
-    final = np.clip(final * 255, 0, 255).astype(np.uint8)
 
+    # Extra cool white-grey tint (makes it look more natural silver-grey)
+    final = final + (np.array([12, 10, 8], dtype=np.float32)/255.0 * comb[..., None] * 0.25)
+
+    final = np.clip(final * 255, 0, 255).astype(np.uint8)
     result = Image.fromarray(final)
-    result = result.filter(ImageFilter.UnsharpMask(radius=0.6, percent=60, threshold=1))
+    
+    # Better sharpening
+    result = result.filter(ImageFilter.UnsharpMask(radius=0.8, percent=80, threshold=2))
+    
     return result
 
-
 # ====================== MAIN PROCESSING ======================
 @timed("Total Processing")
 def process_face_whitening(input_image: Image.Image):
@@ -231,7 +245,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=["*"])
@@ -241,33 +254,30 @@ async def startup():
     loop = asyncio.get_event_loop()
     await loop.run_in_executor(executor, load_face_parser)
     await loop.run_in_executor(executor, load_beard_model)
-
     logger.info("✅ Models loaded")
     logger.info("Running light warmup...")
     dummy = Image.new("RGB", (256, 256))
     _ = process_face_whitening(dummy)
     logger.info("✅ Server Ready!")
 
-
 @app.post("/age-face")
 async def age_face(file: UploadFile = File(...)):
     start_total = time.perf_counter()
     contents = await file.read()
     img = Image.open(io.BytesIO(contents)).convert("RGB")
-    
+   
     loop = asyncio.get_event_loop()
     result = await loop.run_in_executor(executor, process_face_whitening, img)
-    
+   
     buf = io.BytesIO()
-    result.save(buf, format="JPEG", quality=90, optimize=True)
+    result.save(buf, format="JPEG", quality=92, optimize=True)
     buf.seek(0)
-    
+   
     total_time = (time.perf_counter() - start_total) * 1000
     logger.info(f"✅ Total Request Time: {total_time:.1f} ms")
-    
+   
     return StreamingResponse(buf, media_type="image/jpeg")
 
-
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860, workers=1)