import gradio as gr import cv2 import numpy as np import requests import os import tempfile # --- Configuration for Model Files --- # IMPORTANT: These files must be present in your Hugging Face Space repository. # 1. Haar Cascade for face detection (comes with OpenCV) FACE_CASCADE_PATH = cv2.data.haarcascades + 'haarcascade_frontalface_default.xml' # 2. LBF Facial Landmark Model (lbfmodel.yaml) # You need to download 'lbfmodel.yaml' and add it to your Hugging Face Space. # A common source is: https://github.com/kurnianggoro/GSOC2017/blob/master/data/lbfmodel.yaml LBF_MODEL_PATH = "lbfmodel.yaml" # --- Core Image Processing and AI Editing Function --- def edit_image_with_ai(input_image_np_rgb, stability_api_key, prompt_text="perfect white teeth"): """ Processes an input image to detect a mouth, then uses Stability AI to edit it. Args: input_image_np_rgb (numpy.ndarray): The input image in RGB format. stability_api_key (str): The Stability AI API key. prompt_text (str): The prompt for the AI image editing. Returns: tuple: (original_image, edited_image, status_message) original_image is the input image. edited_image is the AI-edited image (or None on failure). status_message provides feedback on the process. """ if input_image_np_rgb is None: return None, None, "Please upload an image or use the webcam." if not stability_api_key: return input_image_np_rgb, None, "⚠️ Please enter your Stability AI API Key." if not prompt_text: return input_image_np_rgb, None, "⚠️ Please enter a prompt for the edit." try: # Convert RGB (from Gradio) to BGR (for OpenCV) img_bgr = cv2.cvtColor(input_image_np_rgb, cv2.COLOR_RGB2BGR) gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY) mouth_mask = np.zeros_like(gray) # Initialize an empty (black) mask # --- Load Face Detector and Facemark Model --- face_cascade = cv2.CascadeClassifier(FACE_CASCADE_PATH) if face_cascade.empty(): # This error is unlikely if OpenCV is correctly installed. raise gr.Error("Fatal Error: Could not load face cascade model. Check OpenCV installation.") if not os.path.exists(LBF_MODEL_PATH): raise gr.Error( f"Error: LBF Facial Landmark Model ('{LBF_MODEL_PATH}') not found. " "Please download it and add it to your Hugging Face Space's root directory." ) facemark = cv2.face.createFacemarkLBF() try: facemark.loadModel(LBF_MODEL_PATH) except cv2.error as e: raise gr.Error(f"Error loading LBF model from '{LBF_MODEL_PATH}': {e}. Ensure the file is valid.") # --- Detect Faces --- faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(50, 50)) # Increased minSize if len(faces) == 0: return input_image_np_rgb, None, "No faces detected in the image. Try a clearer image or different pose." # --- Fit Landmarks and Create Mouth Mask --- ok, landmarks_all_faces = facemark.fit(img_bgr, faces) num_mouths_masked = 0 if ok: for landmarks_one_face in landmarks_all_faces: shape = landmarks_one_face[0].astype(np.int32) # Mouth points (indices 48-59 for outer mouth in 68-point model) mouth_points = shape[48:60] if len(mouth_points) > 2: # Need at least 3 points for a hull mouth_hull = cv2.convexHull(mouth_points) # Dilate the mask slightly to ensure coverage cv2.drawContours(mouth_mask, [mouth_hull], -1, (255), -1) num_mouths_masked += 1 if num_mouths_masked > 0: # Dilate the final mask slightly to ensure edges are included for inpainting kernel = np.ones((5,5),np.uint8) mouth_mask = cv2.dilate(mouth_mask, kernel, iterations = 1) else: return input_image_np_rgb, None, "Facial landmark detection failed. Could not identify facial features." if num_mouths_masked == 0: return input_image_np_rgb, None, "Could not identify or create a mask for any mouth regions." _, mouth_mask_binary = cv2.threshold(mouth_mask, 127, 255, cv2.THRESH_BINARY) # --- Call Stability AI API for Inpainting --- # Use temporary files for sending image and mask to the API with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_image_file, \ tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_mask_file: cv2.imwrite(tmp_image_file.name, img_bgr) # OpenCV saves in BGR cv2.imwrite(tmp_mask_file.name, mouth_mask_binary) tmp_image_path = tmp_image_file.name tmp_mask_path = tmp_mask_file.name edited_image_np_rgb = None # Initialize status_message = "Calling Stability AI API..." try: response = requests.post( "https://api.stability.ai/v2beta/stable-image/edit/inpaint", headers={ "authorization": f"Bearer {stability_api_key}", "accept": "image/*" # Expecting an image in response }, files={ "image": open(tmp_image_path, "rb"), "mask": open(tmp_mask_path, "rb"), }, data={ "prompt": prompt_text, "output_format": "png", # Request PNG output # You can add other parameters like 'seed', 'strength' if desired # e.g., "strength": 0.75, }, timeout=90 # Increased timeout for potentially slow API responses ) if response.status_code == 200: edited_image_bytes = response.content nparr = np.frombuffer(edited_image_bytes, np.uint8) # Decode image, OpenCV reads as BGR edited_image_cv_bgr = cv2.imdecode(nparr, cv2.IMREAD_COLOR) if edited_image_cv_bgr is not None: # Convert BGR to RGB for Gradio display edited_image_np_rgb = cv2.cvtColor(edited_image_cv_bgr, cv2.COLOR_BGR2RGB) status_message = "✅ Image edited successfully!" else: status_message = "⚠️ Stability API returned an invalid image format." return input_image_np_rgb, None, status_message else: # Try to parse JSON error, otherwise use text try: error_detail = response.json() # Stability AI errors are often in a list under 'errors' error_messages = error_detail.get('errors', [response.text]) status_message = f"⚠️ Stability API Error ({response.status_code}): {error_messages[0]}" except requests.exceptions.JSONDecodeError: status_message = f"⚠️ Stability API Error ({response.status_code}): {response.text}" return input_image_np_rgb, None, status_message except requests.exceptions.Timeout: status_message = "⚠️ Stability API request timed out. The server might be busy or the image too large. Please try again later." return input_image_np_rgb, None, status_message except requests.exceptions.RequestException as e: status_message = f"⚠️ Request to Stability API failed: {e}" return input_image_np_rgb, None, status_message except Exception as e: status_message = f"⚠️ An unexpected error occurred during API processing: {e}" return input_image_np_rgb, None, status_message finally: os.remove(tmp_image_path) # Clean up temporary file os.remove(tmp_mask_path) # Clean up temporary file return input_image_np_rgb, edited_image_np_rgb, status_message except gr.Error as e: # Catch Gradio specific errors raised intentionally return input_image_np_rgb, None, str(e) # Display the gr.Error message except Exception as e: # Catch any other unexpected errors during CV processing return input_image_np_rgb, None, f"⚠️ An unexpected error occurred in processing: {e}" # --- Gradio Interface Definition --- with gr.Blocks(theme=gr.themes.Soft(primary_hue=gr.themes.colors.blue, secondary_hue=gr.themes.colors.sky)) as demo: gr.Markdown( "

