Update app.py

app.py

@@ -1,239 +1,33 @@
-import gradio as gr
-import torch
-import cv2
-import numpy as np
-from PIL import Image
-import logging
-from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
+# main_fastapi.py
+from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi.responses import StreamingResponse
+import io
+from bald_processor import make_realistic_bald
 
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+app = FastAPI(title="Make Me Bald API 😎")
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-logger.info(f"Using device: {device}")
+@app.post("/make-bald/")
+async def bald_endpoint(file: UploadFile = File(...)):
+    if not file.content_type.startswith("image/"):
+        raise HTTPException(status_code=400, detail="Sirf image file upload kar!")
 
-print("Loading SegFormer face-parsing model...")
-processor = SegformerImageProcessor.from_pretrained("jonathandinu/face-parsing")
-model = SegformerForSemanticSegmentation.from_pretrained("jonathandinu/face-parsing")
-model.to(device)
-model.eval()
-logger.info("Model loaded!")
-
-hair_class_id = 13
-ear_class_ids = [7, 8]
-
-def make_realistic_bald(input_image: Image.Image) -> tuple[Image.Image, Image.Image, Image.Image]:
-    # (Yeh pura function tera perfect logic wala — same as before, no change)
     try:
-        orig_w, orig_h = input_image.size
-        original_np = np.array(input_image)
-        original_bgr = cv2.cvtColor(original_np, cv2.COLOR_RGB2BGR)
-        logger.info(f"Processing: {orig_w}x{orig_h}")
-
-        MAX_PROCESS_DIM = 2048
-        scale_factor = 1.0
-        working_np = original_np
-        working_bgr = original_bgr
-        working_h, working_w = orig_h, orig_w
-
-        if max(orig_w, orig_h) > MAX_PROCESS_DIM:
-            logger.info(f"Downscaling to max {MAX_PROCESS_DIM}px")
-            scale_factor = MAX_PROCESS_DIM / max(orig_w, orig_h)
-            working_w = int(orig_w * scale_factor)
-            working_h = int(orig_h * scale_factor)
-            working_np = cv2.resize(original_np, (working_w, working_h), cv2.INTER_AREA)
-            working_bgr = cv2.cvtColor(working_np, cv2.COLOR_RGB2BGR)
-
-        pil_working = Image.fromarray(working_np)
-        inputs = processor(images=pil_working, return_tensors="pt").to(device)
-        with torch.no_grad():
-            outputs = model(**inputs)
-            logits = outputs.logits
-
-        upsampled_logits = torch.nn.functional.interpolate(
-            logits, size=(working_h, working_w), mode="bilinear", align_corners=False
-        )
-        parsing = upsampled_logits.argmax(dim=1).squeeze(0).cpu().numpy()
-
-        hair_mask = (parsing == hair_class_id).astype(np.uint8)
-
-        ears_mask = np.zeros_like(hair_mask)
-        for cls in ear_class_ids:
-            ears_mask[parsing == cls] = 1
-
-        ear_y, ear_x = np.where(ears_mask)
-        if len(ear_y) > 0:
-            ear_top_y = ear_y.min()
-            ear_height = ear_y.max() - ear_top_y + 1
-            kernel_v = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (11, 30))
-            ears_protected = cv2.dilate(ears_mask, kernel_v, iterations=2)
-
-            top_margin = max(8, int(ear_height * 0.12))
-            top_start = max(0, ear_top_y - top_margin)
-
-            ear_x_min, ear_x_max = ear_x.min(), ear_x.max()
-            ear_width = ear_x_max - ear_x_min + 1
-            x_margin = int(ear_width * 0.35)
-            protected_left = max(0, ear_x_min - x_margin)
-            protected_right = min(working_w, ear_x_max + x_margin)
-
-            limited_top_mask = np.zeros_like(ears_mask)
-            limited_top_mask[top_start:ear_top_y + 8, protected_left:protected_right] = 1
-            kernel_h = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (17, 5))
-            limited_top_mask = cv2.dilate(limited_top_mask, kernel_h, iterations=1)
-
-            ears_protected = np.logical_or(ears_protected, limited_top_mask).astype(np.uint8)
-
-            hair_above_ears = np.zeros_like(hair_mask)
-            above_ear_line = max(0, ear_top_y - int(ear_height * 0.65))
-            hair_above_ears[:above_ear_line, :] = hair_mask[:above_ear_line, :]
-            ears_protected[hair_above_ears == 1] = 0
+        contents = await file.read()
+        from PIL import Image
+        image = Image.open(io.BytesIO(contents)).convert("RGB")
+        bald_img = make_realistic_bald(image)
+        buf = io.BytesIO()
+        bald_img.save(buf, format="JPEG")
+        buf.seek(0)
+        return StreamingResponse(buf, media_type="image/jpeg", headers={"Content-Disposition":"attachment; filename=bald.jpg"})
+    except ValueError as ve:
+        if str(ve) == "NO_HAIR_DETECTED":
+            raise HTTPException(400, detail="NO_HAIR_DETECTED")
         else:
