Update app.py

app.py

@@ -5,7 +5,7 @@ import numpy as np
 import cv2
 import traceback
 import gc
-from PIL import Image, ImageFilter, ImageEnhance, ImageDraw, ImageFont
+from PIL import Image, ImageFilter, ImageEnhance
 from torchvision.transforms import functional as TF
 from scipy.ndimage import label
 import antialiased_cnns
@@ -15,7 +15,10 @@ from transformers import SegformerImageProcessor, SegformerForSemanticSegmentati
 from ultralytics import YOLO
 from gfpgan import GFPGANer
 import urllib.request
-import gradio as gr
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.responses import StreamingResponse
+from fastapi.middleware.cors import CORSMiddleware
+import io
 
 # ========================= CONFIG =========================
 AGING_MODEL_PATH = "face_aging_model/best_unet_model.pth"
@@ -96,6 +99,7 @@ def load_aging_model():
     if age_model is not None:
         return age_model
     print("Loading UNet aging model...")
+
     class DownLayer(nn.Module):
         def __init__(self, in_ch, out_ch):
             super().__init__()
@@ -161,9 +165,11 @@ def load_aging_model():
     state = torch.load(AGING_MODEL_PATH, map_location=DEVICE, weights_only=True)
     age_model.load_state_dict(state)
     age_model.eval()
+
     if DEVICE.type == "cuda" and int(torch.__version__.split('.')[0]) >= 2:
         print("Compiling UNet with torch.compile...")
         age_model = torch.compile(age_model, mode="reduce-overhead")
+
     print("✅ Aging model loaded!")
     return age_model
 
@@ -177,9 +183,11 @@ def load_face_parser():
     face_parser = SegformerForSemanticSegmentation.from_pretrained("jonathandinu/face-parsing")
     face_parser.to(DEVICE)
     face_parser.eval()
+
     if DEVICE.type == "cuda" and int(torch.__version__.split('.')[0]) >= 2:
         print("Compiling Segformer with torch.compile...")
         face_parser = torch.compile(face_parser, mode="reduce-overhead")
+
     print("✅ Face parser loaded!")
     return face_processor, face_parser
 
@@ -196,12 +204,12 @@ def get_lips_mask(pil_image: Image.Image) -> np.ndarray:
     img_np = np.array(pil_image)
     h, w = img_np.shape[:2]
     lips_mask = np.zeros((h, w), dtype=np.uint8)
-   
+  
     with mp_face_mesh.FaceMesh(static_image_mode=True, max_num_faces=1, refine_landmarks=True,
                                min_detection_confidence=0.5) as face_mesh:
         rgb_image = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
         results = face_mesh.process(rgb_image)
-       
+      
         if results.multi_face_landmarks:
             for face_landmarks in results.multi_face_landmarks:
                 lip_landmarks = [61, 146, 91, 181, 84, 17, 314, 405, 321, 375, 291, 308, 324, 318, 402, 317, 14, 87, 178, 88, 95]
@@ -240,11 +248,11 @@ def get_beard_mask(pil_image: Image.Image) -> np.ndarray:
         model = load_beard_model()
         results = model(temp_path, device=DEVICE.type, conf=0.25, iou=0.5, verbose=False,
                         half=True if DEVICE.type == "cuda" else False)
-       
+      
         img_np = np.array(pil_image)
         h, w = img_np.shape[:2]
         beard_mask = np.zeros((h, w), dtype=np.uint8)
-       
+      
         if results[0].masks is not None:
             for i, cls in enumerate(results[0].boxes.cls):
                 if int(cls) == 0: # beard class
@@ -252,7 +260,7 @@ def get_beard_mask(pil_image: Image.Image) -> np.ndarray:
                     mask = cv2.resize(mask, (w, h))
                     mask = (mask > 0.4).astype(np.uint8) * 255
                     beard_mask = cv2.bitwise_or(beard_mask, mask)
-       
+      
         if np.sum(beard_mask) > 0:
             beard_mask_float = beard_mask.astype(np.float32) / 255.0
             beard_mask_float = cv2.dilate(beard_mask_float, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7)), iterations=2)
@@ -262,7 +270,7 @@ def get_beard_mask(pil_image: Image.Image) -> np.ndarray:
             beard_mask_float = cv2.GaussianBlur(beard_mask_float, (7, 7), 2)
             beard_mask_float = np.clip(beard_mask_float, 0, 1)
             return beard_mask_float
-       
+      
         return np.zeros((h, w), dtype=np.float32)
     finally:
         if os.path.exists(temp_path):
