Update app.py

app.py

@@ -1,143 +1,247 @@
+# ================================================
+# END-TO-END TIMING LOGGING (NO HARDCODE)
+# ================================================
 import os
 import torch
 import numpy as np
 import cv2
-import time
-import asyncio
-import io
 import traceback
 import gc
-from PIL import Image
+import time
+from PIL import Image, ImageFilter
 from transformers import SegformerImageProcessor
 from fastapi import FastAPI, File, UploadFile, HTTPException
 from fastapi.responses import StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
+import io
+import asyncio
 from concurrent.futures import ThreadPoolExecutor
 import logging
 import onnxruntime as ort
 
-logging.basicConfig(level=logging.INFO)
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-PROCESS_SIZE = 256
-ONNX_PATH = "models/segformer_face_parsing.onnx"
+PROCESS_SIZE = 384
+onnx_path = "models/segformer_face_parsing.onnx"
 os.makedirs("models", exist_ok=True)
 
 executor = ThreadPoolExecutor(max_workers=1)
 face_processor = None
 ort_session = None
 
-def ensure_onnx_exists():
-    if os.path.exists(ONNX_PATH):
-        return True
-    logger.info("Exporting ONNX (once, takes ~30s)...")
+def convert_to_onnx():
+    if os.path.exists(onnx_path):
+        return
+    logger.info("⚙️ Converting Segformer to ONNX (first time only)")
     from transformers import SegformerForSemanticSegmentation
     model = SegformerForSemanticSegmentation.from_pretrained("jonathandinu/face-parsing")
     model.eval()
-    dummy = torch.randn(1, 3, 192, 192)
-    torch.onnx.export(
-        model, dummy, ONNX_PATH,
-        input_names=["pixel_values"],
-        output_names=["logits"],
-        dynamic_axes={"pixel_values": {0: "batch", 2: "height", 3: "width"},
-                      "logits": {0: "batch", 2: "height", 3: "width"}},
-        opset_version=14,
-        do_constant_folding=False
-    )
-    logger.info("ONNX ready")
-    return True
+    dummy_input = torch.randn(1, 3, 192, 192)
+    torch.onnx.export(model, dummy_input, onnx_path,
+                      input_names=["pixel_values"],
+                      output_names=["logits"],
+                      dynamic_axes={"pixel_values": {0: "batch", 2: "height", 3: "width"}},
+                      opset_version=14, do_constant_folding=True)
+    logger.info("✅ ONNX conversion done")
 
 def load_face_parser():
     global face_processor, ort_session
     if ort_session is not None:
         return
-    ensure_onnx_exists()
+    t0 = time.time()
+    convert_to_onnx()
     face_processor = SegformerImageProcessor.from_pretrained("jonathandinu/face-parsing")
-    opts = ort.SessionOptions()
-    opts.intra_op_num_threads = 1
-    ort_session = ort.InferenceSession(ONNX_PATH, opts, providers=['CPUExecutionProvider'])
-    logger.info("ONNX loaded")
+    sess_options = ort.SessionOptions()
+    sess_options.intra_op_num_threads = 1
+    sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+    ort_session = ort.InferenceSession(onnx_path, sess_options, providers=['CPUExecutionProvider'])
+    logger.info(f"✅ ONNX loaded in {time.time()-t0:.2f} sec")
 
-def get_hair_mask(pil_image):
+def get_hair_and_exclude_masks(pil_image: Image.Image):
     load_face_parser()
     orig_w, orig_h = pil_image.size