👄 Your Best Smile With Keen Dental 😄

" "

Upload an image or use your webcam. The app detects faces, " " and uses AI to change your teeth based on your prompt.

" ) gr.Markdown( "**Instructions:**\n" "1. Enter your Stability AI API key (e.g., `sk-xxxxxxxx`). Get one from [Stability AI Developer Platform](https://platform.stability.ai/).\n" "2. Upload an image or use your webcam.\n" "3. Click 'Smile!'\n\n" ) with gr.Row(): api_key_input = gr.Textbox( label="🔑 Stability AI API Key", placeholder="Enter your sk-xxxxxxxx API key here", type="password", lines=1, elem_id="api_key_input" ) with gr.Row(): input_image = gr.Image(label="🖼️ Input Image", type="numpy", sources=["upload", "webcam"], height=400, interactive=True) output_image = gr.Image(label="✨ Edited Image", type="numpy", height=400, interactive=False) process_button = gr.Button("🚀 Smile!", variant="primary") status_output = gr.Textbox(label="ℹ️ Status", interactive=False, lines=2) process_button.click( fn=edit_image_with_ai, inputs=[input_image, api_key_input], outputs=[input_image, output_image, status_output], # Display original again for reference api_name="edit_mouth" # For API access if needed ) # To run this locally for testing (optional): if __name__ == "__main__": demo.launch(debug=True) # For Hugging Face Spaces, save this as 'app.py'. # Gradio will automatically call demo.launch().