@@ -280,25 +288,25 @@ def clean_mask(mask, min_area=150):
 def get_hair_mask_segformer(pil_image: Image.Image) -> np.ndarray:
     processor, parser = load_face_parser()
     inputs = processor(images=pil_image, return_tensors="pt").to(DEVICE)
-   
+  
     with torch.no_grad():
         outputs = parser(**inputs)
-   
+  
     logits = outputs.logits
     upsampled = torch.nn.functional.interpolate(logits, size=pil_image.size[::-1], mode="bilinear", align_corners=False)
     probs = torch.softmax(upsampled, dim=1)[0]
     hair_prob = probs[13].cpu().numpy()
-   
+  
     hair_mask = (hair_prob > 0.12).astype(np.uint8)
-   
+  
     face_classes = list(range(1, 6)) + list(range(8, 13)) + [17, 18]
     parsing = upsampled.argmax(dim=1).squeeze(0).cpu().numpy()
     face_mask = np.isin(parsing, face_classes).astype(np.uint8)
-   
+  
     kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
     face_mask = cv2.dilate(face_mask, kernel, iterations=1)
     hair_mask = hair_mask * (1 - face_mask)
-   
+  
     hair_mask = cv2.morphologyEx(hair_mask, cv2.MORPH_OPEN, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3)), iterations=1)
     hair_mask = cv2.morphologyEx(hair_mask, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (11, 11)), iterations=2)
     hair_mask = clean_mask(hair_mask, min_area=100)
