devbucky/rhinoplasty
0
1import gradio as gr2import cv23import numpy as np4from PIL import Image5import mediapipe as mp6import google.genai as genai7import os8import replicate9import base6410from io import BytesIO11 12# 30 user-provided landmark indices outlining the nose13USER_NOSE_POINTS = [14 107, 9, 336, 285, 413, 464, 465, 412, 399, 456,15 420, 429, 279, 327, 326, 2, 97, 98, 64, 129,16 49, 131, 198, 217, 114, 188, 245, 193, 55, 107]17 18def detect_face_first(image_rgb):19 mp_face_detection = mp.solutions.face_detection20 21 with mp_face_detection.FaceDetection(22 model_selection=1,23 min_detection_confidence=0.324 ) as face_detection:25 results = face_detection.process(image_rgb)26 27 if not results.detections:28 return False, 0, "No face detected in the image."29 30 num_faces = len(results.detections)31 if num_faces > 1:32 return False, num_faces, f"Multiple faces detected! Found {num_faces} faces. Please upload an image with ONLY ONE person."33 34 return True, 1, "Face detected successfully."35 36def is_profile_face(landmarks, image_width, image_height):37 lm = landmarks.landmark38 left_eye_indices = [33, 133, 160, 159, 158]39 right_eye_indices = [362, 263, 385, 386, 387]40 41 left_eye_visible = 042 right_eye_visible = 043 44 for idx in left_eye_indices:45 if idx < len(lm):46 if lm[idx].visibility > 0.5:47 left_eye_visible += 148 49 for idx in right_eye_indices:50 if idx < len(lm):51 if lm[idx].visibility > 0.5:52 right_eye_visible += 153 54 total_visible = left_eye_visible + right_eye_visible55 if total_visible < 3:56 return True57 58 visibility_ratio = abs(left_eye_visible - right_eye_visible) / max(total_visible, 1)59 return visibility_ratio > 0.660 61def generate_nose_mask(image: Image.Image):62 """Generate nose mask. Returns (mask_image, info_message)"""63 if image is None:64 return None, "No image uploaded."65 66 img_bgr = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)67 img_rgb = np.array(image)68 h, w = img_bgr.shape[:2]69 70 if w < 100 or h < 100:71 return None, "Image is too small. Minimum size is 100x100 pixels."72 73 face_detected, num_faces, detection_msg = detect_face_first(img_rgb)74 if not face_detected:75 return None, detection_msg76 77 mp_face_mesh = mp.solutions.face_mesh78 results = None79 80 with mp_face_mesh.FaceMesh(81 static_image_mode=True,82 refine_landmarks=True,83 max_num_faces=1,84 min_detection_confidence=0.385 ) as face_mesh:86 results = face_mesh.process(img_rgb)87 88 if not results or not results.multi_face_landmarks:89 with mp_face_mesh.FaceMesh(90 static_image_mode=True,91 refine_landmarks=False,92 max_num_faces=1,93 min_detection_confidence=0.194 ) as face_mesh_relaxed:95 results = face_mesh_relaxed.process(img_rgb)96 97 if not results or not results.multi_face_landmarks:98 with mp_face_mesh.FaceMesh(99 static_image_mode=True,100 refine_landmarks=False,101 max_num_faces=1,102 min_detection_confidence=0.01103 ) as face_mesh_ultra_relaxed:104 results = face_mesh_ultra_relaxed.process(img_rgb)105 106 if not results or not results.multi_face_landmarks:107 with mp_face_mesh.FaceMesh(108 static_image_mode=False,109 refine_landmarks=False,110 max_num_faces=1,111 min_detection_confidence=0.01,112 min_tracking_confidence=0.01113 ) as face_mesh_video:114 results = face_mesh_video.process(img_rgb)115 116 if not results or not results.multi_face_landmarks:117 return None, "Face detected but could not extract facial landmarks."118 119 lm = results.multi_face_landmarks[0]120 is_profile = is_profile_face(lm, w, h)121 face_type = "side profile" if is_profile else "frontal face"122 123 nose_points = USER_NOSE_POINTS124 min_points_required = 3 if is_profile else 5125 126 pts = []127 for idx in nose_points:128 if idx < len(lm.landmark):129 landmark = lm.landmark[idx]130 x_px = int(landmark.x * w)131 y_px = int(landmark.y * h)132 133 if is_profile:134 if -10 <= x_px < w + 10 and -10 <= y_px < h + 10:135 x_px = max(0, min(w - 1, x_px))136 y_px = max(0, min(h - 1, y_px))137 pts.append((x_px, y_px))138 else:139 if 0 <= x_px < w and 0 <= y_px < h:140 pts.append((x_px, y_px))141 142 if len(pts) < min_points_required:143 fallback_indices = [1, 2, 4, 5, 6, 19, 94, 168, 195]144 for idx in fallback_indices:145 if idx < len(lm.landmark):146 landmark = lm.landmark[idx]147 x_px = int(landmark.x * w)148 y_px = int(landmark.y * h)149 if -10 <= x_px < w + 10 and -10 <= y_px < h + 10:150 x_px = max(0, min(w - 1, x_px))151 y_px = max(0, min(h - 1, y_px))152 pts.append((x_px, y_px))153 154 if len(pts) < 3:155 return None, f"Unable to generate mask: Only {len(pts)} nose landmarks detected."156 157 hull = cv2.convexHull(np.array(pts, dtype=np.int32))158 mask = np.zeros((h, w), dtype=np.uint8)159 cv2.fillPoly(mask, [hull], 255)160 161 kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))162 mask = cv2.dilate(mask, kernel, iterations=2)163 mask = cv2.GaussianBlur(mask, (5, 5), 0)164 _, mask = cv2.threshold(mask, 128, 255, cv2.THRESH_BINARY)165 166 success_msg = f"✓ Detected {face_type} with {len(pts)} nose landmarks."167 return Image.fromarray(mask), success_msg168 169def generate_canny_edge(image: Image.Image, low_threshold=100, high_threshold=200):170 """Generate Canny edge map. Returns (canny_image, info_message)"""171 if image is None:172 return None, "No image provided."173 174 try:175 img_array = np.array(image)176 177 if len(img_array.shape) == 3:178 gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)179 else:180 gray = img_array181 182 blurred = cv2.GaussianBlur(gray, (5, 5), 1.4)183 canny = cv2.Canny(blurred, low_threshold, high_threshold)184 canny_3channel = cv2.cvtColor(canny, cv2.COLOR_GRAY2RGB)185 canny_pil = Image.fromarray(canny_3channel)186 187 return canny_pil, "✓ Canny edge map generated."188 189 except Exception as e:190 return None, f"Error: {str(e)}"191 192def generate_gemini_prompt(image: Image.Image, user_modification: str, api_key: str):193 """Use Gemini to generate detailed prompt. Returns (prompt, status)"""194 if image is None:195 return None, "No image provided."196 197 if not user_modification or user_modification.strip() == "":198 return None, "Please describe desired nose modifications."199 200 if not api_key or api_key.strip() == "":201 return None, "Please provide Gemini API key."202 203 try:204 genai.configure(api_key=api_key)205 model = genai.GenerativeModel('gemini-2.0-flash-exp')206 207 system_prompt = f"""You are an expert AI prompt engineer for rhinoplasty simulation using Stable Diffusion inpainting.208The user wants: "{user_modification}"209Analyze this face image and create an OPTIMAL BALANCED PROMPT following this exact structure:210**STRUCTURE:**211[NOSE MODIFICATION DETAILS] + [STRICT PRESERVATION] + [QUALITY KEYWORDS]212**REQUIREMENTS:**2131. **Modification Section**: Describe the nose changes in natural, specific, visual terms (not medical jargon)214 - Be concrete: "reduce bridge height by making it straighter and lower"215 - Use visual descriptors with CAPS for emphasis on key structural features216 - Reference the user's request: {user_modification}217 218 **CRITICAL TIP SHAPE RULES:**219 - If user wants "button nose", "rounded tip", or "fuller tip": USE CAPS like "ROUNDED BULBOUS tip", "FULLER WIDER tip (NOT narrow, NOT pointed)"220 - If user wants "narrow tip", "refined tip", "pointed tip": USE "narrower, more refined tip"221 - For "upturned": emphasize "lifted upward" but also specify tip roundness if needed222 - For "Roman": emphasize "narrow pointed tip extending downward"223 - For "Nubian": emphasize "WIDER FULLER ROUNDED tip with expanded width"224 225 **AVOID conflicting words**: Don't use "refined" or "delicate" when describing tips that should be fuller/rounder2262. **Preservation Section**: Use STRONG language to preserve identity227 - MUST include: "CRITICAL: preserve exact skin texture, skin tone, facial structure, eyes, lips, cheeks, expression, lighting, shadows, and background completely unchanged"228 - Use words like: "identical", "exact same", "preserve", "keep unchanged"2293. **Quality Section**: Add photorealism keywords230 - Include: "photorealistic, high resolution, professional photography, natural blending, 8k, sharp focus, detailed skin texture"231**EXAMPLE OUTPUTS:**232For button/rounded nose:233"Nose modification: create a button nose with SHORT COMPACT structure, ROUNDED BULBOUS TIP that is FULLER and WIDER (explicitly NOT narrow or pointed), low gentle bridge with soft slope, maintain the rounded ball-like quality of the tip. CRITICAL: preserve exact skin texture, skin tone, facial structure, eyes, lips, expression, lighting, and background completely unchanged. Photorealistic, high resolution, professional photography, natural blending, 8k, sharp focus, detailed skin texture."234For refined/narrow nose:235"Nose refinement: reduce the bridge height for a straighter profile, narrow and refine the nasal tip to be more delicate and pointed, slightly lift the tip angle. CRITICAL: preserve exact skin texture, skin tone, facial structure, eyes, lips, expression, lighting, and background completely unchanged. Photorealistic, high resolution, professional photography, natural blending, 8k, sharp focus, detailed skin texture."236**OUTPUT RULES:**237- Maximum 3 sentences total238- Natural language (avoid medical terms like "dorsum", "alar base", "osteotomy")239- Focus on VISUAL outcomes, not surgical procedures240- Use CAPS for critical structural features that need emphasis241- Balance specificity with clarity242- Do NOT use bullet points or lists, use flowing prose243Generate the prompt now:"""244 245 response = model.generate_content([system_prompt, image])246 generated_prompt = response.text.strip()247 248 return generated_prompt, "✓ Gemini analysis completed."249 250 except Exception as e:251 return None, f"Gemini error: {str(e)}"252 253def image_to_data_uri(image: Image.Image):254 """Convert PIL Image to data URI for Replicate"""255 buffered = BytesIO()256 image.save(buffered, format="PNG")257 img_str = base64.b64encode(buffered.getvalue()).decode()258 return f"data:image/png;base64,{img_str}"259 260# ============================================261# MODEL-SPECIFIC FUNCTIONS (3 MODELS ONLY)262# ============================================263 264def run_sdxl_inpainting(original_image, mask_image, prompt, replicate_api_key):265 """Model 1: Flux Fill Dev - Uses: image + mask + prompt"""266 try:267 os.environ["REPLICATE_API_TOKEN"] = replicate_api_key.strip()268 269 w, h = original_image.size270 mask_image = mask_image.resize((w, h))271 272 image_uri = image_to_data_uri(original_image)273 mask_uri = image_to_data_uri(mask_image)274 275 output = replicate.run(276 "black-forest-labs/flux-fill-dev",277 input={278 "image": image_uri,279 "mask": mask_uri,280 "prompt": prompt,281 "guidance": 33,282 "steps": 28,283 "output_format": "png",284 "output_quality": 95285 }286 )287 288 if output:289 import requests290 if isinstance(output, list):291 image_url = str(output[0]) if hasattr(output[0], '__str__') else output[0].url292 else:293 image_url = str(output)294 295 response = requests.get(image_url)296 result_image = Image.open(BytesIO(response.content))297 return result_image, "✓ Flux Fill Dev completed!"298 else:299 return None, "No output from model."300 301 except Exception as e:302 return None, f"Error: {str(e)}"303 304def run_flux_fill_pro(original_image, mask_image, prompt, replicate_api_key):305 """Model 2: Flux Fill Pro - Uses: image + mask + prompt (requires min 256x256)"""306 try:307 os.environ["REPLICATE_API_TOKEN"] = replicate_api_key.strip()308 309 w, h = original_image.size310 311 if w < 256 or h < 256:312 scale_factor = max(256 / w, 256 / h)313 new_w = int(w * scale_factor)314 new_h = int(h * scale_factor)315 original_image = original_image.resize((new_w, new_h))316 mask_image = mask_image.resize((new_w, new_h))317 status_msg = f"✓ Flux Fill Pro completed! (Upscaled from {w}x{h} to {new_w}x{new_h} to meet 256px minimum)"318 else:319 mask_image = mask_image.resize((w, h))320 status_msg = "✓ Flux Fill Pro completed!"321 322 image_uri = image_to_data_uri(original_image)323 mask_uri = image_to_data_uri(mask_image)324 325 output = replicate.run(326 "black-forest-labs/flux-fill-pro",327 input={328 "image": image_uri,329 "mask": mask_uri,330 "prompt": prompt,331 "steps": 25,332 "guidance": 3.0,333 }334 )335 336 if output:337 import requests338 if isinstance(output, list):339 image_url = str(output[0]) if hasattr(output[0], '__str__') else output[0].url340 else:341 image_url = str(output)342 343 response = requests.get(image_url)344 result_image = Image.open(BytesIO(response.content))345 return result_image, status_msg346 else:347 return None, "No output from model."348 349 except Exception as e:350 return None, f"Error: {str(e)}"351 352def run_bria_genfill(original_image, mask_image, prompt, replicate_api_key):353 """Model 3: Bria GenFill - Uses: image + mask + prompt"""354 try:355 os.environ["REPLICATE_API_TOKEN"] = replicate_api_key.strip()356 357 w, h = original_image.size358 mask_image = mask_image.resize((w, h))359 360 image_uri = image_to_data_uri(original_image)361 mask_uri = image_to_data_uri(mask_image)362 363 enhanced_prompt = f"{prompt}, seamlessly blended with facial features, matching existing skin tone and texture, natural transitions, photorealistic medical photography quality"364 365 output = replicate.run(366 "bria/genfill",367 input={368 "image": image_uri,369 "mask": mask_uri,370 "prompt": enhanced_prompt,371 }372 )373 374 if output:375 import requests376 if isinstance(output, list):377 image_url = str(output[0]) if hasattr(output[0], '__str__') else output[0].url378 else:379 image_url = str(output)380 381 response = requests.get(image_url)382 result_image = Image.open(BytesIO(response.content))383 return result_image, "✓ Bria GenFill completed!"384 else:385 return None, "No output from model."386 387 except Exception as e:388 return None, f"Error: {str(e)}"389 390# ============================================391# PIPELINE392# ============================================393 394def rhinoplasty_pipeline(image, modification_text, use_gemini_enhancer, gemini_api_key, replicate_api_key, model_choice):395 """Complete rhinoplasty simulation pipeline"""396 397 status_log = []398 status_log.append(f"🎯 Selected Model: {model_choice}\n")399 400 if image is None:401 return None, "No image uploaded."402 403 # Step 1: Generate nose mask404 status_log.append("Step 1: Generating nose mask...")405 mask, mask_msg = generate_nose_mask(image)406 status_log.append(f" {mask_msg}")407 408 if mask is None:409 return None, "\n".join(status_log)410 411 # Step 2: Generate canny edge (kept for future use)412 status_log.append("\nStep 2: Generating canny edge map (for future ControlNet use)...")413 canny, canny_msg = generate_canny_edge(image)414 status_log.append(f" {canny_msg}")415 416 # Step 3: Prepare prompt417 if use_gemini_enhancer:418 status_log.append("\nStep 3: Analyzing face with Gemini...")419 prompt, gemini_msg = generate_gemini_prompt(image, modification_text, gemini_api_key)420 status_log.append(f" {gemini_msg}")421 422 if prompt is None:423 return None, "\n".join(status_log)424 425 status_log.append(f"\nEnhanced Prompt:\n{prompt}")426 else:427 status_log.append("\nStep 3: Using structured prompt (Gemini disabled)...")428 prompt = f"""Nose modification: {modification_text}, creating a natural and harmonious appearance. CRITICAL: preserve exact skin texture, skin tone, facial structure, eyes, lips, cheeks, eyebrows, expression, lighting, shadows, hair, and background completely unchanged. Keep the person's identity identical. Photorealistic, high resolution, professional photography, natural blending, 8k, sharp focus, detailed skin texture, natural lighting."""429 status_log.append(f" ✓ Structured prompt prepared.")430 status_log.append(f"\nPrompt Used:\n{prompt}")431 432 # Step 4: Run selected model433 status_log.append(f"\n\nStep 4: Running {model_choice}...")434 status_log.append(" (This may take 20-60 seconds...)")435 436 if model_choice == "Flux Fill Dev (Recommended)":437 result, result_msg = run_sdxl_inpainting(image, mask, prompt, replicate_api_key)438 elif model_choice == "Flux Fill Pro (Premium)":439 result, result_msg = run_flux_fill_pro(image, mask, prompt, replicate_api_key)440 elif model_choice == "Bria GenFill (Budget)":441 result, result_msg = run_bria_genfill(image, mask, prompt, replicate_api_key)442 else:443 return None, "Invalid model selection."444 445 status_log.append(f" {result_msg}")446 447 if result is None:448 return None, "\n".join(status_log)449 450 status_log.append(f"\n💰 Estimated cost: ~$0.01-0.03")451 452 return result, "\n".join(status_log)453 454# ============================================455# GRADIO INTERFACE456# ============================================457 458default_gemini_key = os.getenv("GEMINI_API_KEY", "")459default_replicate_key = os.getenv("REPLICATE_API_TOKEN", "")460 461iface = gr.Interface(462 fn=rhinoplasty_pipeline,463 api_name="predict",464 inputs=[465 gr.Image(type="pil", label="Upload Face Image"),466 gr.Textbox(467 label="Desired Nose Modification",468 placeholder="Examples: 'button nose with ROUNDED FULLER tip' | 'Roman nose with dorsal hump' | 'upturned nose'",469 lines=2,470 info="Use CAPS for emphasis on key features like 'ROUNDED', 'FULLER', 'WIDER'"471 ),472 gr.Checkbox(473 label="Use Gemini Prompt Enhancer",474 value=True,475 info="Enable AI-powered prompt enhancement (requires Gemini API key)"476 ),477 gr.Textbox(478 label="Gemini API Key (Optional - only if enhancer enabled)",479 placeholder="Get free at: https://makersuite.google.com/app/apikey",480 value=default_gemini_key,481 type="password",482 lines=1483 ),484 gr.Textbox(485 label="Replicate API Key (Required)",486 placeholder="Get at: https://replicate.com/account/api-tokens",487 value=default_replicate_key,488 type="password",489 lines=1490 ),491 gr.Dropdown(492 label="Select AI Model",493 choices=[494 "Flux Fill Dev (Recommended)",495 "Flux Fill Pro (Premium)",496 "Bria GenFill (Budget)"497 ],498 value="Flux Fill Dev (Recommended)",499 info="Choose your preferred model"500 )501 ],502 outputs=[503 gr.Image(label="Result"),504 gr.Textbox(label="Processing Log", lines=18)505 ],506 title="🔬Simulator",507 description="""508 **Simulation**509 """,510 examples=[],511)512 513if __name__ == "__main__":514 # Launch with explicit settings for Hugging Face Spaces515 iface.launch(516 server_name="0.0.0.0",517 server_port=7860,518 show_error=True,519 quiet=False520 )