Update app.py

app.py

@@ -1,122 +1,121 @@
-from fastapi import FastAPI, UploadFile, File, HTTPException
-from fastapi.responses import StreamingResponse
-import io
-import logging
+import gradio as gr
 import torch
 import cv2
 import numpy as np
 from PIL import Image
+import logging
 from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
-import gradio as gr
-import os
-# Logging setup
+
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
-# Global variables
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
-processor = None
-model = None
+logger.info(f"Using device: {device}")
+
+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]
-app = FastAPI(
-    title="Make Me Bald API 😎",
-    description="Upload photo → Get realistic bald version! 🧑‍🦲",
-    version="1.0"
-)
-@app.on_event("startup")
-async def startup_event():
-    global processor, model
-    logger.info(f"Loading SegFormer model on {device}...")
-    processor = SegformerImageProcessor.from_pretrained("jonathandinu/face-parsing")
-    model = SegformerForSemanticSegmentation.from_pretrained("jonathandinu/face-parsing")
-    model.to(device)
-    model.eval()
-    logger.info("Model loaded successfully!")
-def make_realistic_bald(image_bytes: bytes) -> bytes:
-    """
-    Tera original processor code - bytes input leke bald image bytes return karta hai
-    (local version se direct replace kiya, halo reduce wale changes bhi rakhe)
-    """
+
+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:
-        # Convert bytes to PIL Image
-        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-        orig_w, orig_h = image.size
-        original_np = np.array(image)
+        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 image: {orig_w}x{orig_h}")
-        # ── SPEED FIX: Resize large images for processing ─────────────────────────────
-        MAX_PROCESS_DIM = 2048 # ← adjust 1536–2304 if needed
+        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"Large image → downscaling to max {MAX_PROCESS_DIM}px...")
+            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), interpolation=cv2.INTER_AREA)
+            working_np = cv2.resize(original_np, (working_w, working_h), cv2.INTER_AREA)
             working_bgr = cv2.cvtColor(working_np, cv2.COLOR_RGB2BGR)
-        # Segmentation
+
         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)
-        # Ear protection (same as your original)
+
         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
         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), :]
             )
-        # Smoothing (tera original wala – halo kam karne ke liye blur adjust kiya)
+
         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 detected (resized): {hair_pixels:,}")
-        # Extended mask (tera original)
+        logger.info(f"Hair pixels (resized): {hair_pixels:,}")
+
         final_mask = hair_mask_final.copy()
         use_extended_mask = False
         if hair_pixels > 380000:
-            logger.info(" VERY LARGE hair → creating extended mask...")
+            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)
@@ -125,50 +124,44 @@ def make_realistic_bald(image_bytes: bytes) -> bytes:
             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
-            extra_pixels = np.sum(extended)
-            logger.info(f" Extended mask size: {extra_pixels:,} (+{extra_pixels - hair_pixels:,})")
             final_mask = extended
-        # Inpainting - tera original radii logic
+
         if use_extended_mask or hair_pixels > 420000:
-            radius = 18
-            inpaint_flag = cv2.INPAINT_TELEA
+            radius, flag = 18, cv2.INPAINT_TELEA
         elif hair_pixels > 220000:
-            radius = 15
-            inpaint_flag = cv2.INPAINT_TELEA
+            radius, flag = 15, cv2.INPAINT_TELEA
         else:
-            radius = 10
-            inpaint_flag = cv2.INPAINT_NS
-        logger.info(f" Inpainting with radius={radius}")
-        inpainted_bgr = cv2.inpaint(working_bgr, final_mask * 255,
-                                   inpaintRadius=radius, flags=inpaint_flag)
+            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]
-        # Color matching (tera original)
+
         if use_extended_mask or hair_pixels > 280000:
-            logger.info(" Applying skin color correction...")
-            regions = [
-                (0.18, 0.30, 0.34, 0.66),
-                (0.32, 0.47, 0.35, 0.65)
-            ]
+            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 = int(working_h * y1r)
-                y2 = int(working_h * y2r)
-                x1 = int(working_w * x1r)
-                x2 = int(working_w * x2r)
+                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)
@@ -179,71 +172,68 @@ def make_realistic_bald(image_bytes: bytes) -> bytes:
                     diff = target_color - current_mean
                     corrected = np.clip(bald_area + diff * strength, 0, 255).astype(np.uint8)
                     result_small[final_mask == 1] = corrected
-        # Final blending (tera original)
+
         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
+                result_small[final_mask == 1], 0.65, blurred_bald[final_mask == 1], 0.35, 0
             )
-        # Upscale if resized
+
         if scale_factor < 1.0:
-            logger.info("Upscaling final result to original size...")
+            logger.info("Upscaling to original size")
             result = cv2.resize(result_small, (orig_w, orig_h), interpolation=cv2.INTER_LANCZOS4)
         else:
             result = result_small
-        # Convert back to bytes (for API return)
-        _, buffer = cv2.imencode('.jpg', cv2.cvtColor(result, cv2.COLOR_RGB2BGR), [int(cv2.IMWRITE_JPEG_QUALITY), 92])
-        return buffer.tobytes()
-    except Exception as e:
-        logger.error(f"Bald processing failed: {str(e)}", exc_info=True)
-        raise ValueError(f"Processing error: {str(e)}")
-@app.post("/make-bald/")
-async def bald_endpoint(file: UploadFile = File(...)):
-    logger.info("=== REQUEST AAYI /make-bald/ PE ===")
-    logger.info(f"Filename: {file.filename} | Content-Type: {file.content_type} | Size: {file.size / 1024:.2f} KB")
-    if not file.content_type.startswith("image/"):
-        raise HTTPException(status_code=400, detail="Sirf image file upload kar bhai! (jpeg/png etc.)")
-    try:
-        contents = await file.read()
-        logger.info(f"Image read successful, size: {len(contents) / 1024:.2f} KB")
-        bald_bytes = make_realistic_bald(contents)
-        logger.info(f"Bald processing done, output size: {len(bald_bytes) / 1024:.2f} KB")
-        bald_io = io.BytesIO(bald_bytes)
-        bald_io.seek(0)
-        return StreamingResponse(
-            bald_io,
-            media_type="image/jpeg",
-            headers={"Content-Disposition": "attachment; filename=bald_version.jpg"}
-        )
-    except ValueError as ve:
-        error_detail = str(ve).strip()
-        logger.warning(f"ValueError: {error_detail}")
-        if "NO_HAIR" in error_detail.upper() or "NO_HAIR_DETECTED" in error_detail.upper():
-            raise HTTPException(status_code=400, detail="NO_HAIR_DETECTED")
-        raise HTTPException(status_code=400, detail=error_detail or "Processing mein kuch galat hua")
+
+        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
+
     except Exception as e:
-        logger.error(f"Unexpected error: {str(e)}", exc_info=True)
-        raise HTTPException(status_code=500, detail=f"Server error: {str(e)}")
-@app.get("/")
-def home():
-    return {
-        "message": "Bald banne aaya? 😏",
-        "how_to_use": "POST request bhejo /make-bald/ pe with form-data key 'file' aur image attach karo.",
-        "example": "curl -X POST -F 'file=@your_photo.jpg' https://seniordev22-space.hf.space/make-bald/ -o bald.jpg"
-    }
-# Gradio dummy for HF Spaces
-def dummy_fn():
-    return "API chal raha hai! cURL ya Postman se /make-bald/ pe POST karo."
-gr_interface = gr.Interface(
-    fn=dummy_fn,
-    inputs=None,
-    outputs="text",
-    title="Make Me Bald API 🧑‍🦲",
-    description="Ye sirf info page hai. Actual bald banane ke liye:\n\ncurl -X POST -F 'file=@photo.jpg' https://seniordev22-space.hf.space/make-bald/ -o bald.jpg"
-)
-# Mount Gradio on root path
-app = gr.mount_gradio_app(app, gr_interface, path="/")
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+        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.
+    """)
+
+# NO demo.launch() here — HF Spaces handles it automatically!