@@ -310,22 +318,22 @@ def apply_hair_and_beard_color(image: Image.Image, hair_mask: np.ndarray, beard_
     combined_mask = np.maximum(hair_mask, beard_mask)
     if np.sum(combined_mask) == 0:
         return image
-   
+  
     combined_mask = cv2.GaussianBlur(combined_mask, (BLUR_RADIUS*2+1, BLUR_RADIUS*2+1), BLUR_RADIUS)
     combined_mask = np.clip(combined_mask, 0, 1)
-   
+  
     if EDGE_SMOOTHING:
         combined_mask = cv2.bilateralFilter(combined_mask.astype(np.float32), 9, 75, 75)
         combined_mask = np.clip(combined_mask, 0, 1)
-   
+  
     combined_mask = np.clip(combined_mask * 1.2, 0, 1)
-   
+  
     img_np = np.array(image).astype(np.float32)
     target_color = np.array([255, 255, 255], dtype=np.float32)
     gray = cv2.cvtColor(img_np.astype(np.uint8), cv2.COLOR_RGB2GRAY).astype(np.float32) / 255.0
     lum_factor = 0.6 + 0.4 * gray
     white_layer = target_color * lum_factor[..., np.newaxis]
-   
+  
     alpha = ALPHA_HAIR
     result = (1 - alpha * combined_mask[..., np.newaxis]) * img_np + (alpha * combined_mask[..., np.newaxis]) * white_layer
     result = np.clip(result, 0, 255).astype(np.uint8)
@@ -348,34 +356,10 @@ def enhance_texture(img: Image.Image) -> Image.Image:
     img = ImageEnhance.Sharpness(img).enhance(SHARPNESS_BOOST)
     return img
 
-def create_comparison(orig, raw_aged, final):
-    W = 640
-    def rsz(img):
-        ratio = img.height / img.width if img.width else 1
-        return img.resize((W, int(W * ratio)), Image.LANCZOS)
-   
-    o, r, f = rsz(orig), rsz(raw_aged), rsz(final)
-    H = max(o.height, r.height, f.height)
-    canvas = Image.new("RGB", (W*3, H), (255, 255, 255))
-    canvas.paste(o, (0, (H - o.height)//2))
-    canvas.paste(r, (W, (H - r.height)//2))
-    canvas.paste(f, (W*2, (H - f.height)//2))
-   
-    draw = ImageDraw.Draw(canvas)
-    try:
-        font = ImageFont.truetype("arial.ttf", 28)
-    except:
-        font = ImageFont.load_default()
-   
-    draw.text((W//4, 8), "Original", (0, 0, 0), font=font)
-    draw.text((W + W//5, 8), "Aged Raw", (0, 0, 0), font=font)
-    draw.text((W*2 + W//6, 8), "Final Result", (0, 0, 0), font=font)
-    return canvas
-
 # ================== MAIN PROCESSING FUNCTION ==================
-def process_face_aging(input_image: Image.Image):
+def process_face_aging(input_image: Image.Image) -> Image.Image:
     if input_image is None:
-        raise gr.Error("Please upload a clear photo of a young person!")
+        raise ValueError("Please provide a valid image!")
 
     try:
         print(f"→ Processing image: {input_image.size}")
@@ -387,27 +371,27 @@ def process_face_aging(input_image: Image.Image):
 
         src_age = torch.full((1, SAFE_IMG_SIZE, SAFE_IMG_SIZE), SOURCE_AGE / 100.0)
         tgt_age = torch.full((1, SAFE_IMG_SIZE, SAFE_IMG_SIZE), TARGET_AGE / 100.0)
+
         cond_input = torch.cat([rgb_tensor, src_age, tgt_age], dim=0).unsqueeze(0).to(DEVICE)
 
         with torch.no_grad():
             aging_net = load_aging_model()
             raw_output = aging_net(cond_input).squeeze(0)
-            raw_aged = TF.to_pil_image(raw_output.clamp(0, 1)).resize((ow, oh), Image.LANCZOS)
-           
+
             alpha = WRINKLE_STRENGTH
             blended = (1 - alpha) * rgb_tensor.unsqueeze(0) + alpha * raw_output
             blended = blended.clamp(0, 1).squeeze(0)
+
             final_aged = TF.to_pil_image(blended).resize((ow, oh), Image.LANCZOS)
-           
             final_aged = enhance_texture(final_aged)
             final_aged = post_correct_aged(orig, final_aged)
 
         print(" Generating hair mask...")
         hair_mask = get_hair_mask_segformer(final_aged)
-       
+
         print(" Generating beard mask...")
         beard_mask = get_beard_mask(final_aged)
-       
+
         print(" Applying white hair & beard...")
         final_img = apply_hair_and_beard_color(final_aged, hair_mask, beard_mask)
 
@@ -424,44 +408,52 @@ def process_face_aging(input_image: Image.Image):
                 except Exception as e:
                     print(f" GFPGAN error: {e}")
 
-        comparison = create_comparison(orig, raw_aged, final_img)
-       
         print("✓ Processing completed!")
         gc.collect()
-        return final_img, comparison
+        return final_img
 
     except Exception as e:
         print(f"❌ Error: {str(e)}")
         traceback.print_exc()
-        raise gr.Error(f"Processing failed: {str(e)}")
+        raise
+
+# ================== FASTAPI SETUP ==================
+app = FastAPI(title="Face Aging + White Hair & Beard API")
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+@app.post("/age-face")
+async def age_face(file: UploadFile = File(...)):
+    if not file.content_type.startswith("image/"):
+        raise HTTPException(status_code=400, detail="Only image files allowed")
 
-# ================== GRADIO INTERFACE ==================
-with gr.Blocks(theme=gr.themes.Soft(), title="👴 Face Aging + White Hair & Beard Generator") as demo:
-    gr.Markdown("# 👴 Face Aging + White Hair & Beard Generator")
-    gr.Markdown("Upload a clear photo of a young person.<br>This tool will age them to ~90 years with realistic wrinkles and add natural white hair & beard.")
+    contents = await file.read()
 
-    with gr.Row():
-        input_img = gr.Image(type="pil", label="Upload Young Face Photo", height=450)
+    try:
+        input_image = Image.open(io.BytesIO(contents)).convert("RGB")
+        result_image = process_face_aging(input_image)
+
+        # Convert result to bytes
+        img_byte_arr = io.BytesIO()
+        result_image.save(img_byte_arr, format="PNG")
+        img_byte_arr.seek(0)
 
-    with gr.Row():
-        output_img = gr.Image(type="pil", label="Final Aged Result (with White Hair & Beard)", height=450)
-        comparison_img = gr.Image(type="pil", label="Comparison: Original | Raw Aged | Final", height=450)
+        return StreamingResponse(img_byte_arr, media_type="image/png")
 
-    btn = gr.Button("🚀 Generate Aged Face", variant="primary")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Processing failed: {str(e)}")
+    finally:
+        gc.collect()
 
-    btn.click(
-        fn=process_face_aging,
-        inputs=input_img,
-        outputs=[output_img, comparison_img],
-        queue=True,
-        concurrency_limit=2   # ← Yeh line concurrency_count ki jagah use hui
-    )
 
+# For local testing
 if __name__ == "__main__":
-    print("Starting Face Aging App...")
-    demo.queue(max_size=8).launch(   # ← concurrency_count add kar do
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-        debug=False
-    )
+    import uvicorn
+    print("Starting FastAPI server...")
+    uvicorn.run(app, host="0.0.0.0", port=7860)