-    # Always feed 192x192 to ONNX
-    img_model = pil_image.resize((192,192), Image.LANCZOS)
-    inputs = face_processor(images=img_model, return_tensors="pt")
-    pixel_vals = inputs["pixel_values"].numpy().astype(np.float32)
-    logits = torch.from_numpy(ort_session.run(["logits"], {"pixel_values": pixel_vals})[0])
-    up = torch.nn.functional.interpolate(logits, size=(192,192), mode="bilinear")
+    
+    t0 = time.time()
+    img_small = pil_image.resize((192, 192), Image.LANCZOS)
+    inputs = face_processor(images=img_small, return_tensors="pt")
+    pixel_values = inputs["pixel_values"].numpy().astype(np.float32)
+    logger.debug(f"    - Preprocess: {time.time()-t0:.3f}s")
+    
+    t0 = time.time()
+    ort_inputs = {"pixel_values": pixel_values}
+    ort_outs = ort_session.run(["logits"], ort_inputs)
+    logits = torch.from_numpy(ort_outs[0])
+    logger.debug(f"    - ONNX inference: {time.time()-t0:.3f}s")
+    
+    t0 = time.time()
+    up = torch.nn.functional.interpolate(logits, size=(192, 192), mode="bilinear", align_corners=False)
     probs = torch.softmax(up, dim=1)[0]
-    strong = (probs[13].cpu().numpy() > 0.055).astype(np.float32)
-    soft = (probs[13].cpu().numpy() > 0.022).astype(np.float32)
-    hair = np.maximum(strong, soft * 0.68)
-    # face subtraction
+    logger.debug(f"    - Softmax+upsample: {time.time()-t0:.3f}s")
+    
+    # Hair mask
+    t0 = time.time()
+    strong_hair = (probs[13].cpu().numpy() > 0.055).astype(np.float32)
+    soft_hair = (probs[13].cpu().numpy() > 0.022).astype(np.float32)
+    hair = np.maximum(strong_hair, soft_hair * 0.68)
     parsing = up.argmax(dim=1).squeeze(0).cpu().numpy()
-    face_cls = list(range(1,6)) + list(range(8,13)) + [17,18]
-    face_mask = np.isin(parsing, face_cls).astype(np.float32)
-    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3,3))
-    face_mask = cv2.dilate(face_mask, kernel, iterations=1)
-    h,w = face_mask.shape
-    forehead = np.zeros_like(face_mask)
+    face_cls = list(range(1, 6)) + list(range(8, 13)) + [17, 18]
+    face_m = np.isin(parsing, face_cls).astype(np.float32)
+    kernel_face = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3,3))
+    face_m = cv2.dilate(face_m, kernel_face, iterations=1)
+    h, w = face_m.shape
+    forehead = np.zeros_like(face_m)
     forehead[:int(h*0.3), :] = 1.0
-    face_mask = face_mask * (1 - forehead*0.45)
-    hair = hair * (1 - face_mask)
-    hair = cv2.morphologyEx(hair, cv2.MORPH_CLOSE, kernel, iterations=1)
-    hair = cv2.GaussianBlur(hair, (3,3), 0.8)
-    hair = cv2.resize(hair, (orig_w, orig_h))
-    return hair
+    face_m = face_m * (1 - forehead * 0.45)
+    hair = hair * (1 - face_m)
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3,3))
+    hair = cv2.morphologyEx(hair, cv2.MORPH_CLOSE, kernel, iterations=2)
+    hair = cv2.GaussianBlur(hair, (5,5), 1.5)
+    hair = cv2.resize(hair, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR)
+    logger.debug(f"    - Hair postprocess: {time.time()-t0:.3f}s")
+    
+    # Exclude mask
+    t0 = time.time()
+    nose = (probs[2].cpu().numpy() > 0.5).astype(np.float32)
+    lip_up = (probs[11].cpu().numpy() > 0.5).astype(np.float32)
+    lip_low = (probs[12].cpu().numpy() > 0.5).astype(np.float32)
+    exclude = np.clip(nose + lip_up + lip_low, 0, 1)
+    exclude = cv2.resize(exclude, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)
+    kernel_ex = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5,5))
+    exclude = cv2.dilate(exclude, kernel_ex, iterations=2)
+    logger.debug(f"    - Exclude mask: {time.time()-t0:.3f}s")
+    
+    return hair, exclude
 