-            ears_protected = np.zeros_like(hair_mask)
-
-        hair_mask_final = hair_mask.copy()
-        hair_mask_final[ears_protected == 1] = 0
-
-        if hair_mask[:int(working_h * 0.25), :].sum() > 60:
-            hair_mask_final[:int(working_h * 0.25), :] = np.maximum(
-                hair_mask_final[:int(working_h * 0.25), :], hair_mask[:int(working_h * 0.25), :]
-            )
-
-        kernel_s = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (13, 13))
-        hair_mask_final = cv2.morphologyEx(hair_mask_final, cv2.MORPH_CLOSE, kernel_s, iterations=2)
-        hair_mask_final = cv2.dilate(hair_mask_final, kernel_s, iterations=1)
-
-        blurred = cv2.GaussianBlur(hair_mask_final.astype(np.float32), (9, 9), 3)
-        hair_mask_final = (blurred > 0.28).astype(np.uint8)
-
-        kernel_edge = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
-        hair_mask_final = cv2.dilate(hair_mask_final, kernel_edge, iterations=1)
-
-        hair_pixels = np.sum(hair_mask_final)
-        logger.info(f"Hair pixels (resized): {hair_pixels:,}")
-
-        final_mask = hair_mask_final.copy()
-        use_extended_mask = False
-        if hair_pixels > 380000:
-            logger.info("Large hair → extended mask")
-            use_extended_mask = True
-            big_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (25, 25))
-            extended = cv2.dilate(hair_mask_final, big_kernel, iterations=1)
-            upper = np.zeros_like(hair_mask_final)
-            upper_end = int(working_h * 0.48)
-            upper[:upper_end, :] = 1
-            extended = np.logical_or(extended, upper).astype(np.uint8)
-            extended[ears_protected == 1] = 0
-
-            if np.mean(working_np) < 110:
-                hsv = cv2.cvtColor(working_np, cv2.COLOR_RGB2HSV)
-                dark_lower = np.array([0, 0, 0])
-                dark_upper = np.array([180, 70, 90])
-                dark_mask = cv2.inRange(hsv, dark_lower, dark_upper)
-                extended = np.logical_or(extended, (dark_mask > 127)).astype(np.uint8)
-
-            extended = cv2.morphologyEx(extended, cv2.MORPH_CLOSE, kernel_s, iterations=1)
-            extended[int(working_h * 0.75):, :] = 0
-            final_mask = extended
-
-        if use_extended_mask or hair_pixels > 420000:
-            radius, flag = 18, cv2.INPAINT_TELEA
-        elif hair_pixels > 220000:
-            radius, flag = 15, cv2.INPAINT_TELEA
-        else:
-            radius, flag = 10, cv2.INPAINT_NS
-
-        logger.info(f"Inpainting radius={radius}")
-        inpainted_bgr = cv2.inpaint(working_bgr, final_mask * 255, inpaintRadius=radius, flags=flag)
-        inpainted_rgb = cv2.cvtColor(inpainted_bgr, cv2.COLOR_BGR2RGB)
-
-        result_small = working_np.copy()
-        result_small[final_mask == 1] = inpainted_rgb[final_mask == 1]
-
-        if use_extended_mask or hair_pixels > 280000:
-            logger.info("Skin color correction")
-            regions = [(0.18, 0.30, 0.34, 0.66), (0.32, 0.47, 0.35, 0.65)]
-            colors = []
-            for y1r, y2r, x1r, x2r in regions:
-                y1, y2 = int(working_h * y1r), int(working_h * y2r)
-                x1, x2 = int(working_w * x1r), int(working_w * x2r)
-                if y2 > y1 + 40 and x2 > x1 + 80:
-                    crop = working_np[y1:y2, x1:x2]
-                    if crop.size > 0:
-                        colors.append(np.median(crop, axis=(0,1)).astype(np.float32))
-
-            if colors:
-                target_color = np.mean(colors, axis=0)
-                brightness = np.mean(target_color)
-                strength = 0.82 if brightness > 145 else 0.62 if brightness < 85 else 0.74
-                bald_area = result_small[final_mask == 1].astype(np.float32)
-                if len(bald_area) > 200:
-                    current_mean = bald_area.mean(axis=0)
-                    diff = target_color - current_mean
-                    corrected = np.clip(bald_area + diff * strength, 0, 255).astype(np.uint8)
-                    result_small[final_mask == 1] = corrected
-
-        if hair_pixels > 90000 or use_extended_mask:
-            blurred_bald = cv2.GaussianBlur(result_small, (5, 5), 0.8)
-            result_small[final_mask == 1] = cv2.addWeighted(
-                result_small[final_mask == 1], 0.65, blurred_bald[final_mask == 1], 0.35, 0
-            )
-
-        if scale_factor < 1.0:
-            logger.info("Upscaling to original size")
-            result = cv2.resize(result_small, (orig_w, orig_h), interpolation=cv2.INTER_LANCZOS4)
-        else:
-            result = result_small
-
-        result_pil = Image.fromarray(result)
-
-        comparison = np.hstack((original_np, result))
-        comparison_pil = Image.fromarray(comparison)
-
-        final_mask_big = cv2.resize(final_mask.astype(np.uint8) * 255, (orig_w, orig_h), cv2.INTER_NEAREST) > 127
-        mask_vis = np.zeros_like(original_np)
-        mask_vis[final_mask_big] = [255, 70, 70]
-        mask_overlay = cv2.addWeighted(original_np, 0.78, mask_vis, 0.22, 0)
-        mask_pil = Image.fromarray(mask_overlay)
-
-        return result_pil, comparison_pil, mask_pil
-
+            raise HTTPException(400, detail=str(ve))
     except Exception as e:
-        logger.error(f"Error: {str(e)}", exc_info=True)
-        raise gr.Error(f"Processing failed: {str(e)}. Try smaller image.")
-
-with gr.Blocks(title="Make Me Bald 🧑‍🦲", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# Realistic Bald Maker 🔥")
-    gr.Markdown("Upload face photo → get bald version with natural skin blending. Ears protected, no weird halos!")
-
-    with gr.Row():
-        input_img = gr.Image(type="pil", label="Your Photo", sources=["upload", "webcam"])
-        output_bald = gr.Image(label="Bald Version")
-    
-    with gr.Row():
-        comparison = gr.Image(label="Before vs After")
-        mask_overlay = gr.Image(label="Hair Mask Overlay (red = removed area)")
-
-    btn = gr.Button("Make Bald 😎", variant="primary")
-
-    btn.click(
-        fn=make_realistic_bald,
-        inputs=input_img,
-        outputs=[output_bald, comparison, mask_overlay],
-        api_name="make_bald"
-    )
-
-    gr.Examples(
-        examples=[["example1.jpg"], ["example2.jpg"]],  # agar examples folder mein daale to
-        inputs=input_img,
-        label="Try these examples"
-    )
-
-    gr.Markdown("""
-    **Tips:**
-    - Best results on clear front-facing photos.
-    - Large images auto-resized for speed (then upscaled).
-    - If no hair detected → try another photo.
-    """)
+        raise HTTPException(500, detail=f"Processing failed: {str(e)}")
 
-# NO demo.launch() here — HF Spaces handles it automatically!
+@app.get("/")
+def home():
+    return {"message": "Bald banne aaya? POST /make-bald/ pe image daal!"}