from fastapi import FastAPI, UploadFile, File, HTTPException from fastapi.responses import StreamingResponse import io import logging import torch import gradio as gr # Logging setup (Space logs mein clear dikhega) logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # ── Global variables (model loading) ─────────────────────────────────────────── device = "cuda" if torch.cuda.is_available() else "cpu" processor = None model = None hair_class_id = 13 ear_class_ids = [7, 8] @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!") app = FastAPI( title="Make Me Bald API 😎", description="Upload photo → Get realistic bald version! 🧑‍🦲", version="1.0" ) def make_realistic_bald(image_bytes: bytes) -> bytes: """ Main bald processing function - takes bytes, returns bald image bytes (Updated to reduce halo/shadow artifacts) """ 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) original_bgr = cv2.cvtColor(original_np, cv2.COLOR_RGB2BGR) logger.info(f"Processing image: {orig_w}x{orig_h}") # Resize for processing (speed + memory) 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: 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) # 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 logic (same as yours) 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) ears_protected = np.zeros_like(hair_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 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), :] ) # Sharper mask: reduced blur kernel_s = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7)) hair_mask_final = cv2.morphologyEx(hair_mask_final, cv2.MORPH_CLOSE, kernel_s, iterations=1) hair_mask_final = cv2.dilate(hair_mask_final, kernel_s, iterations=1) blurred = cv2.GaussianBlur(hair_mask_final.astype(np.float32), (5, 5), 1.0) hair_mask_final = (blurred > 0.45).astype(np.uint8) # higher threshold → sharper edges kernel_edge = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3)) 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 (same logic) final_mask = hair_mask_final.copy() use_extended_mask = False if hair_pixels > 380000: logger.info("Very large hair area → using extended mask") use_extended_mask = True big_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (21, 21)) 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 extended = cv2.morphologyEx(extended, cv2.MORPH_CLOSE, kernel_s, iterations=1) extended[int(working_h * 0.75):, :] = 0 final_mask = extended # Inpainting - reduced radius for less halo radius = 8 if use_extended_mask or hair_pixels > 300000 else 5 inpaint_flag = cv2.INPAINT_TELEA # better boundary preservation logger.info(f"Inpainting with radius={radius}") inpainted_bgr = cv2.inpaint(working_bgr, final_mask * 255, inpaintRadius=radius, flags=inpaint_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] # Light color matching (reduced strength) if use_extended_mask or hair_pixels > 200000: logger.info("Applying light skin color correction") regions = [(0.20, 0.35, 0.35, 0.65), (0.35, 0.50, 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) strength = 0.45 # reduced to avoid artifacts 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 # Sharpen to remove residual blur/halo sharpen_kernel = np.array([[-1,-1,-1], [-1,9,-1], [-1,-1,-1]]) result_small = cv2.filter2D(result_small, -1, sharpen_kernel) # Upscale if resized 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 # Convert back to bytes _, 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") 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 (HF Spaces compatibility ke liye) app = gr.mount_gradio_app(app, gr_interface, path="/") if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=7860)