-def apply_grey_hair(image, hair_mask):
-    comb = hair_mask
-    comb = cv2.GaussianBlur(comb, (5,5), 1)
-    img = np.array(image).astype(np.float32)/255.0
+def apply_strong_grey_hair(image: Image.Image, hair_mask: np.ndarray, beard_mask: np.ndarray) -> Image.Image:
+    t0 = time.time()
+    comb = np.maximum(hair_mask, beard_mask)
+    if np.sum(comb) < 100:
+        logger.warning("⚠️ Small mask area")
+    comb = cv2.GaussianBlur(comb, (7,7), 2)
+    img = np.array(image).astype(np.float32) / 255.0
     hsv = cv2.cvtColor((img*255).astype(np.uint8), cv2.COLOR_RGB2HSV).astype(np.float32)
-    hsv[:,:,1] = hsv[:,:,1] * (1 - 0.7*comb)
-    hsv[:,:,2] = np.clip(hsv[:,:,2] + (70*comb), 100, 230)
-    result = cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2RGB).astype(np.float32)/255.0
-    comb_3ch = np.stack([comb,comb,comb], axis=2)
-    final = result*comb_3ch + img*(1-comb_3ch)
+    hsv_hair = hsv.copy()
+    saturation_factor = 0.8
+    brightness_boost = 90
+    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
+    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)
+    hair_mask_binary = (hair_mask > 0.5)
+    if np.sum(hair_mask_binary) > 100:
+        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:
+        mean_hair_lab = np.array([128,0,0])
+        std_hair_lab = np.array([30,10,10])
+    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)
+        std_beard_lab = np.maximum(std_beard_lab, 1e-5)
+        beard_norm = (beard_pixels_lab - mean_beard_lab) / std_beard_lab
+        beard_transfer = beard_norm * std_hair_lab + mean_hair_lab
+        beard_transfer = np.clip(beard_transfer, 0, 255)
+        img_lab_transfer = img_lab.copy()
+        img_lab_transfer[beard_mask_binary] = beard_transfer
+    else:
+        img_lab_transfer = img_lab
+    final = cv2.cvtColor(img_lab_transfer.astype(np.uint8), cv2.COLOR_LAB2RGB).astype(np.float32) / 255.0
+    hair_mask_3ch = np.stack([hair_mask, hair_mask, hair_mask], axis=2)
+    final = hair_grey * hair_mask_3ch + final * (1 - hair_mask_3ch)
+    comb_3ch = np.stack([comb, comb, comb], axis=2)
+    final = final * comb_3ch + img * (1 - comb_3ch)
+    warm = np.array([5,3,0], dtype=np.float32)/255.0
+    final = final + (warm * comb[..., None] * 0.2)
     final = np.clip(final*255, 0, 255).astype(np.uint8)
-    return Image.fromarray(final)
+    result = Image.fromarray(final)
+    result = result.filter(ImageFilter.UnsharpMask(radius=0.5, percent=50, threshold=0))
+    logger.debug(f"    - Color transfer: {time.time()-t0:.3f}s")
+    return result
 
-def process_image(input_image):
+def process_face_whitening(input_image: Image.Image):
+    total_start = time.time()
+    logger.info("="*50)
+    logger.info("🖼️ Processing new image")
+    
+    t0 = time.time()
     orig = input_image.convert("RGB")
     ow, oh = orig.size
     img_resized = orig.resize((PROCESS_SIZE, PROCESS_SIZE), Image.LANCZOS)
