Spaces:
Running
Running
| import json | |
| import os | |
| import time | |
| import uuid | |
| import tempfile | |
| from PIL import Image, ImageDraw, ImageFont | |
| import gradio as gr | |
| import base64 | |
| import mimetypes | |
| from google import genai | |
| from google.genai import types | |
| def save_binary_file(file_name, data): | |
| with open(file_name, "wb") as f: | |
| f.write(data) | |
| def generate(text, file_name, api_key, model="gemini-2.0-flash-exp"): | |
| # Initialize client using provided api_key (or fallback to env variable) | |
| client = genai.Client(api_key=(api_key.strip() if api_key and api_key.strip() != "" | |
| else os.environ.get("GEMINI_API_KEY"))) | |
| try: | |
| print("Uploading file to Gemini API...") | |
| files = [ client.files.upload(file=file_name) ] | |
| contents = [ | |
| types.Content( | |
| role="user", | |
| parts=[ | |
| types.Part.from_uri( | |
| file_uri=files[0].uri, | |
| mime_type=files[0].mime_type, | |
| ), | |
| types.Part.from_text(text=text), | |
| ], | |
| ), | |
| ] | |
| generate_content_config = types.GenerateContentConfig( | |
| temperature=0, # Lower temperature for more consistent, conservative results | |
| top_p=0.92, | |
| max_output_tokens=8192, | |
| response_modalities=["image", "text"], | |
| response_mime_type="text/plain", | |
| # Additional parameters to encourage subtle, natural results | |
| safety_settings=[ | |
| { | |
| "category": "HARM_CATEGORY_DANGEROUS_CONTENT", | |
| "threshold": "BLOCK_MEDIUM_AND_ABOVE" | |
| } | |
| ] | |
| ) | |
| text_response = "" | |
| image_path = None | |
| # Create a temporary file to potentially store image data | |
| with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp: | |
| temp_path = tmp.name | |
| print("Sending request to Gemini API...") | |
| # Add a timeout to prevent indefinite waiting | |
| start_time = time.time() | |
| max_wait_time = 60 # Maximum wait time in seconds | |
| try: | |
| stream = client.models.generate_content_stream( | |
| model=model, | |
| contents=contents, | |
| config=generate_content_config, | |
| ) | |
| for chunk in stream: | |
| # Check for timeout | |
| if time.time() - start_time > max_wait_time: | |
| print("Gemini API request timed out after", max_wait_time, "seconds") | |
| break | |
| if not chunk.candidates or not chunk.candidates[0].content or not chunk.candidates[0].content.parts: | |
| continue | |
| candidate = chunk.candidates[0].content.parts[0] | |
| # Check for inline image data | |
| if candidate.inline_data: | |
| save_binary_file(temp_path, candidate.inline_data.data) | |
| print(f"Smile enhancement image generated: {temp_path}") | |
| image_path = temp_path | |
| # If an image is found, we assume that is the desired output. | |
| break | |
| else: | |
| # Accumulate text response if no inline_data is present. | |
| text_response += chunk.text + "\n" | |
| print("Received text response from Gemini API") | |
| except Exception as e: | |
| print(f"Error during content generation: {str(e)}") | |
| # Continue with the function, returning empty responses | |
| except Exception as e: | |
| print(f"Error in Gemini API setup: {str(e)}") | |
| return None, f"Error: {str(e)}" | |
| finally: | |
| # Always clean up files | |
| try: | |
| if 'files' in locals() and files: | |
| del files | |
| except: | |
| pass | |
| return image_path, text_response | |
| def assess_image_quality(original_image, enhanced_image): | |
| """ | |
| Assesses the quality of the enhanced image based on specific criteria. | |
| Returns a tuple of (is_acceptable, feedback_message) | |
| """ | |
| try: | |
| # Check if enhanced image exists | |
| if enhanced_image is None: | |
| return False, "No enhanced image generated" | |
| # Image dimension checks | |
| if enhanced_image.size[0] < 100 or enhanced_image.size[1] < 100: | |
| return False, "Enhanced image appears to be too small or improperly sized" | |
| # Check that the enhanced image has similar dimensions to the original | |
| # This helps ensure facial proportions are maintained | |
| width_diff = abs(original_image.size[0] - enhanced_image.size[0]) | |
| height_diff = abs(original_image.size[1] - enhanced_image.size[1]) | |
| # If dimensions are significantly different, it suggests the image proportions changed | |
| if width_diff > 20 or height_diff > 20: | |
| return False, "Enhanced image dimensions differ significantly from original, suggesting facial proportions may have changed" | |
| # Check image has proper RGB channels for natural skin tones | |
| if enhanced_image.mode != 'RGB': | |
| return False, "Enhanced image does not have the correct color mode" | |
| # For now, we'll do basic checks and assume the model follows guidelines | |
| return True, "Image passes quality assessment criteria" | |
| except Exception as e: | |
| print(f"Error in quality assessment: {str(e)}") | |
| # Default to not accepting the image if assessment fails | |
| return False, f"Assessment error: {str(e)}" | |
| def process_smile_enhancement(input_image, max_attempts=2): | |
| try: | |
| if input_image is None: | |
| return None, "", "" | |
| # Get API key from environment variable | |
| gemini_api_key = "AIzaSyCVzRDxkuvtaS1B22F_F-zl0ehhXR0nuU8" | |
| if not gemini_api_key: | |
| print("Error: GEMINI_API_KEY not found in environment variables") | |
| return [input_image], "", "API key not configured" | |
| # Save the input image to a temporary file | |
| with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp: | |
| input_path = tmp.name | |
| input_image.save(input_path) | |
| print(f"Input image saved to {input_path}") | |
| # Initialize attempt counter and result variables | |
| current_attempt = 0 | |
| result_img = None | |
| feedback_history = [] | |
| max_processing_time = 120 # Maximum time in seconds for overall processing | |
| start_processing_time = time.time() | |
| while current_attempt < max_attempts: | |
| # Check if overall processing time exceeded | |
| if time.time() - start_processing_time > max_processing_time: | |
| print(f"Overall processing time exceeded {max_processing_time} seconds") | |
| break | |
| current_attempt += 1 | |
| print(f"Starting processing attempt {current_attempt}/{max_attempts}...") | |
| # Create a comprehensive prompt for true smile enhancement that affects facial features naturally | |
| # Adjust prompt based on previous attempts if needed | |
| prompt = """ | |
| Create a naturally enhanced smile that suits this specific person's face and character. Make the following personalized improvements: | |
| - Slightly enhance the existing teeth while PRESERVING their natural color, spacing, and individual characteristics | |
| - DO NOT make teeth perfectly white or perfectly aligned - keep some natural variation and character | |
| - Create subtle, natural smile lines around the eyes (crow's feet) appropriate for this person's age and face | |
| - Slightly raise the cheeks WITHOUT widening the face | |
| - Add a slight narrowing of the eyes that happens in genuine smiles | |
| - Create subtle dimples ONLY if they already exist in the original image | |
| - Enhance the overall joyful expression while maintaining the person's unique facial structure | |
| IMPORTANT GUIDELINES: | |
| - PRESERVE THE PERSON'S NATURAL DENTAL CHARACTERISTICS - teeth should still look like THEIR teeth, just slightly enhanced | |
| - Keep teeth coloration natural and appropriate for the person - avoid unnaturally white teeth | |
| - Maintain slight natural imperfections in tooth alignment that give character to the smile | |
| - Create a genuine, authentic-looking smile that affects the entire face naturally | |
| - ABSOLUTELY CRITICAL: DO NOT widen the face or change face width at all | |
| - Preserve the person's identity completely (extremely important) | |
| - Preserve exact facial proportions and face width of the original image | |
| - Maintain natural-looking results appropriate for the person's age and face structure | |
| - Keep teeth proportionate to the face - avoid making them too large or prominent | |
| - Maintain proper tooth-to-face ratio and ensure teeth fit naturally within the mouth | |
| - Keep the original background, lighting, and image quality intact | |
| - Ensure the enhanced smile looks natural, genuine, and believable | |
| - Create a smile that looks like a moment of true happiness for THIS specific person | |
| - Remember that not everyone has or wants perfect white teeth - the enhancements should SUIT THE INDIVIDUAL | |
| - If teeth are enhanced, maintain their natural characteristics while making subtle improvements | |
| """ | |
| # If not the first attempt, add previous feedback to the prompt | |
| if current_attempt > 1 and feedback_history: | |
| prompt += """ | |
| IMPORTANT FEEDBACK FROM PREVIOUS ATTEMPT: | |
| """ + " ".join(feedback_history) + """ | |
| Please address these issues in this new attempt. | |
| """ | |
| # Process silently | |
| print(f"Processing attempt {current_attempt}/{max_attempts}...") | |
| # Set timeout for individual API call | |
| api_call_timeout = time.time() + 45 # 45 second timeout for API call | |
| try: | |
| # Process the image using Google's Gemini model with timeout | |
| image_path, text_response = generate(text=prompt, file_name=input_path, api_key=gemini_api_key) | |
| # Check if API call timeout occurred | |
| if time.time() > api_call_timeout: | |
| print("API call timeout occurred") | |
| feedback_history.append("API call timed out, trying again with simplified request.") | |
| continue | |
| print(f"API response received: Image path: {image_path is not None}, Text length: {len(text_response)}") | |
| if image_path: | |
| # Load and convert the image if needed | |
| try: | |
| result_img = Image.open(image_path) | |
| if result_img.mode == "RGBA": | |
| result_img = result_img.convert("RGB") | |
| print("Successfully loaded generated image") | |
| # Assess the quality of the enhanced image | |
| is_acceptable, assessment_feedback = assess_image_quality(input_image, result_img) | |
| print(f"Image quality assessment: {is_acceptable}, {assessment_feedback}") | |
| if is_acceptable: | |
| # Return the enhanced image with success message | |
| success_message = "Successfully loaded generated image\nImage quality assessment: True, Image passes quality assessment criteria" | |
| return [result_img], "", success_message | |
| else: | |
| # Image didn't pass quality assessment, add feedback for next attempt | |
| feedback_history.append(assessment_feedback) | |
| # If we've reached max attempts, return the best result so far | |
| if current_attempt >= max_attempts: | |
| print("Max attempts reached, returning best result") | |
| return [result_img], "", "" | |
| except Exception as img_error: | |
| print(f"Error processing the generated image: {str(img_error)}") | |
| feedback_history.append(f"Error with image: {str(img_error)}") | |
| else: | |
| # No image was generated, only text response | |
| print("No image was generated, only text response") | |
| feedback_history.append("No image was generated in the previous attempt.") | |
| # If we've reached max attempts, return the original image | |
| if current_attempt >= max_attempts: | |
| print("Max attempts reached, returning original image") | |
| return [input_image], "", "" | |
| except Exception as gen_error: | |
| print(f"Error during generation attempt {current_attempt}: {str(gen_error)}") | |
| feedback_history.append(f"Error during processing: {str(gen_error)}") | |
| # If we've reached max attempts, return the original image | |
| if current_attempt >= max_attempts: | |
| return [input_image], "", "" | |
| # Return the original image as a fallback without messages | |
| print("Returning original image as fallback") | |
| return [input_image], "", "" | |
| except Exception as e: | |
| # Return the original image silently on error | |
| print(f"Overall error in process_smile_enhancement: {str(e)}") | |
| return [input_image], "", "" | |
| # Create a clean interface with minimal UI elements | |
| with gr.Blocks(title="Smile Enhancement") as demo: | |
| with gr.Row(): | |
| with gr.Column(): | |
| image_input = gr.Image( | |
| type="pil", | |
| label=None, | |
| image_mode="RGB", | |
| elem_classes="upload-box" | |
| ) | |
| submit_btn = gr.Button("Enhance Smile with Natural Expressions", elem_classes="generate-btn") | |
| with gr.Column(): | |
| output_gallery = gr.Gallery(label=None) | |
| # Show feedback text to display assessment results | |
| feedback_text = gr.Textbox(label="Status", visible=True) | |
| # Hidden element for structure | |
| output_text = gr.Textbox(visible=False) | |
| submit_btn.click( | |
| fn=process_smile_enhancement, | |
| inputs=[image_input], | |
| outputs=[output_gallery, output_text, feedback_text] | |
| ) | |
| demo.queue(max_size=50).launch(share=True) |