-    hair = get_hair_mask(img_resized)
-    result = apply_grey_hair(img_resized, hair)
-    final = result.resize((ow, oh), Image.LANCZOS)
-    return final
+    logger.info(f"📥 Step 1 - Load & resize: {time.time()-t0:.2f}s (to {PROCESS_SIZE}x{PROCESS_SIZE})")
+    
+    t0 = time.time()
+    hair_mask, exclude_mask = get_hair_and_exclude_masks(img_resized)
+    logger.info(f"🎭 Step 2 - Mask generation: {time.time()-t0:.2f}s (hair sum: {np.sum(hair_mask):.0f})")
+    
+    t0 = time.time()
+    beard_mask = np.zeros_like(hair_mask)
+    logger.info(f"🧔 Step 3 - Beard mask (skipped): {time.time()-t0:.2f}s")
+    
+    t0 = time.time()
+    result_resized = apply_strong_grey_hair(img_resized, hair_mask, beard_mask)
+    logger.info(f"🎨 Step 4 - Color transfer: {time.time()-t0:.2f}s")
+    
+    t0 = time.time()
+    final_img = result_resized.resize((ow, oh), Image.LANCZOS)
+    logger.info(f"📤 Step 5 - Resize to original: {time.time()-t0:.2f}s ({ow}x{oh})")
+    
+    processing_time = time.time() - total_start
+    logger.info(f"⚙️ Core processing time: {processing_time:.2f} seconds")
+    return final_img, processing_time
 
-app = FastAPI()
+# ================================================
+# FASTAPI APP WITH END-TO-END TIMING
+# ================================================
+app = FastAPI(title="Grey Hair API (Accurate Timing)")
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
 
-sem = asyncio.Semaphore(1)
-
 @app.on_event("startup")
 async def startup():
+    t0 = time.time()
     loop = asyncio.get_event_loop()
     await loop.run_in_executor(executor, load_face_parser)
-    logger.info("Server ready")
+    logger.info(f"🔥 Server ready in {time.time()-t0:.2f} seconds")
 
 @app.post("/age-face")
 async def age_face(file: UploadFile = File(...)):
-    await sem.acquire()
-    try:
-        data = await file.read()
-        img = Image.open(io.BytesIO(data)).convert("RGB")
-        loop = asyncio.get_event_loop()
-        out = await loop.run_in_executor(executor, process_image, img)
-        buf = io.BytesIO()
-        out.save(buf, format="JPEG", quality=90)
-        buf.seek(0)
-        return StreamingResponse(buf, media_type="image/jpeg")
-    except Exception as e:
-        logger.error(traceback.format_exc())
-        raise HTTPException(500, str(e))
-    finally:
-        sem.release()
-        gc.collect()
+    request_start = time.time()
+    logger.info("🚀 Request received")
+    
+    # Step A: Read file
+    t0 = time.time()
+    contents = await file.read()
+    read_time = time.time() - t0
+    logger.info(f"📁 File read: {read_time:.3f}s ({len(contents)} bytes)")
+    
+    # Step B: Decode image
+    t0 = time.time()
+    img = Image.open(io.BytesIO(contents)).convert("RGB")
+    decode_time = time.time() - t0
+    logger.info(f"🖼️ Image decode: {decode_time:.3f}s ({img.size[0]}x{img.size[1]})")
+    
+    # Step C: Process (core)
+    loop = asyncio.get_event_loop()
+    result_img, core_time = await loop.run_in_executor(executor, process_face_whitening, img)
+    
+    # Step D: Encode to JPEG
+    t0 = time.time()
+    buf = io.BytesIO()
+    result_img.save(buf, format="JPEG", quality=92)
+    encode_time = time.time() - t0
+    logger.info(f"💾 JPEG encode: {encode_time:.3f}s")
+    buf.seek(0)
+    
+    # Step E: Send response
+    total_time = time.time() - request_start
+    logger.info(f"📡 Total end-to-end time: {total_time:.2f} seconds (core: {core_time:.2f}, overhead: {total_time-core_time:.2f})")
+    logger.info("="*50)
+    
+    return StreamingResponse(buf, media_type="image/jpeg")
 
 if __name__ == "__main__":
     import uvicorn